Compute Aligned Training (CAT)

• Compute Aligned Training (CAT) is a framework that redefines model training by optimizing for the transformed distribution produced by deployment-time inference operators.
• It employs an operator-level view with a diagonal approximation to reweight gradients, enabling both supervised fine-tuning and reinforcement learning to target final system output.
• CAT’s methodology improves high-budget inference performance while addressing variance control and bias–variance tradeoffs, with benefits demonstrated across language and protein design tasks.

Searching arXiv for the requested topic and closely related work. Searching for "Compute Aligned Training" on arXiv. Compute Aligned Training (CAT) is a post-training framework for generative models in which the optimization target is the distribution induced by the deployment-time inference procedure, rather than the raw one-sample base policy. In the formulation introduced in "Compute Aligned Training: Optimizing for Test Time Inference" (Ousherovitch et al., 27 Apr 2026), CAT addresses a mismatch between standard supervised fine-tuning (SFT) or reinforcement learning (RL)—which optimize individual samples under a base policy—and modern test-time strategies such as Pass@$\!N$, majority vote, filtering, reranking, and Best-of-$N$, which consume multiple samples and return a single final output. The framework therefore treats inference as an operator on the base policy and derives training losses that maximize performance after that operator is applied.

1. Operator-level view of inference

CAT begins from a distinction between the base policy

$\pi_\theta(\cdot \mid x)$

and the effective deployment policy

$$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$

where $\mathcal{T}$ is a test-time strategy (TTS) and $\phi$ denotes strategy hyperparameters such as sample budget $N$, threshold $k$, or verifier settings (Ousherovitch et al., 27 Apr 2026). The central claim is that evaluation increasingly occurs under $\tilde{\pi}_\theta$, not under $\pi_\theta$. Standard training is therefore misaligned whenever deployment uses search, aggregation, or filtering.

The practical motivation is clearest in search-based settings. Under Pass@$N$0, once the correct answer already appears with moderate probability, further pushing that probability toward $N$1 can have little marginal value, because success is already likely across $N$2 draws. Under majority vote, the relevant objective is not that the correct answer dominate the full probability mass, but that it exceed the strongest rival often enough to win plurality. CAT formalizes this intuition as a marginal-utility reweighting problem: updates should be concentrated where they matter for the final system returned by the chosen inference operator.

This operator view also reframes “test-time scaling” as an optimization target rather than merely an evaluation trick. A plausible implication is that CAT treats extra inference compute as part of the model-definition pipeline: the deployable system is the composition of a policy and a TTS operator, not the policy alone.

2. Exact objectives and diagonal approximation

For SFT, CAT replaces ordinary cross-entropy on the base policy with negative log-likelihood under the transformed policy: $N$3 If

$N$4

the exact gradient is

$N$5

For RL, the CAT objective is expected reward under the transformed policy: $N$6 Its exact gradient weights each sample by its effect on deployment-time utility rather than by one-sample utility alone (Ousherovitch et al., 27 Apr 2026).

Because the full Jacobian over $N$7 is generally impractical, CAT uses a diagonal approximation in which $N$8 is treated mainly as a function of the target’s own base probability $N$9, ignoring off-diagonal dependence on rivals. Under this approximation, SFT reduces to a scalar rescaling of the ordinary cross-entropy gradient: $\pi_\theta(\cdot \mid x)$0 RL analogously becomes

$\pi_\theta(\cdot \mid x)$1

The paper writes the exact gradient as $\pi_\theta(\cdot \mid x)$2 and proves a sufficient descent condition: $\pi_\theta(\cdot \mid x)$3 It further distinguishes orthogonal strategies, for which success depends only on the target probability and the approximation is exact, from competitive strategies, in which increasing rival probability harms the target and the diagonal approximation is conservative. Pass@$\pi_\theta(\cdot \mid x)$4 is orthogonal; majority vote is analyzed via a threshold relaxation rather than exact plurality (Ousherovitch et al., 27 Apr 2026).

3. Compute-aligned supervised fine-tuning

In CAT-SFT, the training target is not “maximize the probability of this demonstration under one sample,” but “maximize the probability that the deployed inference strategy returns the correct answer.” For reasoning tasks with chain-of-thought, the relevant marginal answer probability is

$\pi_\theta(\cdot \mid x)$5

but this is intractable. In practice, CAT uses the probability of the provided reasoning trace $\pi_\theta(\cdot \mid x)$6 as a lower-bound proxy for $\pi_\theta(\cdot \mid x)$7, while retaining the operator-aware weight $\pi_\theta(\cdot \mid x)$8 (Ousherovitch et al., 27 Apr 2026).

For Pass@$\pi_\theta(\cdot \mid x)$9, the transformed success probability is

$$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$0

yielding the CAT-SFT weight

$$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$1

This suppresses gradients on examples that are already likely to succeed within $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$2 trials. The implementation uses a log-space approximation when $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$3: $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$4

For majority vote, the paper uses a threshold relaxation: success occurs when the correct answer appears at least $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$5 times in $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$6 samples,

$$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$7

The corresponding SFT weight is

$$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$8

This concentrates learning near the consensus boundary. The paper also analyzes its variance: as $$\tilde{\pi}_\theta(y \mid x)=\mathcal{T}(\pi_\theta(\cdot \mid x),\phi)(y),$$9, $\mathcal{T}$0, implying variance inflation that scales approximately as $\mathcal{T}$1 for simple majority. That analysis is used to explain why majority-vote SFT exhibits a bias–variance tradeoff and why moderate $\mathcal{T}$2 can outperform very large $\mathcal{T}$3 (Ousherovitch et al., 27 Apr 2026).

The paper additionally sketches a token-level contrastive auxiliary loss for majority vote based on the token margin

$\mathcal{T}$4

with

$\mathcal{T}$5

This is not the main CAT algorithm, but it reflects the same principle: majority-vote alignment depends on controlling strong rivals, not only increasing the correct trace.

4. Compute-aligned reinforcement learning

CAT-RL optimizes reward after the deployment operator rather than reward of a single rollout. Under the diagonal approximation, the generic policy-gradient form becomes

$\mathcal{T}$6

Thus CAT-RL can be interpreted as a strategy-specific per-sample modulation of reward or advantage (Ousherovitch et al., 27 Apr 2026).

For Pass@$\mathcal{T}$7 with binary reward,

$\mathcal{T}$8

The paper reports that this exact RL-aligned weight can have severe variance at larger $\mathcal{T}$9, especially on hard examples where $\phi$0. To stabilize GRPO/PPO, it often uses a normalized “log-weighted” variant,

$\phi$1

which is biased but substantially more stable.

For majority vote RL, the threshold-relaxed weight is

$\phi$2

Unlike the SFT weight, this acts as a boundary-localized spotlight: it is largest where a marginal increase in probability is most likely to flip consensus. In RL, $\phi$3 need not be fixed; the paper estimates it from rollout batches using the empirical frequency of the strongest rival.

For Best-of-$\phi$4 RL, if $\phi$5 is the mass of outputs with reward strictly lower than $\phi$6, the probability that a given output wins is

$\phi$7

with weight

$\phi$8

This makes BoN training explicitly quantile-sensitive: only sufficiently high-ranked samples receive meaningful gradient.

CAT is integrated into PPO- and GRPO-style methods by multiplying the policy term by the CAT weight after reward or advantage normalization, rather than embedding the weight into the raw reward. For GRPO, the aligned objective is written as

$\phi$9

The paper also uses batch-level weight normalization,

$N$0

to stabilize overall step size (Ousherovitch et al., 27 Apr 2026).

5. Strategy-specific behavior and empirical results

The empirical program is organized around three claims: CAT works beyond Pass@$N$1, beyond SFT, and beyond LLMs (Ousherovitch et al., 27 Apr 2026). On MATH with Mistral-7B under SFT, a Pass@64-aligned model reduces Pass@1 from $N$2 to $N$3 but increases Pass@64 from $N$4 to $N$5. This is the canonical CAT tradeoff: single-sample quality can decline while high-budget inference improves substantially.

Majority-vote SFT shows a similar but more delicate pattern. The baseline achieves Maj@64 $N$6, whereas CAT reaches $N$7 for MajVote $N$8, $N$9 for $k$0, and $k$1 for $k$2. The fact that $k$3 is best is explicitly attributed to the bias–variance tradeoff and to sensitivity to the threshold hyperparameter $k$4. This is consistent with the theoretical variance analysis of majority-vote weighting.

In RL on MATH, standard GRPO trained for one-sample reward reaches Pass@32 $k$5, while CAT Pass@16 training reaches $k$6. For majority-vote RL, standard RL reaches Maj@16 $k$7, while CAT reaches $k$8 for Maj@4 training and $k$9 for Maj@8 training. The RL ablations are particularly important: at larger $\tilde{\pi}_\theta$0, the exact Pass@$\tilde{\pi}_\theta$1 RL estimator becomes unstable, whereas the normalized estimator provides the strongest scaling. This suggests that in CAT the variance-control machinery is structurally important, not merely a convenience.

The framework is also evaluated beyond language. On ProtGPT2 for protein design, Best-of-$\tilde{\pi}_\theta$2-aligned RL improves expected max reward at $\tilde{\pi}_\theta$3 from $\tilde{\pi}_\theta$4 under standard RL to $\tilde{\pi}_\theta$5 in an unconditional “valley of death” hydrophobicity landscape. In a conditional complementary-hydrophobicity task, standard RL has stronger $\tilde{\pi}_\theta$6 reward ($\tilde{\pi}_\theta$7) than BoN models ($\tilde{\pi}_\theta$8–$\tilde{\pi}_\theta$9), but at $\pi_\theta$0 it scales to only $\pi_\theta$1, whereas BoN models reach about $\pi_\theta$2–$\pi_\theta$3. The paper interprets this as a distinction between average-quality optimization and upper-tail optimization.

A related diagnostic is the alignment coefficient

$\pi_\theta$4

which is used to predict crossover points between models trained for different Pass@$\pi_\theta$5 budgets. The reported agreement with empirical scaling curves supports the idea that CAT is fundamentally a gradient-support matching method.

6. Practical implications, limitations, and related directions

CAT is designed to avoid the cost of simulating full test-time strategies during training. The paper contrasts its analytical weighting scheme with “full TTS rollouts,” where training would generate all $\pi_\theta$6 samples per prompt. That alternative scales training compute directly with the target inference budget and creates poor credit assignment; for example, in Pass@$\pi_\theta$7, once multiple rollouts are already correct, leave-one-out advantages can collapse to zero. CAT instead uses closed-form or approximated derivatives of the operator-induced success probability, so in SFT it adds essentially no overhead beyond computing sequence probability and a scalar weight, and in RL it requires only estimates of quantities such as $\pi_\theta$8, $\pi_\theta$9, or $N$00 (Ousherovitch et al., 27 Apr 2026).

The main limitation is that CAT is exact only for some operators. For Pass@$N$01, the diagonal approximation is exact. For majority vote and Best-of-$N$02, it is approximate and can leave performance on the table or become unstable. The paper therefore emphasizes warmup with ordinary CE, log-space computation, clipping of sequence probabilities and weights, dynamic threshold estimation, history buffers for quantile estimates, lower learning rates in high-variance regimes, and batch-level weight normalization. Another limitation is representational: CAT is most natural when the deployment strategy admits a tractable transformed policy $N$03 or at least tractable derivatives. More complex procedures such as genetic algorithms, Monte Carlo tree search, or richer agentic systems are explicitly identified as open territory.

A nearby but distinct line of work is "Compute as Teacher: Turning Inference Compute Into Reference-Free Supervision" (Jayalath et al., 17 Sep 2025). That method also treats extra inference-time compute as a source of training signal, but does so by synthesizing a teacher signal from multiple rollouts and then optimizing toward it, rather than by analytically reweighting gradients under a deployment operator. This suggests a broader research program in which inference compute can be aligned with training either through operator-aware objectives or through compute-generated supervision.

The acronym CAT is highly overloaded elsewhere in machine learning, where it also denotes unrelated methods such as Corpus Aware Training, Contrastive Adapter Training, Continuous Adversarial Training, and Closed-loop Adversarial Training. In the specific sense of (Ousherovitch et al., 27 Apr 2026), however, Compute Aligned Training refers to a framework for optimizing models for the policy that will actually be used at test time—namely, the base generator composed with a compute-amplifying inference procedure.