HyperAIHyperAI

Command Palette

Search for a command to run...

TTPO: 테스트 시점 정책 최적화

초록

강화 학습(RL)과 온폴리시 자기 증류(OPSD)와 같은 최근의 주요 사후 학습 기법들은 대규모 언어 모델의 수학적 추론 능력을 빠르게 향상시켜 왔으나, 정답 레이블에 의존하기 때문에 테스트 시점 학습(TTT)에는 적용할 수 없다. 정답 레이블을 다수결 의사 레이블로 대체하는 것은 자연스러운 대안이지만, 잘못된 투표가 교사 모델을 오염시키고 모든 토큰을 오도할 수 있어 취약하다. 우리는 이러한 실패 양상이 비대칭적임을 관찰했다. 의사 레이블과 불일치하는 롤아웃은 투표 자체의 정확성과 무관하게 일반적으로 틀린 경우가 많다는 것이다. 이 관찰에 기반하여, 우리는 일치하는 롤아웃은 OPSD로 증류하고 불일치하는 롤아웃은 그룹 RL로 페널티를 부여하는 비대칭 목적 함수인 테스트 시점 정책 최적화(TTPO)를 제안한다. 토큰 수준 선택은 두 갈래를 더욱 정교화한다. 증류는 이미 수렴된 위치의 가중치를 낮추고, RL은 확신을 가진 오류만 페널티를 부여한다. 두 업데이트 모두 빈번한 의사 레이블 오류 하에서도 견고하게 유지되며, 다수결 라우팅은 모델이 개선됨에 따라 더욱 정밀한 자기 지도를 제공한다. 레이블 없이도 TTPO는 5개의 경시대회 수준 벤치마크에서 레이블 지도 OPSD와 동등한 성능을 달성하고, TTT 환경에서 Qwen3-1.7B의 성능을 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