# Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

arXiv: https://arxiv.org/abs/2602.11220 · HTML: https://mlsanity.com/paper/2602.11220/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Web, search, and research agents, one of 16 directions in LLM agents and tool use.
- Builds on 33 of its 58 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Placed in

- LLM agents and tool use / [Web, search, and research agents](https://mlsanity.com/agents/web-search-and-research-agents/)

## Builds on (most cited first, 30 of 33)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment](https://mlsanity.com/paper/2303.16634/) (2023)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Gradient Episodic Memory for Continual Learning](https://mlsanity.com/paper/1706.08840/) (2017)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models](https://mlsanity.com/paper/2403.13372/) (2024)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [Fine-tuning Aligned Language Models Compromises Safety, Even When Users Do Not Intend To!](https://mlsanity.com/paper/2310.03693/) (2023)
- [AGIEval: A Human-Centric Benchmark for Evaluating Foundation Models](https://mlsanity.com/paper/2304.06364/) (2023)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [ProofWriter: Generating Implications, Proofs, and Abductive Statements over Natural Language](https://mlsanity.com/paper/2012.13048/) (2020)
- [Transformers as Soft Reasoners over Language](https://mlsanity.com/paper/2002.05867/) (2020)
- [How Abilities in Large Language Models are Affected by Supervised Fine-tuning Data Composition](https://mlsanity.com/paper/2310.05492/) (2023)
- [AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset](https://mlsanity.com/paper/2504.16891/) (2025)
- [Understanding Catastrophic Forgetting in Language Models via Implicit Inference](https://mlsanity.com/paper/2309.10105/) (2023)
- [Mitigating Catastrophic Forgetting in Large Language Models with Self-Synthesized Rehearsal](https://mlsanity.com/paper/2403.01244/) (2024)
- [Verbalized Sampling: How to Mitigate Mode Collapse and Unlock LLM Diversity](https://mlsanity.com/paper/2510.01171/) (2025)
- [Revisiting Catastrophic Forgetting in Large Language Model Tuning](https://mlsanity.com/paper/2406.04836/) (2024)
- [On-Policy RL Meets Off-Policy Experts: Harmonizing Supervised Fine-Tuning and Reinforcement Learning via Dynamic Weighting](https://mlsanity.com/paper/2508.11408/) (2025)
- [On the Algorithmic Bias of Aligning Large Language Models with RLHF: Preference Collapse and Matching Regularization](https://mlsanity.com/paper/2405.16455/) (2024)
- [InsCL: A Data-efficient Continual Learning Paradigm for Fine-tuning Large Language Models with Instructions](https://mlsanity.com/paper/2403.11435/) (2024)
- [The Best Instruction-Tuning Data are Those That Fit](https://mlsanity.com/paper/2502.04194/) (2025)
- [The Price of Format: Diversity Collapse in LLMs](https://mlsanity.com/paper/2505.18949/) (2025)

## Abstract (arXiv)

Large language models are commonly adapted to downstream tasks through supervised fine-tuning (SFT), but substantial distribution mismatch between downstream supervision and a model's generation distribution can intensify catastrophic forgetting. Data rewriting offers a data-centric way to narrow this mismatch before SFT. Existing methods, however, typically sample rewrites from a prompt-induced conditional distribution, which need not align with the backbone's natural question-answering generation distribution, and fixed templates can reduce output diversity. We formulate data rewriting as a policy-learning problem and train a lightweight LoRA rewriting policy with reinforcement learning. The policy optimizes question-answering-style distributional alignment and semantic diversity under a hard task-consistency gate, producing verified supervision for downstream SFT. Across three instruction-tuned backbones, the resulting models attain downstream gains broadly comparable to standard SFT while reducing degradation on non-downstream benchmarks in every evaluated setting. Additional experiments on logical reasoning and medical question answering provide preliminary evidence that a rewriting policy can be reused across domains for the same backbone.