Optimal Token Baseline (OTB)

Last updated: 12/25/2025.

Optimal Token Baseline (OTB) is dynamic token-level baseline for variance reduction. It weights updates based on "Realized Energy"—essentially, how much uncertainty has accumulated up to that specific token. It downweights the noisy parts and trusts the clear signals. Read Optimal Token Baseline blog for more details.

The method: OTB

OTB builds a dynamic baseline that adapts to each token by tracking the “Realized Energy”—the uncertainty that has accumulated up to that token. It downweights the noisy parts and trusts the clear signals.
Unlike standard group means (which average over the padding EOS token ineffectively), OTB handles this naturally by computing baselines only over valid tokens.

Logit-Gradient Proxy

Computing true uncertainty requires expensive backward passes (calculating gradient norms per token). Instead, OTB introduces the Logit-Gradient Proxy: the realized energy can be estimated entirely from forward probabilities.
This means zero extra backward calls and effectively no additional runtime overhead.

Mechanics at a glance

For each prompt group of size N, OTB computes rewards-to-go G_t and cumulative variance weights W_t. The optimal baseline per token is

B*_t = (Σ_i G_t^{(i)} · W_t^{(i)}) / (Σ_i W_t^{(i)} + ε),
W_t = Σ_{j=1}^t (1 - 2π_j + Σπ_j²),
Σπ_j² = exp(logsumexp(2·logits_j) - 2·logsumexp(logits_j)).

The final advantage is (G_t - B*_t) · mask_t, so padding tokens stay at zero.

Integration in VERL

AdvantageEstimator.OPTIMAL_TOKEN_BASELINE registers compute_optimal_token_baseline_advantage, invoked whenever algorithm.adv_estimator is set to optimal_token_baseline.
ActorRolloutRefWorker.compute_log_prob emits an additional tensor sum_pi_squared (Σπ² per token) when actor.calculate_sum_pi_squared=True. This requires disabling fused log-prob kernels, because they do not surface logits.
Trainers assert sum_pi_squared exists, regroup trajectories by non_tensor_batch["uid"], and run the OTB calculation. If rollout IS is active, they rescale the weights by rollout_is_weights**2 before aggregating.
In Ulysses sequence-parallel setups, the actor gathers, unpads, and returns Σπ² in the same way it handles log-probabilities, so OTB supports sharded sequence-parallel models out of the box.
sum_pi_squared_checkpointing is available to trade compute for memory when Σπ² tensors become large (e.g., lengthy chain-of-thought reasoning).

Configuration checklist

actor_rollout_ref.actor.calculate_sum_pi_squared: true (mandatory).
actor_rollout_ref.model.use_fused_kernels: false (required until fused kernels emit logits).
algorithm.adv_estimator: optimal_token_baseline.
Group sampling (actor_rollout_ref.rollout.n > 1) to unlock OTB’s variance reduction; with n=1 the baseline collapses to returns.

Example OmegaConf overlay:

algorithm:
  adv_estimator: optimal_token_baseline

actor_rollout_ref:
  actor:
    calculate_sum_pi_squared: true
    sum_pi_squared_checkpointing: false # optional memory saver
  rollout:
    n: 8

Example script

examples/otb_trainer/run_qwen2_5-7b.sh.

Gradient Variance Proxy Metrics

All gradient-variance analysis in the Optimal Token Baseline work starts from the variance identity

Var(ĝ) = E[||ĝ||²] - ||E[ĝ]||²,

which states that the variance of any stochastic gradient equals the mean squared magnitude minus the squared norm of its expectation.

For a trajectory τ, the policy-gradient estimator is

ĝ(τ) = ∇ log π_θ(τ) · A(τ),        A(τ) = R(τ) - B.

The logit-gradient proxy approximates the squared gradient norm without an extra backward pass:

||ĝ(τ)||² ≈ Ŵ(τ) · A(τ)²,

where Ŵ(τ) is the realized energy built. Given a mini-batch {τ_i} of size N, we decompose its statistics into three diagnostics:

Signal strength (squared norm of the mean gradient)
```
S = || (1/N) · Σ ĝ(τ_i) ||²
```

Total power (signal + noise)

P_total = (1/N) · Σ Ŵ(τ_i) · A(τ_i)²

Pure noise (estimated variance of the batch mean)
```
Var_proxy = (1/(N-1)) · (P_total - S)
```

verl/trainer/ppo/metric_utils.py#L306 implements these diagnostics via compute_variance_proxy_metrics, emitting variance_proxy/proxy1_signal_strength, variance_proxy/proxy2_total_power, and variance_proxy/proxy3_pure_noise.

Tracking these metrics provides a forward-only, low-overhead view of gradient health for any advantage estimator that supplies sum_pi_squared.