Step Rejection Fine-Tuning: A Practical Distillation Recipe
Igor Slinko, Ilia Zavidnyi, Egor Bogomolov, Yaroslav Zharov
TLDR
Step Rejection Fine-Tuning (SRFT) improves LLM agent training by leveraging partially correct, unresolved trajectories, outperforming standard RFT.
Key contributions
- Introduces Step Rejection Fine-Tuning (SRFT) for LLM agents to utilize unresolved trajectories.
- SRFT employs a critic LLM to assess step correctness and mask loss for erroneous steps.
- Allows models to learn from errors in context without reproducing them, fostering recovery.
- Achieves a 3.7% higher resolution rate on SWE-bench Verified than RFT, reaching 32.2%.
Why it matters
Standard LLM agent training discards valuable, partially correct data. SRFT offers a practical method to leverage these unresolved trajectories, significantly boosting performance on complex tasks like SWE-bench. This approach allows models to learn from mistakes, making them more robust and efficient.
Original Abstract
Rejection Fine-Tuning (RFT) is a standard method for training LLM agents, where unsuccessful trajectories are discarded from the training set. In the context of SWE-bench tasks, this corresponds to filtering out runs where the submitted patch does not pass the tests. However, this approach discards unresolved trajectories, even though they form a large portion of all trajectories for hard tasks and even then may be partially correct. In this work, we propose Step Rejection Fine-Tuning (SRFT) - a practical way to leverage these unresolved trajectories. For this, we employ a critic LLM to assess the correctness of each step in a trajectory. Consequently, during training, we mask the loss for erroneous steps while retaining them in the context window. This way we ensure the model learns to recover from errors without reproducing them. Evaluation on SWE-bench Verified shows that while RFT improves the resolution rate by 2.4% by excluding unresolved trajectories, SRFT improves it by 3.7% by filtering them instead of discarding completely, reaching the total resolution rate of 32.2%.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.