HyperAIHyperAI

Command Palette

Search for a command to run...

TTPO: テスト時方策最適化

概要

強化学習(RL)やオンポリシー自己蒸留(OPSD)といった近年の主要なポストトレーニング手法は、大規模言語モデルの数学的推論能力を急速に向上させてきたが、いずれも正解ラベルに依存するためテスト時学習(TTT)には適用できない。正解ラベルを多数決による擬似ラベルで置き換えるのは自然な代替案だが、これは脆弱である。すなわち、誤った投票が教師モデルを汚染し、すべてのトークンを誤導してしまう。我々は、この失敗モードが非対称であることを見出した。擬似ラベルと一致しないロールアウトは、投票自体が正しいか否かにかかわらず、通常は誤っている。この観察に基づき、我々はテスト時方策最適化(TTPO)を提案する。これは、一致するロールアウトをOPSDによって蒸留し、一致しないロールアウトをグループ化RLによって罰するという非対称な目的関数である。トークンレベルの選択により両分岐はさらに洗練される。蒸留では既に収束した位置の重みを下げ、RLでは確信度の高い誤りのみを罰する。擬似ラベル誤りが頻繁に生じる状況でも、両方の更新は十分に根拠づけられたまま維持され、多数決による経路選択はモデルの改善に伴ってより強固な自己教師信号をもたらす。ラベルを一切用いずに、TTPOは5つの競技レベルのベンチマークにおいてラベル教師ありOPSDと同等の性能を達成し、Qwen3-1.7BをTTTにおいて38.0%から45.2%へ向上させ、思考なしの設定では+25.2%から+36.4%の改善を示し、強いタスク横断汎化を実現する。コードはhttps://github.com/ZJU-REAL/TTPOで公開している

One-sentence Summary

Test-Time Policy Optimization (TTPO), proposed by researchers at Zhejiang University and Alibaba Group, is an asymmetric test-time training method that distills majority-vote agreeing rollouts via on-policy self-distillation, penalizes disagreeing rollouts with Grouped RL, refines both branches through token-level selection, and enables label-free mathematical reasoning that matches label-supervised OPSD while lifting Qwen3-1.7B from 38.0%38.0\%38.0% to 45.2%45.2\%45.2% on competition benchmarks.

Key Contributions

  • Majority-vote pseudo-labels remain useful for test-time training despite frequent errors: though wrong on ∼85% of competition-level prompts, ∼79% of disagreeing rollouts are wrong too, so penalizing disagreement stays correct while distillation does not.
  • TTPO, an asymmetric objective, distills agreeing rollouts toward an answer-conditioned teacher and penalizes disagreeing rollouts with GRPO penalties, with token-level selection in both branches.
  • Trained without any labels, TTPO matches or exceeds label-supervised OPSD on five competition-level benchmarks, raises Qwen3-1.7B from 38.0% to 45.2% in test-time training, and demonstrates strong cross-task generalization.

Introduction

The authors tackle test-time training for LLM mathematical reasoning, where a model must improve on unlabeled problems by learning from its own sampled solutions. Prior work either uses coarse sequence-level rewards from majority-vote pseudo-labels (e.g., TTRL), which reinforce errors when the pseudo-label is wrong, or relies on dense token-level distillation that requires ground-truth answers and is highly sensitive to label noise. The authors propose Test-Time Policy Optimization (TTPO), an asymmetric objective that combines GRPO penalties on rollouts that disagree with the pseudo-label (a robust negative signal) with on-policy self-distillation on rollouts that agree, conditioning the teacher on the rollouts’ own answer to stay safe even under incorrect pseudo-labels. Token-level selection sharpens both branches, yielding label-free training that matches or exceeds supervised on-policy distillation across model scales and competition benchmarks.

Method

The authors propose an asymmetric test-time training framework designed to optimize language models on competition-level problems without ground-truth labels. The core pipeline begins by sampling multiple reasoning trajectories for each test problem and extracting their final answers. These answers are clustered by mathematical equivalence to generate a majority-vote pseudo-label. The trajectories are then partitioned into two distinct sets: positive samples that agree with the pseudo-label, and negative samples that disagree.

As shown in the framework diagram below:

In the test-time training setting, majority-vote pseudo-labels are frequently incorrect. The authors observe that naively applying self-distillation to all trajectories using a corrupted teacher propagates errors across every token. To minimize the blast radius of these pseudo-label errors, the framework employs an asymmetric design. While distilling positive samples toward the pseudo-label-conditioned teacher, the model applies a reinforcement learning penalty to negative samples. This approach is highly robust because penalizing a disagreeing rollout remains correct for the vast majority of cases, relying solely on the disagreement rather than the potentially flawed pseudo-label content.

For the positive samples, the framework leverages On-Policy Self-Distillation (OPSD). The teacher distribution is constructed by conditioning the same model on the pseudo-label as privileged information, while the student distribution processes the original prompt. The teacher and student distributions are defined as:

qt(a^)=πθ([x;a^]teacher,y<t)q_t^{(\hat{a})} = \pi_\theta(\cdot \mid [x; \hat{a}]_{\text{teacher}}, y_{<t})qt(a^)=πθ([x;a^]teacher,y<t) pt=πθ(xstudent,y<t)p_t = \pi_\theta(\cdot \mid x_{\text{student}}, y_{<t})pt=πθ(xstudent,y<t)

The authors apply a forward Kullback-Leibler divergence loss with a specialized per-token weighting scheme:

LOPSD(k)=1Tkt=1Tkw(t)KL(qt(a^)pt)\mathcal{L}_{\text{OPSD}}(k) = \frac{1}{T_k} \sum_{t=1}^{T_k} w(t) \cdot \text{KL}(q_t^{(\hat{a})} \| p_t)LOPSD(k)=Tk1t=1Tkw(t)KL(qt(a^)pt)

Not all tokens offer equal learning value, so the weighting mechanism down-weights positions where the student has already converged. It measures two complementary signals: the student entropy and the teacher-student divergence. These signals are normalized and combined using a Soft-OR operation:

w(t)=H^(t)+Δ^(t)H^(t)Δ^(t)w(t) = \hat{H}(t) + \hat{\Delta}(t) - \hat{H}(t) \cdot \hat{\Delta}(t)w(t)=H^(t)+Δ^(t)H^(t)Δ^(t)

This assigns high weight when the student is either uncertain or confidently wrong, and near-zero weight when the student is already confident and aligned with the teacher.

For the negative samples, the framework applies Group Relative Policy Optimization (GRPO). Each trajectory receives a binary reward based on the majority-vote classification, and advantages are computed group-relatively across all rollouts for a given problem. Since negative samples inherently receive negative advantages, this branch effectively penalizes the disagreeing trajectories:

LGRPO(k)=AkTkt=1Tkm(t)logπθ(yk(t)x,yk(<t))\mathcal{L}_{\text{GRPO}}(k) = - \frac{A_k}{T_k} \sum_{t=1}^{T_k} m(t) \cdot \log \pi_\theta(y_k^{(t)} \mid x, y_k^{(<t)})LGRPO(k)=TkAkt=1Tkm(t)logπθ(yk(t)x,yk(<t))

To prevent collateral damage from false penalties on locally correct tokens within these failed trajectories, the authors introduce a token masking scheme. Tokens are scored based on their negative log-probability and normalized certainty:

s(t)=logπθ(yk(t)x,yk(<t))(1H^(t))s(t) = - \log \pi_\theta(y_k^{(t)} \mid x, y_k^{(<t)}) \cdot (1 - \hat{H}(t))s(t)=logπθ(yk(t)x,yk(<t))(1H^(t))

This scoring prioritizes genuinely anomalous outputs where the model confidently produced unlikely content, while naturally excluding locally correct high-probability tokens. A binary mask is then constructed by selecting the top half of tokens by score:

m(t)=1[s(t)median({s(t)}t=1Tk)]m(t) = \mathbf{1}[s(t) \ge \text{median}(\{s(t')\}_{t'=1}^{T_k})]m(t)=1[s(t)median({s(t)}t=1Tk)]

The final unified objective combines both branches, balancing their respective contributions with a weighting factor λ\lambdaλ. The overall loss is computed as the sum of the OPSD loss over the positive samples and the scaled GRPO loss over the negative samples, normalized by the batch size:

LTTPO=1B(kPLOPSD(k)+λkNLGRPO(k))\mathcal{L}_{\text{TTPO}} = \frac{1}{|\mathcal{B}|} \left( \sum_{k \in \mathcal{P}} \mathcal{L}_{\text{OPSD}}(k) + \lambda \sum_{k \in \mathcal{N}} \mathcal{L}_{\text{GRPO}}(k) \right)LTTPO=B1(kPLOPSD(k)+λkNLGRPO(k))

Experiment

The evaluation uses Qwen3 models fine-tuned with LoRA on five competition math benchmarks, comparing TTPO against label-dependent and label-free baselines in both supervised and test-time training settings. TTPO’s asymmetric objective—forward KL on positive samples and GRPO on negatives—combined with majority-vote pseudo-labels, consistently matches or surpasses label-dependent methods and significantly outperforms prior label-free approaches, allowing a 4B model to match an 8B base model. Ablations validate the token-level selection and update strategy, while analysis shows that a thinking teacher with short-answer conditioning provides robust guidance, the method generalizes across tasks, and pseudo-labels can even outperform ground-truth by sustaining a beneficial positive-negative split and fostering self-evolving improvement.

TTPO, using only majority-vote pseudo-labels, consistently outperforms label-dependent methods OPSD and GRPO across three model scales on the OpenThoughts training data. On the 4B model, TTPO matches the average performance of the untrained 8B base model, effectively amplifying reasoning capacity to that of a model twice its size. These results demonstrate that pseudo-labels can substitute for ground-truth annotations without sacrificing performance. TTPO surpasses OPSD and GRPO across all model scales despite using no ground-truth labels, with average gains widening at larger sizes. TTPO on Qwen3-4B achieves an average score of 58.6, equaling the 8B base model and effectively doubling reasoning capacity per parameter. Majority-vote pseudo-labels prove sufficient to replace ground-truth supervision, enabling TTPO to exceed even the label-dependent OPSD baseline.

On the Qwen3-1.7B model, the label-free test-time training method TTPO substantially outperforms both TTRL and OPSD-TTT, achieving an average score of 45.2 across three reasoning benchmarks. This represents a 7.2-point improvement over the base model and demonstrates that dense, token-level guidance from a thinking teacher with answer-conditioned pseudo-labels is more effective than binary rewards or deterministic self-distillation. TTPO reaches a 45.2 average on Qwen3-1.7B, exceeding OPSD-TTT by 3.3 points and TTRL by 5.4 points. The gap over TTRL highlights the benefit of dense distributional guidance over binary reward signals. TTPO's improvement over OPSD-TTT shows that majority-vote pseudo-labels with asymmetric distillation extract more value than greedy self-distillation.

The full token-level selection method outperforms variants that omit either positive-sample weighting or negative-sample masking. Removing positive weighting dilutes the gradient signal by uniformly distilling all tokens, while removing negative masking indiscriminately penalizes correct reasoning steps and allows noisy tokens to dominate updates. Combining both mechanisms focuses distillation on informative tokens and restricts penalties to error sites, yielding the best results. Dropping positive-sample weighting leads to lower scores across all benchmarks, as uniform distillation weakens the contribution of informative tokens. Removing negative-sample masking also degrades performance, with a pronounced drop on two benchmarks, because it penalizes locally correct steps and introduces noise from anomalous tokens. The full method benefits from complementary effects: positive weighting concentrates distillation on high-value tokens, while negative masking prevents collateral damage and noise injection.

When the teacher model uses a thinking mode, providing only the answer as privileged information slightly improves performance over no privileged information, while a full trajectory reduces performance by interfering with independent reasoning. With a non-thinking teacher, injecting a short answer severely degrades performance, and a full trajectory merely matches the no-privilege baseline. The thinking-teacher plus answer setting thus offers the best trade-off between guidance and robustness. With a thinking teacher, an answer-only hint yields the highest performance, while a full trajectory degrades it by dominating the context. With a non-thinking teacher, a short answer drastically lowers performance, and a trajectory provides no benefit over no privileged information. The thinking-teacher with answer configuration achieves the best balance of guidance and robustness. Disabling the teacher's thinking mode consistently harms performance across all types of privileged information.

The experiments evaluate TTPO, a test-time training method using majority-vote pseudo-labels from a thinking teacher, on reasoning benchmarks across multiple model scales, comparing against label-dependent (OPSD, GRPO) and label-free (TTRL, OPSD-TTT) baselines. Results demonstrate that pseudo-labels can fully substitute ground-truth annotations, enabling TTPO to match the performance of models twice its size, while dense token-level guidance with asymmetric distillation significantly outperforms binary rewards or greedy self-distillation. Ablations show that combining positive-sample weighting and negative-sample masking is essential for effective distillation, and that a thinking teacher given only answer-level privileged information achieves the best balance of guidance and robustness without interfering with the student's reasoning.


AIでAIを構築

アイデアからローンチまで — 無料のAIコーディング支援、すぐに使える環境、最高のGPU価格でAI開発を加速。

AI コーディング補助
すぐに使える GPU
最適な料金体系

HyperAI Newsletters

最新情報を購読する
北京時間 毎週月曜日の午前9時 に、その週の最新情報をメールでお届けします
メール配信サービスは MailChimp によって提供されています