Adaptive gradient sampling inspired by uncertainty-aware reduced-order models can reduce the number of expensive function evaluations needed in zeroth-order LLM optimization.

An adaptive gradient sampling strategy, informed by uncertainty estimates from a reduced-order surrogate model (e.g., Gaussian Process or Bayesian neural network approximating the LLM loss landscape), will reduce the total number of LLM forward-pass evaluations required to reach a target optimization quality (e.g., within 5% of the best-known solution) by at least 30% compared to uniform random gradient estimation (e.g., standard SPSA or ZO-SGD with fixed sampling) when optimizing discrete or continuous prompt/instruction parameters for a large language model (≥1B parameters).

PHASE 1 — Surrogate Feasibility (Days 1–14): Establish that a reduced-order model can predict LLM evaluation outcomes with sufficient accuracy to guide sampling. Sample 200 random points in the optimization landscape for 2 tasks, fit GP and BNN surrogates, measure R² and calibration.

PHASE 2 — Ablation of Sampling Strategies (Days 15–35): Compare 4 methods: (A) uniform ZO-SGD baseline, (B) uncertainty-guided sampling (exploit low-uncertainty regions), (C) uncertainty-guided exploration (sample high-uncertainty regions), (D) combined adaptive strategy balancing exploration/exploitation. Use 3 tasks × 20 runs × 500 evaluation budget.

PHASE 3 — Full Benchmark (Days 36–60): Run all methods on 5 diverse NLP tasks with 3 LLM sizes (1B, 7B, 13B). Primary metric: number of LLM calls to reach 95% of best-observed objective. Secondary: final objective quality at fixed budget (500 calls).

PHASE 4 — Analysis and Ablation (Days 61–75): Sensitivity analysis on surrogate type (GP vs. BNN vs. random forest), dimensionality of search space, and noise level. Identify boundary conditions empirically.

# Adaptive Gradient Sampling for ZO-LLM Optimization
# Architecture Overview

import torch
import gpytorch
from dataclasses import dataclass
from typing import Callable, List, Tuple

@dataclass
class AdaptiveZOConfig:
    budget: int = 500          # Max LLM evaluations
    dim: int = 100             # Search space dimensionality
    delta: float = 0.01        # Perturbation magnitude
    k_directions: int = 10     # Directions per gradient estimate
    n_candidates: int = 1000   # Candidate directions to score
    beta_init: float = 2.0     # UCB exploration weight (initial)
    beta_final: float = 0.5    # UCB exploration weight (final)
    surrogate_retrain_freq: int = 10  # Retrain every N evaluations
    warmup_evals: int = 20     # Random evals before surrogate active

class ExactGPSurrogate(gpytorch.models.ExactGP):
    """Gaussian Process surrogate for LLM loss landscape."""
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=train_x.shape[-1])
        )
    def forward(self, x):
        return gpytorch.distributions.MultivariateNormal(
            self.mean_module(x), self.covar_module(x)
        )

class AdaptiveZOOptimizer:
    def __init__(self, llm_eval_fn: Callable, config: AdaptiveZOConfig):
        self.eval_fn = llm_eval_fn  # Black-box LLM evaluation
        self.cfg = config
        self.eval_history: List[Tuple] = []  # (direction, f_plus, f_minus)
        self.surrogate = None
        self.n_evals = 0
        
    def _compute_beta(self) -> float:
        """Anneal exploration weight over optimization."""
        progress = self.n_evals / self.cfg.budget
        return self.cfg.beta_init * (1 - progress) + self.cfg.beta_final * progress
    
    def _fit_surrogate(self, X: torch.Tensor, y: torch.Tensor):
        """Fit GP surrogate to accumulated evaluations."""
        likelihood = gpytorch.likelihoods.GaussianLikelihood()
        model = ExactGPSurrogate(X, y, likelihood)
        model.train(); likelihood.train()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
        for _ in range(100):  # Training iterations
            optimizer.zero_grad()
            output = model(X)
            loss = -mll(output, y)
            loss.backward()
            optimizer.step()
        model.eval(); likelihood.eval()
        return model, likelihood
    
    def _score_candidates(self, candidates: torch.Tensor,
                          model, likelihood) -> torch.Tensor:
        """Score candidate directions via UCB acquisition."""
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            pred = likelihood(model(candidates))
            mu = pred.mean
            sigma = pred.variance.sqrt()
        beta = self._compute_beta()
        return mu + beta * sigma  # UCB: higher = more promising to evaluate
    
    def _select_directions(self, x_current: torch.Tensor) -> torch.Tensor:
        """Select k perturbation directions adaptively."""
        # Generate random candidate directions
        candidates = torch.randn(self.cfg.n_candidates, self.cfg.dim)
        candidates = candidates / candidates.norm(dim=1, keepdim=True)  # Normalize
        
        if self.surrogate is None or self.n_evals < self.cfg.warmup_evals:
            # Warmup: uniform random selection
            idx = torch.randperm(self.cfg.n_candidates)[:self.cfg.k_directions]
        else:
            model, likelihood = self.surrogate
            # Feature: concatenate current point with direction
            features = torch.cat([
                x_current.unsqueeze(0).expand(self.cfg.n_candidates, -1),
                candidates
            ], dim=1)
            scores = self._score_candidates(features, model, likelihood)
            # Select top-k by UCB score
            idx = torch.topk(scores, self.cfg.k_directions).indices
        
        return candidates[idx]
    
    def _estimate_gradient(self, x: torch.Tensor,
                           directions: torch.Tensor) -> torch.Tensor:
        """ZO gradient estimate via 2-point SPSA."""
        grad_estimate = torch.zeros_like(x)
        for d in directions:
            f_plus = self.eval_fn(x + self.cfg.delta * d)
            f_minus = self.eval_fn(x - self.cfg.delta * d)
            self.n_evals += 2
            # Update surrogate training data
            self.eval_history.append((d.numpy(), f_plus, f_minus))
            # Accumulate gradient
            grad_estimate += (f_plus - f_minus) / (2 * self.cfg.delta) * d
        return grad_estimate / len(directions)
    
    def _update_surrogate(self):
        """Retrain surrogate on accumulated data."""
        if len(self.eval_history) < 10:
            return
        # Construct training data from evaluation history
        X_list, y_list = [], []
        for (d, f_plus, f_minus) in self.eval_history:
            X_list.append(d)
            y_list.append((f_plus + f_minus) / 2)  # Use mean as target
        X = torch.tensor(X_list, dtype=torch.float32)
        y = torch.tensor(y_list, dtype=torch.float32)
        # Normalize
        y = (y - y.mean()) / (y.std() + 1e-8)
        self.surrogate = self._fit_surrogate(X, y)
    
    def optimize(self, x_init: torch.Tensor) -> Tuple[torch.Tensor, List]:
        """Main optimization loop."""
        x = x_init.clone().requires_grad_(False)
        # Adam state
        m, v = torch.zeros_like(x), torch.zeros_like(x)
        lr, beta1, beta2, eps = 0.01, 0.9, 0.999, 1e-8
        trajectory = []
        
        t = 0
        while self.n_evals < self.cfg.budget:
            t += 1
            # Select directions adaptively
            directions = self._select_directions(x)
            # Estimate gradient
            grad = self._estimate_gradient(x, directions)
            # Adam update
            m = beta1 * m + (1 - beta1) * grad
            v = beta2 * v + (1 - beta2) * grad ** 2
            m_hat = m / (1 - beta1 ** t)
            v_hat = v / (1 - beta2 ** t)
            x = x - lr * m_hat / (v_hat.sqrt() + eps)
            # Retrain surrogate periodically
            if self.n_evals % self.cfg.surrogate_retrain_freq == 0:
                self._update_surrogate()
            # Log trajectory
            trajectory.append({'n_evals': self.n_evals, 'x': x.clone()})
        
        return x, trajectory

# Evaluation harness
def run_comparison(task_name: str, llm_model, n_runs: int = 20):
    results = {'adaptive': [], 'uniform': []}
    for seed in range(n_runs):
        torch.manual_seed(seed)
        x0 = torch.randn(100)  # Random initialization
        # Adaptive method
        adaptive_opt = AdaptiveZOOptimizer(
            llm_eval_fn=lambda x: evaluate_llm(llm_model, x, task_name),
            config=AdaptiveZOConfig()
        )
        _, traj_adaptive = adaptive_opt.optimize(x0.clone())
        # Uniform baseline (beta=0 disables UCB, uses random selection)
        uniform_opt = AdaptiveZOOptimizer(
            llm_eval_fn=lambda x: evaluate_llm(llm_model, x, task_name),
            config=AdaptiveZOConfig(beta_init=0.0, beta_final=0.0,
                                    warmup_evals=500)  # Always random
        )
        _, traj_uniform = uniform_opt.optimize(x0.clone())
        results['adaptive'].append(traj_adaptive)
        results['uniform'].append(traj_uniform)
    return compute_efficiency_metrics(results)