Compute Aligned Training: Optimizing for Test Time Inference

1. Problem Statement: The Training-Inference Misalignment

The paper addresses a fundamental disconnect in modern Large Language Model (LLM) development: training objectives are optimized for single-sample performance, while deployment increasingly relies on test-time compute strategies (search).

• The Status Quo: Standard post-training methods like Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) optimize for the likelihood of a single correct output (e.g., minimizing Cross-Entropy or maximizing expected reward of one sample).
• The Conflict: Modern inference strategies (e.g., Pass@N, Majority Vote, Best-of-N) generate multiple candidate solutions and aggregate them to find the best answer.
  • Example (Pass@N): If a model has a 50% chance of being correct, increasing it to 90% yields diminishing returns for a Pass@100 metric (where success is already near-certain). Standard training, however, continues to push for 100% likelihood, wasting capacity on "easy" samples.
  • Example (Majority Vote): A model needs only a plurality (e.g., 35% probability) to win against rivals. Standard training pushes for $p \to 1.0$, encouraging "landslide" victories rather than robust consensus, which can degrade performance when the model is forced to act as an ensemble.
• Consequence: This misalignment causes policy collapse (reduced diversity) and prevents models from effectively utilizing increased test-time compute budgets.

2. Methodology: Compute Aligned Training (CAT)

The authors propose Compute Aligned Training (CAT), a framework that aligns training objectives with the specific inference strategy used at test time.

Core Conceptualization

The authors conceptualize inference strategies not as post-hoc heuristics, but as differentiable operators ($T$) that transform the base policy $\pi_\theta$ into an effective test-time distribution $\tilde{\pi}_\theta$:
$$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$$
where $\phi$ represents hyperparameters (e.g., sample budget $N$).

The Unified Loss Function

Instead of minimizing the negative log-likelihood of the ground truth under the base policy ($\pi_\theta$), CAT minimizes the loss under the effective system-level policy ($\tilde{\pi}_\theta$):
$$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$$

Gradient Scaling Mechanism

To make this computationally feasible without sampling the full output space, the authors derive a diagonal gradient approximation. They show that the gradient of the CAT objective can be approximated by scaling the standard gradient by a weight factor $w(p, \phi)$:
$$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$$
Here, $p$ is the base probability of the correct answer, and $w$ represents the marginal utility of increasing $p$ with respect to the test-time strategy.

Key Scaling Factors Derived:

1. Pass@N (SFT & RL):
   • $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
   • Behavior: Suppresses gradients for easy samples (high $p$) where success is already guaranteed, focusing learning on "hard" samples where marginal gains improve coverage.
2. Majority Vote (SFT & RL):
   • $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
   • Behavior: Concentrates gradient mass near the consensus threshold ($p \approx k/N$). It ignores samples that are already winning or hopelessly losing, focusing only on the "tipping point."
3. Best-of-N (RL):
   • $w \propto N (P_{<y})^{N-1}$
   • Behavior: Creates "winner-take-all" dynamics. Gradients reinforce only the highest-reward samples in a batch, encouraging the model to explore high-variance, high-reward tails of the distribution rather than collapsing to a safe mean.

Theoretical Guarantees

The paper proves that for Competitive Strategies (where increasing a rival's probability hurts success), the diagonal approximation is a conservative underestimate of the true gradient. This ensures the optimizer never steps in the wrong direction, acting as a natural stabilizer against policy collapse.

3. Key Contributions

1. Unified Operator Framework: Formalizes test-time strategies as differentiable operators, deriving a general mechanism to align training with arbitrary inference strategies (SFT and RL).
2. Extension Beyond Pass@N: Generalizes previous work (e.g., Direct Coverage Optimization) to Majority Vote and Best-of-N, covering both SFT and RL paradigms.
3. Modality Agnosticism: Demonstrates the framework's applicability beyond LLMs to Protein Language Models (PLMs), showing improved performance in complex, multi-modal reward landscapes.
4. Variance Analysis: Provides a theoretical and empirical analysis of gradient variance inflation, proposing normalization techniques (e.g., Log-Weighted estimators) to stabilize training.

4. Experimental Results

The authors evaluated CAT on the MATH benchmark (using Mistral-7B) and Protein Generation (using ProtGPT2).

A. Language Models (MATH Benchmark)

• Pass@N (SFT): Models trained with CAT for $N=64$ sacrificed single-sample accuracy (Pass@1 dropped by ~2.3%) but achieved a 7.8% absolute improvement in Pass@64 coverage compared to standard SFT.
• Majority Vote (SFT): CAT models consistently outperformed baselines across all inference budgets. A "sweet spot" was observed at $N=16$, where the model balanced bias and variance optimally.
• Reinforcement Learning (GRPO): Standard RL suffered from policy collapse, saturating quickly (Pass@32 $\approx$ 35.8%). CAT-trained models preserved the necessary variance for search, achieving 40.0% at Pass@32.
• Warmup Effect: Using CAT for the SFT warmup phase significantly improved the subsequent RL performance, even boosting single-shot (Pass@1) accuracy by enriching the model's reasoning paths.

B. Protein Generation (Protein Language Models)

• Unconditional (Valley of Death): In a multi-modal reward landscape (local optimum at $h \approx 0.35$, global optimum at $h \approx 0.75$, separated by a "valley" of aggregation), standard RL got stuck in the local optimum. Best-of-N CAT successfully "jumped" the valley, finding the global maximum by maintaining high distributional variance.
• Conditional Generation: CAT models learned to generate sequences with inverse hydrophobicity to inputs, demonstrating superior scaling and alignment with the relational task compared to standard RL.

5. Significance and Implications

• Paradigm Shift: The paper argues that training objectives must be designed with the downstream inference procedure in mind. As test-time compute becomes the primary lever for scaling model capability, training must optimize for the aggregated outcome, not the single sample.
• Efficiency: CAT achieves these gains without increasing training compute costs. It simply re-weights existing gradients based on the marginal utility of the sample to the test-time strategy.
• Stability vs. Exploration: The framework provides a principled way to trade off single-sample precision for distributional diversity, enabling models to effectively utilize search procedures like Best-of-N and Majority Vote without collapsing into mode-seeking behavior.
• Generalizability: By proving efficacy in both LLMs and Protein LMs, the paper suggests that "Compute Aligned Training" is a universal principle for any generative model utilizing test-time search.

In conclusion, Compute Aligned Training bridges the gap between learning and search, ensuring that models are not just "smart" at generating one answer, but are optimized to be "smart" at finding the best answer when given the computational budget to search.