PUST: Decoupling LLM Post-Training for Reusable & Efficient Alignment

TL;DR

Post-training just got modular. Proxy-guided Update Signal Transfer (PUST) separates the expensive "trial-and-error" phase of RLHF from the actual model update. By using a small model (proxy) to find the right direction and transferring that "update signal" to a larger model, PUST makes post-training asynchronous, cacheable, and remarkably cost-effective.

The Bottleneck: The Coupling of Exploration and Alignment

In standard Reinforcement Learning from Human Feedback (RLHF), your primary model — say, a 70B parameter beast — has to handle everything. It generates thousands of tokens (exploration) and then updates its weights based on rewards (alignment). This is:

1. Expensive: Sampling from a 70B model is slow and GPU-intensive.
2. Fragile: If you want to train a different model, you have to do it all over again.
3. Monolithic: You cannot "save" the progress of exploration and apply it elsewhere.

The authors of PUST argue that reward-oriented optimization and distribution alignment serve distinct roles and shouldn't be trapped in the same loop.

The PUST Framework: Signal over Distribution

The core insight of PUST is that we shouldn't distill the absolute distribution of a teacher (which leads to the capability bottleneck of the teacher). Instead, we should extract the relative improvement signal — the "delta" showing which tokens were encouraged or suppressed during RL.

The 3-Step Workflow

1. Proxy Exploration: A lightweight proxy (e.g., 1.7B) undergoes GRPO to discover high-reward behaviors.
2. Signal Extraction: We calculate the log-ratio difference between the base proxy and the optimized proxy. This is the "Directional Update Signal."
3. Signal Transfer: This signal is "pasted" onto the large primary model’s space.

The Math of Transfer

To prevent the primary model from over-updating (drifting too far from its original capabilities), PUST introduces Anchor-Based Update Calibration. The loss function effectively minimizes the KL divergence between the primary model and a dynamic target distribution created by the proxy's signal, adjusted by a calibration coefficient $\lambda$.

${\mathcal{L}}_{\mathrm{proxy}}(heta)=-\mathbb{E}\left[\sum {\pi }_{heta}\left(\log \frac{{\pi }_{\varphi }^{+}}{{\pi }_{\varphi }}-\lambda \log \frac{{\pi }_{heta}}{{\pi }_{\mathrm{ref}}}\right)\right]$

Does it Work? Weak-to-Strong Magic

The most striking result is the efficacy of transferring signals from a "weak" proxy to a "strong" primary model. In the Math domain, a 4B proxy boosted an 8B primary model by over 30 points on average across AIME benchmarks.

Key Finding: The primary model often achieves higher accuracy than the proxy that guided it. This suggests that the signal acts as a "trigger" for latent capabilities already present in the larger model.

Reusability and the "Cache" of Intelligence

Because the update signal is decoupled, PUST allows for:

• Asynchronous Generation: Collect signals once, use them forever.
• Cross-Model Transfer: A signal from a Qwen-4B can improve a Qwen-1.7B and a Qwen-8B simultaneously.
• Agnostic Scaling: You can tune the "intensity" of the feedback by simply changing the $\lambda$ coefficient without re-running the RL loop.

Critical Insight: The Modular Future of RLHF

PUST fundamentally redefines training data as "environmental information" that can be compressed into a "transferable update signal."

Limitations:

• The Scale Gap: If the gap between the proxy and primary is too large, the signals become noisy.
• Static Delta: The signal is currently fixed after extraction. Future work on "Adaptive Calibration" (making $\lambda$ dynamic based on token-entropy) will be the next frontier.

Conclusion

PUST is a significant step toward making LLM post-training a modular engineering discipline. By treating improvement signals as cacheable assets, it lowers the barrier for high-performance reasoning optimization, moving us away from monolithic, "all-or-nothing" RLHF runs toward a more flexible, scalable paradigm.