Function-Space Temporal Self-Distillation for Policy Optimization

AI summary

Overview

Research area: Reinforcement learning post-training of large language models (RLVR / verifiable rewards), specifically on-policy self-distillation for reasoning and tool use. Technical level: Advanced — assumes familiarity with policy-gradient methods (GRPO), KL/JSD divergences, top-K distribution truncation, EMA teachers, and experience replay. Scope: A single-policy training framework that replaces a persistent teacher model with sparse, stored historical next-token distributions ("function-space" targets), evaluated on four SciKnowEval reasoning subsets plus ToolAlpaca with Qwen3-8B and Olmo3-7B-Instruct.

What This Paper Is About

Standard RLVR collapses an entire response into one scalar reward, making it hard to know which token-level decisions were good or bad. Self-distillation fixes this by comparing the model's unassisted view of a response with a "privileged" view that has seen a hint, a successful sibling rollout, or the reference answer — but because both views come from the same evolving policy, the teacher is a moving target that needs extra stabilization (EMA copies, trust regions). The paper asks whether temporal stability and reliability-aware supervision can be obtained without keeping a second model in memory, and answers with FSPO: the method freezes selected privileged predictions as sparse distributions and replays them later as fixed targets.

Key Contributions

1. Function-space temporal consolidation. Instead of storing parameters, FSPO stores the top-K next-token distributions produced by earlier privileged teachers (plus pooled tail mass) and replays them as fixed targets (L_temp). This gives temporal separation across optimization steps with a single resident policy and no persistent EMA teacher.
2. Privileged disagreement correction. Tokens where the unassisted student's log-probability exceeds the privileged teacher's by more than a margin γ are flagged as "locally overvalued" and explicitly penalized through L_neg, going beyond plain distribution matching.
3. Reliability-adaptive reference grounding. A smoothed estimate of per-problem rollout success sets a weight w_p that strengthens cross-entropy supervision on the verified reference solution when the policy repeatedly fails, and fades it out as the policy becomes consistently correct.
4. Quality- and coverage-aware temporal memory. Retained trajectories are selected by a retention score combining verifier quality with structural-signature coverage (coverage weight 0.5), and replay is sampled evenly across structural groups; targets are refreshed after H steps for a fraction ζ every F steps to avoid stale supervision.
5. A unified cyclic objective in which the weight ρ_s shifts from current privileged adaptation to historical consolidation within each cycle of length C.

Main Findings

• Best average accuracy at the 5-hour budget. On Qwen3-8B, FSPO reaches Avg@5h of 76.2 versus 72.0 for on-policy SDPO, 70.4 for GRPO, and 63.3 for on-policy GRPO. On Olmo3-7B-Instruct, FSPO reaches 72.4 versus 68.0 (SDPO), 63.2 (GRPO), and 60.8 (on-policy GRPO) — roughly a 4 percentage-point average gain over the strongest baseline.
• Best 5-hour result on every scientific reasoning task. On Qwen3-8B: Chemistry 83.4, Physics 76.0, Biology 71.2, Materials 81.7. On Olmo3-7B-Instruct: Chemistry 83.3, Physics 70.1, Biology 62.3, Materials 79.9.
• Tool use is the narrow margin. FSPO matches the strongest Qwen3-8B tool-use result at 5 hours (68.5) and improves on Olmo3-7B-Instruct (66.6 vs 65.0 GRPO / 62.1 SDPO), but SDPO is stronger at the 1-hour budget on both backbones (68.0 and 60.8 for SDPO vs 66.6 and 57.6 for FSPO).
• Faster convergence and cheaper steps. FSPO obtains most of its final gain early and reduces per-step cost over wall-clock time, because the policy increasingly produces shorter responses, whereas GRPO and SDPO generate progressively longer trajectories.
• Less catastrophic forgetting. On five prior-capability benchmarks, FSPO shows consistently higher retention than baselines while improving the trained task; the authors attribute this to replaying fixed historical predictions.
• Memory reduction. Estimated peak memory for Olmo3-7B-Instruct drops from about 70 GB (GRPO with a colocated rollout copy; SDPO with a persistent EMA teacher) to about 56 GB, since FSPO reuses the model being trained for both generation and privileged evaluation. The cache's peak overhead is about 5 MB. Projecting the same accounting, one eliminated bf16 parameter copy equals roughly 12 TB of aggregate memory at 6T parameters.
• Ablations show the components are complementary and interdependent. On Olmo3-7B-Instruct over Biology and Tool Use, the full objective averages 64.5, versus 56.0 for distillation only, 60.1 without temporal consolidation, 61.7 without reference grounding, and 54.1 without both temporal and reference terms — below distillation alone, implying aggressive negative correction only helps when the self-teacher is stabilized. Zero-shot base is 27.8.
• What is stored matters more than that something is stored. Sparse function-space targets (64.5) beat parameter snapshots every H steps (58.8), EMA teachers at μ=0.99 (63.2) and μ=0.999 (61.3), a frozen initial policy (56.5), and trajectory-level CE replay (58.5).
• Demonstration hierarchy and memory policy both matter. Verified sibling → retained trajectory → reference (64.5) beats reference-only (40.5), sibling-only (59.0), and dropping the retained trajectory (60.7). Quality+coverage retention (64.5) beats reward-only (61.9), coverage-only (57.0), uniform sampling (60.6), and random retention with uniform sampling (57.1).
• Adaptive grounding and scheduling beat fixed alternatives. Reliability-adaptive w_p (64.5) beats constant w_p=1 (62.6) and binary weighting (63.0). Cyclic adaptation→consolidation (64.5) beats fixed ρ=0.4 (62.8) and a monotonic ramp (62.5); reversing the cycle order drops to 59.7.

Methodology in Plain English

For each problem, the current policy samples G responses and scores them with a verifier. To build a teacher for a given rollout, the method picks the strongest available hint in a fixed order: another rollout of the same problem that verifiably succeeded, otherwise a previously retained successful trajectory, otherwise the ground-truth reference solution. The same model is then run twice over the same response prefix — once with just the question (student) and once with the hint appended (teacher) — and the teacher's output is treated as a fixed target.

Two efficiency tricks keep this cheap. First, the comparison only covers the student's top-K tokens plus a single "everything else" bucket, so the full vocabulary is never matched. Second, tokens where the student's confidence far exceeds the teacher's are treated separately and pushed down with a hinge penalty, rather than being averaged into the general distribution-matching loss.

Separately, the system keeps a memory. When a rollout is verified, the method stores its token sequence, score, a structural signature, and the teacher's top-K predictions over that sequence, then throws the teacher away. At later training steps, current responses are compared against those frozen stored predictions, which is what removes the moving-target problem: the target literally cannot move. Memories are chosen to be both high-quality and structurally diverse, replayed in balanced groups, and occasionally refreshed after a delay so they do not become useless.

Finally, for problems the policy keeps failing, the loss leans more on the verified reference solution, with that extra weight computed from a smoothed success rate and clipped to [0,1]. The overall loss blends current privileged distillation, historical consolidation, and reference grounding, with the balance cycling from the first toward the second within each cycle — memory is pre-seeded with reference trajectories so consolidation can start from step one.

Why This Matters

Impact on research. The paper reframes a stabilization problem as a memory problem: rather than adding a second parameter set to slow the teacher down (EMA, target networks, trust regions), it stores the teacher's outputs and lets the parameters move freely. It also connects LLM self-distillation to the older function-space / logit-replay literature in continual learning, and shows that dense token-level targets can double as a forgetting regularizer.

Real-world applications:

• Post-training reasoning and math/science assistants where verifiable answers exist but gradient-time memory is the binding constraint.
• Tool-using agents, where the paper's structural signatures (invoked action sequences) are used to preserve diverse calling strategies rather than one dominant pattern.
• Continual fine-tuning pipelines that must keep serving existing capabilities, where the reported retention advantage and the replayed historical targets apply directly.
• Memory-constrained or multi-tenant serving fleets, where removing one model-sized copy translates to large aggregate savings as parameter counts grow.

Industry relevance. Removing a persistent teacher copy is an operational win: one fewer model to hold resident, no EMA synchronization, and a cache measured in megabytes rather than gigabytes at the 7B scale tested. Any team running GRPO-style training at scale and paying for the rollout/teacher duplication has a concrete reason to look at this.

Future Directions

• Push the scaling argument past estimation. The 12 TB figure at 6T parameters is an accounting projection, not a measurement; validating memory and accuracy at much larger backbones is the obvious open test.
• Stress the 1-hour budget. SDPO beats FSPO at 1 hour on both backbones, so the regime where a persistent EMA teacher wins early in training — and whether warm-starting or a different cycle schedule closes it — is unresolved.
• Diagnose the negative-correction dependency. The ablation shows L_neg without temporal and reference anchors (54.1) performs worse than distillation alone (56.0). Understanding exactly when targeted suppression helps versus destabilizes would make the method safer to apply to new tasks.
• Extend beyond verifiable-reward domains. All evaluation uses verifier-scored tasks (SciKnowEval, ToolAlpaca); whether the retained-trajectory hierarchy and reliability weighting survive noisy or learned reward models is not reported.
• Tune the memory lifecycle. Retention threshold, coverage weight (fixed at 0.5), refresh age H, refresh fraction ζ, and cycle length C were not explored per dataset; adaptive versions of these are left open.

Target Audience

Researchers and engineers working on RL post-training and reasoning models who already understand GRPO-style objectives and want token-level dense supervision without a second model; continual-learning researchers interested in function-space replay at LLM scale; and infrastructure-minded practitioners concerned with the memory cost of on-policy distillation. Readers new to RLVR will find the experimental narrative readable but will need background in divergence-based distillation to follow the method section.

Authors’ abstract

Reinforcement Learning with Verifiable Rewards (RLVR) improves reasoning in language models, but scalar outcomes provide coarse supervision over response trajectories. Self-distillation offers denser token-level guidance by comparing an unassisted policy with a feedback-informed self-teacher, yet the teacher evolves with the student, creating unstable targets and motivating additional regularization. We introduce Function-Space Policy Optimization (FSPO), a single-policy framework that combines current privileged adaptation with function-space temporal consolidation. FSPO stores sparse historical next-token distributions as fixed temporal targets, providing temporal separation without maintaining a persistent teacher model. It further introduces privileged disagreement correction to explicitly suppress locally overvalued actions, reliability-adaptive grounding that strengthens reference supervision when self-generated evidence repeatedly fails, and quality and coverage-aware temporal memory that retains verified behaviors and improves privileged demonstrations. Across reasoning and tool-use benchmarks, FSPO achieves the highest average accuracy under a 5-hour training budget, improving over the strongest baseline by approximately 4% points on average while reducing peak memory from about 70 GB to 56 GB on Olmo3-7B-Instruct.