R e A lign: Structured Revision for Small Language Model Alignment (2025)

Abstract: Aligning small language models with human preferences is challenging, as weak policies struggle to generate informative on-policy samples and suffer from unstable gradients when trained on off-policy signals from stronger models. In this work, we propose ReAlign, a training framework that combines the stability of on-policy learning with the guidance of reviser-assisted supervision. In the ReAlign, we first train a lightweight reviser to improve policy-generated responses using preference-based supervision, conditioned on both the prompt and the initial output. And then, the policy is optimized using a combination of standard on-policy preference pairs and reviser-enhanced pairs constructed as a structured revision task, where the latter provide richer, more learnable feedback. Experimental results on AlpacaEval-2 and Arena-Hard demonstrate that ReAlign significantly boosts alignment performance for weak policies, outperforming strong preference optimization baselines.

## introduction aligning small language models (slms) with human preferences is challenging, especially when using on-policy methods such as reinforcement learning from human feedback (rlhf) @xcite . these methods rely on model-generated outputs; however, small models often produce uniformly low-quality responses in the early training, resulting in weak and noisy supervision signals. two broad strategies have been explored to address this problem. the first involves supervised fine-tuning (sft) on high-quality, humanannotated or model-generated data to bootstrap the model's performance, followed by on-policy optimization @xcite . the second strategy collects high-quality preference data from stronger models and employs off-policy optimization methods such as direct preference optimization (dpo) @xcite @xcite . however, since these responses are sampled independently of the weak model, they often exhibit a distributional mismatch. to mitigate this, an intermediate sft step is commonly applied to reduce the gap and enhance the effectiveness of subsequent off-policy learning @xcite . figure 1a illustrates how the two training routes interact with model scale. for smaller models (e.g., 1b), off-policy training accounts for the majority of performance gains, while on-policy updates contribute only marginal improvements. however, as model capacity increases from 1b to 3b and then to 8b, the incremental benefit of on-policy optimization becomes more pronounced, with the largest gains ultimately observed at 8b. in other words, off-policy data provide high-reward signals that weak models cannot generate on their own, while on-policy optimization grows increasingly effective and reinforces the model's existing strengths @xcite . these observations naturally raise the question: can we integrate the strengths of on-policy and offpolicy methods into a unified framework? unfortunately, directly combining on-policy and off-policy supervision within a standard dpo objective can lead to unstable optimization. as shown in figure 1b , off-policy pairs often lie far outside the policy's current distribution and exhibit large negative log-probability margins and higher variance than on-policy. the corresponding loss term therefore dominates the aggregate gradient, swamping the on-policy signal and driving the model toward regions it cannot reach @xcite . in contrast, on-policy training yields small, positive margins that align more closely with the model's generation behavior @xcite . to address this incompatibility, we introduce a revision task into the policy training process. rather than forcing the policy to imitate responses from strong models, which often lie outside the policy model's distribution, we train the policy to revise its own responses toward improved alternatives. specifically, the policy takes both a prompt and its self-generated response as input, and learns to prefer stronger candidates over weaker ones. by anchoring the learning signal to the model's own outputs, the revision task ensures distributional consistency with on-policy data while enabling supervision from higher-quality responses (see figure 1b ). this setup allows preference information to be integrated more smoothly, even when the initial generations are suboptimal. the revision task serves two key purposes. first, it enables context-aware preference learning by conditioning on the initial response. second, it provides more informative supervision in early training, when the model's outputs are weak and standard on-policy learning lacks useful gradients. our preliminary experiments (figure 1c ) show that training solely on the revision task can yield substantial improvements, suggesting its potential as an effective alignment objective. then the question becomes: how do we obtain the improved responses required for the revision task? a direct approach is to reuse high-quality answers from a strong model. however, these responses are generated independently of the policy and often diverge in style and structure. when paired with self-generated outputs, they tend to produce inconsistent log-probability margins, which destabilize training and reduce the effectiveness. to mitigate this, we introduce a lightweight reviser model that generates improved responses conditioned on the policy's initial output. compared to raw strong-model responses, these revisions are closer to the original generation while maintaining high quality. empirically, preference pairs constructed from such revisions exhibit smaller and more consistent log-probability margins (figure 1b ), and resemble the low-gap, high-quality pairs studied by @xcite . building on this design, we propose realign, a unified training framework that combines on-policy learning with reviser-guided off-policy supervision. re-align trains the policy using both self-sampled preference pairs and revision-task pairs generated by the reviser, enabling more stable and effective alignment, especially for small language models. we evaluate realign on the ultrafeedback dataset @xcite by fine-tuning several small policy models using alignment methods such as sft, dpo, simpo @xcite , wpo @xcite , and simplemix @xcite . the aligned models are then assessed on the alpacaeval-2 @xcite and arena-hard @xcite benchmarks. experimental results demonstrate that realign effectively enhances the ability of slms to generate high-quality responses, making it a promising approach for aligning slms with human preferences. ## methodology we begin by outlining the preliminaries of preference-based alignment. as shown in figure 2 , we then introduce realign, a unified framework the reviser is first trained to refine low-quality responses into high-quality ones using paired responses. in the second stage, the policy is optimized using both its own responses and the revised outputs from the reviser. this dual approach enables continuous improvement of the policy from both self-generated and reviser-enhanced responses. ## main results on-policy vs. off-policy optimization we first compare on-policy and off-policy optimization across all policy models. on alpacaeval-2, offpolicy optimization consistently outperform onpolicy alternatives. for example, dpo-off improves the lc win rate from 33.23% to 45.67% on llama-3.2-3b-instruct, supporting the hypothesis that weak models struggle to generate high-quality on-policy data early in training. in contrast, offpolicy data provides stronger guidance with highquality responses from stronger models. however, on arena-hard, qwen2.5-3b-instruct achieves better performance with on-policy dpo (38.2%) than with off-policy dpo (35.3%). this suggests that on-policy optimization becomes more effective when the policy is sufficiently capable or better aligned with the target distribution, possibly as a result of pretraining differences. this observation aligns with the findings of @xcite . ## ablation studies impact of the reviser reviser outputs yield better performance than strong-model responses. for instance, r-d outperforms s-d (48.00% vs. 45.67%), indicating that reviser-generated responses are both higher quality and more learnable for weak policies. similarly, p-d + r-a surpasses p-d + s-a (50.17% vs. 45.53%), showing that revision supervision is more effective when conditioned on the policy's own outputs. this policy-awareness keeps the learning signal within the model's reachable distribution. model version sft realign (dpo) wr (%) avg. l wr (%) avg. l llama-3b initial 50.0 2137 84.2 2073 revised 81.2 2058 83.9 2130 qwen-3b initial 48.8 1872 74.9 2086 revised 66.2 1981 75.8 2228 llama-1b initial 20.6 2043 59.1 2186 revised 44.6 2130 56.0 2149 ## conclusion in this work, we presented realign, a novel framework for aligning small language models that combines the stability of on-policy learning with the guidance of reviser-assisted supervision. by introducing a revision task grounded in the model's own generations, realign mitigates the distributional mismatch issues commonly seen in off-policy preference optimization and provides richer, more stable learning signals during early training. the integration of a lightweight reviser enables the creation of high-quality, preference-aligned revisions that remain close to the policy's distribution, allowing small language models to benefit from stronger supervision without destabilizing gradients. our empirical results on challenging benchmarks such as alpacaeval-2 and arena-hard demonstrate that re-align consistently outperforms existing baselines across various small model sizes. these results highlight the importance of conditioning supervision on the model's own behavior and suggest that structured revision-based training offers a scalable and effective strategy for preference alignment.