One-sentence Summary

Columbia University and Microsoft Research propose TRIPOST, a training algorithm that endows smaller language models with the ability to self-improve by having them interact with larger LLMs to collect feedback on their own generations and replay that experience for training, yielding up to $7.13\%$7.13% improvement for LLaMA-7B on math and reasoning tasks.

Key Contributions

• TRIPOST transfers self-improvement capability from large language models (LLMs) to smaller models by having the smaller model interact with an LLM to collect feedback and revisions on its own generations, then replaying that experience for training.
• The method constructs improving trajectories from model-LLM interactions, post-processes them, and trains the smaller model with weighted supervised learning to generate both useful feedback and corresponding improvements. LLM demonstrations are conditioned on the smaller model's generations, rather than expert-only examples or predefined rules.
• Experiments on four math and reasoning datasets from Big-Bench Hard show that TRIPOST improves LLaMA-7B performance by up to 7.13%. Analysis indicates that interactive learning from and correcting the model's own mistakes is essential for small models to improve task performance.

Introduction

The work addresses self-improvement for language models, where a model generates feedback on its own outputs and revises them to produce better answers. This is important for math and reasoning tasks, where self-improvement can make model outputs more precise and faithful, but prior work shows that smaller models often struggle to generate useful self-critique or learn effectively from demonstrations produced by much larger models. In particular, smaller models may make different types of errors than large models, so training them on large-model self-improvement demonstrations can be ineffective.

The authors contribute TRIPOST, an iterative training algorithm that helps a small model learn self-improvement by exploring with expert LLMs or Python tools, filtering and rebalancing the resulting training data, and training the model with weighted supervised learning. This approach focuses on learning from the small model’s own mistakes, rather than imitating large-model demonstrations that may not match the smaller model’s behavior. Evaluated on BIG-Bench Hard math and reasoning datasets, TRIPOST improves in-domain and out-of-domain performance compared with training on ground-truth rationales or direct LLM demonstrations.

Dataset

The authors build a dataset for stepwise reasoning and trajectory correction from BIG-Bench Hard.

• Source benchmark: BIG-Bench, a collection of more than 200 text-based tasks. The experiments use four BIG-Bench Hard math and reasoning tasks:

  • Multistep Arithmetic: scriptable task.
  • Word Sorting: scriptable task.
  • Date Understanding: unscriptable task.
  • Logical Deduction: unscriptable task.

• Gold rationale sources:

  • For scriptable tasks, step-by-step rationales and feedback are generated using scripts.
  • For unscriptable tasks, solutions are generated by prompting an LLM with CoT prompts from Suzgun et al. 2022.
  • Only correct LLM-generated rationales are kept for training.
  • For unscriptable tasks, feedback is also generated by prompting Codex/text-davinci-003.
  • Codex is used as the improvement module for all tasks.

• Data composition after interactive trajectory editing:

  • Gold step-by-step demonstrations.
  • Directly correct trajectories generated by the model.
  • Edited trajectories created using the model, feedback, and an improvement module.

• Processing into training examples:

  • If an attempt is incorrect, the data is split into single-step improvement triplets: attempt, feedback, improved attempt.
  • If an attempt is correct and the trajectory ends, the data becomes a terminal pair: attempt, final feedback containing the phrase "the final response is correct".
  • Expert corrections can replace original attempts and feedback, so edited attempts and edited feedback are used as labels.

• Filtering:

  • Improvement triplets are filtered out if they contain incorrect feedback or invalid improvement steps.
  • Filtering rules are described in the appendix.

• Dataset mixture:

  • Terminal correct examples and filtered improvement triplets are combined into one training dataset.
  • A hyperparameter p controls the proportion of improvement triplets in the mixture.
  • Too many improvement triplets cause the model to attempt self-improvement even when the previous attempt is already correct, harming performance.

• Train and evaluation split:

  • Each of the four tasks is divided into seen and unseen subtasks.
  • Simpler questions are mainly categorized as seen subtasks and used for training.
  • Unseen subtasks are used to measure generalization.
  • The split uses task-specific categorization shown in Table 2.

• No cropping strategy or additional metadata schema is specified in the provided text. Exact dataset sizes for each subset are also not reported.

Method

The authors build TRIPOST as an iterative training procedure that teaches a smaller language model $M_\theta$Mθ​ to learn from its own mistakes rather than from expert demonstrations that may not match its current capability. The motivating observation is that smaller models can struggle with direct self-improvement prompting or with imitating self-improvement traces generated by much larger models, especially on mathematical and logical tasks. As shown in the figure below:

To address this capability mismatch, the authors use the smaller model as the primary generator of attempts, then apply external feedback and targeted editing to create training signals that remain grounded in the smaller model’s behavior. Refer to the framework diagram:

The algorithm has three main modules: interactive trajectory editing, data post-processing, and weighted model training. Each iteration begins with the current model producing an initial attempt for a task. The framework then augments that attempt with feedback and correction, filters noisy or incorrect examples, and trains the model on the resulting supervised targets.

A trajectory is represented as a sequence of attempts and feedbacks,

where $x_0^{\mathrm{att}}$x0att​ is the initial attempt generated by $M_\theta$Mθ​, and $x_i^{\mathrm{fb}}$xifb​ and $x_i^{\mathrm{att}}$xiatt​ denote the $i$i-th feedback and updated attempt. A trajectory ends when the final feedback indicates that the answer is correct. Directly correct examples take the compact form

where the attempt is already correct and the feedback confirms correctness. Self-improving examples take the longer form

In the interactive trajectory editing stage, the feedback module FBK inspects an incorrect attempt and identifies the first error step. FBK may be implemented as an LLM-based critic or as a rule-based script, depending on the task. The feedback is intended to be precise and incremental, so the framework does not simply ask an LLM to rewrite the entire solution. Instead, it replaces the first feedback where the smaller model’s feedback and the external feedback disagree with an external feedback token:

Then the improvement module IMP uses the previous attempt and the corrected feedback to generate an improved attempt,

which is appended to the trajectory:

This editing is localized to the part of the solution after the first identified mistake. Earlier correct steps are preserved. This prevents the external module from silently producing a full expert solution and ensures that the supervision reflects a correction path the smaller model can plausibly learn. The process can repeat up to a maximum number of revisions until the final attempt becomes correct. If the edited trajectory does not reach the correct answer, it is discarded.

After trajectory editing, the algorithm collects three kinds of material: gold step-by-step demonstrations, directly correct attempts generated by $M_\theta$Mθ​, and edited self-improving trajectories. To make the supervised learning targets uniform, each piece of data is converted into either a single-step improvement triplet,

if the attempt was incorrect and required improvement, or a terminal correct example,

if the attempt is already correct. In the triplet form, the feedback or improvement step may come from the external FBK or IMP modules when an edit was applied.

The next stage filters and re-balances the collected data. Filtering removes improvement triplets whose feedback or improvement step is judged incorrect by heuristic rules. This is important because noisy feedback can teach the model to produce invalid corrections. The cleaned examples are then combined into a single training set, with a mixing hyperparameter $p$p controlling the proportion of improvement triplets $x_{\mathrm{imp}}$ximp​ relative to directly correct examples $x_{\mathrm{T}}$xT​. This balancing step is central to the training objective: if too many improvement examples are included, the model learns to revise its answer even when its initial attempt is already correct, which can reduce overall performance.

The final stage trains $M_\theta$Mθ​ with weighted supervised learning. The target sequence typically includes the question, the attempt, the feedback, and the updated attempt. The loss is weighted so that the model emphasizes learning how to critique and revise its own generations, while still learning from directly correct attempts. Formally, the training objective can be written as

where $q$q is the question and $y_t$yt​ is the $t$t-th supervised token. Tokens belonging to directly correct examples receive baseline weight $1.0$1.0. Tokens belonging to the feedback part $x_i^{\mathrm{fb}}$xifb​ or the improved attempt part $x_{i+1}^{\mathrm{att}}$xi+1att​ of an improvement triplet receive a larger weight $w > 1.0$w>1.0. The initial attempt in the triplet is kept in context but is not upweighted in the same way. This weighting encourages the decoder to learn the critical operations of identifying an error and producing a corrected continuation.

The authors also explore an automatic balancing variant called TRIPOST-auto. Instead of fixing the proportion $p$p of improvement triplets, this variant uses the natural distribution of generated examples: directly correct attempts become $x_{\mathrm{T}}$xT​, and corrected attempts become $x_{\mathrm{imp}}$ximp​. In this setting, $p$p effectively reflects how often the current model makes mistakes. Since the model is not artificially exposed to too many improvement cases, it tends to attempt self-improvement only when such a revision is more likely to be useful.

Experiment

Evaluations on four BIG-Bench Hard math and reasoning tasks with a LLaMA-7B student model compare TRIPOST against rationale fine-tuning, self-consistent rationale fine-tuning, and LLM demonstration fine-tuning. The results show that TRIPOST distills useful self-improvement behavior into the smaller model, yielding the best overall task performance while improving both initial attempt quality and measurable self-improvement contribution. Additional ablations and the TRIPOST-auto variant indicate that interactive feedback, data filtering, weighted objectives, and the balance of improvement demonstrations are important, with excessive self-improvement training hurting reliability and feedback generation becoming harder as the model improves.

Training smaller models with self-improvement demonstrations from much larger models can be ineffective because the two model scales tend to make different kinds of mistakes. The example shows a smaller model making simple copying or arithmetic slips, while a larger model introduces a different step-level arithmetic error. This capability mismatch suggests that self-improvement demonstrations may not provide a well-matched learning signal for smaller models. Smaller models and larger models can fail through different types of errors on the same math example. Self-improvement demonstrations from larger models can reflect errors that smaller models do not naturally make. Training smaller models on self-improvement demonstrations from larger models can be less reliable than training on ground-truth rationales.

The dataset categories are organized by task complexity, with easier instances designated as seen subtasks for training and more complex instances held out as unseen subtasks. The split spans arithmetic, sorting, date reasoning, and deduction, using task-specific measures such as operation depth, word count, reasoning steps, and option count. This design supports evaluation of generalization to harder variations beyond the training distribution. Seen subtasks correspond to simpler question types, while unseen subtasks use higher complexity ranges. The criteria change by task, comparing structural difficulty in math, sorting, temporal reasoning, and logical choice tasks.

Across the four BIG-Bench hard datasets, TRIPOST-trained models outperform LMSI, ground-truth rationale fine-tuning, and LLM self-improvement demonstration baselines on both seen and unseen tasks. The results generally improve as more self-improvement iterations are added, suggesting that learning from the model's own mistakes strengthens both initial attempts and refinement. The self-improvement demonstration baseline often performs worse than the ground-truth rationale baseline, consistent with a capability gap between the large demonstrator and the smaller learner. TRIPOST achieves the strongest overall performance across all four datasets. Performance tends to increase with additional self-improvement iterations. Unseen-task accuracy is usually lower than seen-task accuracy across methods. Using only LLM self-improvement demonstrations reduces performance relative to ground-truth rationale fine-tuning.

TRIPOST-trained models show positive self-improvement contributions across all four task types, with gains appearing on both seen and unseen subtasks. The largest self-improvement contribution is observed for logical deduction, especially on seen cases, while the other tasks show smaller but consistent gains. Overall performance reflects a combination of directly correct initial attempts and measurable improvement contributions. Self-improvement contributes positively on every task. Logical deduction benefits most from self-improvement. Seen subtasks generally contribute more to self-improvement than unseen subtasks. Directly correct attempts remain a major part of overall accuracy.

TRIPOST-auto automatically weights edited improvement data by current task performance, producing training mixtures with fewer self-improvement examples and more directly correct attempts. It reduces the model's tendency to edit its own attempts, yet often gives stronger overall task performance than fixed re-balancing TRIPOST, especially on less scriptable tasks. TRIPOST-auto shows little or no measurable self-improvement activity across tasks. Its overall performance generally improves across iterations and often exceeds fixed-balancing TRIPOST, particularly on word sorting and logical deduction.

The experiments use difficulty-based seen and unseen splits across arithmetic, sorting, date reasoning, and deduction to evaluate generalization beyond the training distribution. They show that self-improvement demonstrations from much larger models can be poorly matched to smaller learners’ typical mistakes, whereas TRIPOST-trained models outperform LMSI, ground-truth rationale fine-tuning, and LLM self-improvement demonstration baselines on both seen and unseen tasks. Additional iterations generally improve performance and positive self-improvement contributions, especially for logical deduction, while the auto-balanced variant reduces explicit editing behavior but often achieves stronger overall accuracy on less scriptable tasks.