Specializing Smaller Language Models towards Multi-Step Reasoning

In plain words

Large language models can work through a maths word problem step by step, and this skill was long thought to appear only in very big models. The authors ask whether a small model can learn it too. Their recipe: a big model writes worked solutions to school-level maths questions, the solutions that reach the right answer are kept, and a small pretrained model is trained to reproduce them.

The small model gets better at maths, but it pays a price. It forgets much of how to answer general reasoning questions outside maths. Its maths skill also does not transfer equally well to differently worded problems, so the authors pick the saved version of the model by testing it on other maths datasets rather than on the training data.

The authors also report that the small models' ability to reason step by step improves smoothly with size, instead of staying flat until a sudden jump. Their results are self-reported and come from one family of models, so they are best read as evidence for a trade-off, not as a general recipe.

TL;DR

The paper distills chain-of-thought maths solutions from code-davinci-002 (a GPT-3.5-class model) into T5 and FlanT5 students from 250M to 11B parameters. The training data is GSM8K, augmented with 40 teacher solutions per question and filtered for correctness. Four data formats are mixed: in-context and zero-shot, each with or without chain-of-thought. The main objective is distribution matching on the teacher's top-5 token probabilities, with a dynamic-programming alignment between the two tokenizers. Checkpoints are chosen on held-out maths datasets. Specialized FlanT5-XXL gets 27.1% on GSM8K (author-reported, against 16.1% before specialization), but its BIG-Bench Hard chain-of-thought accuracy falls to 0.0 tab. 1. The official code is MIT-licensed, but it differs from the paper in several places (see Paper vs code).

Problem

Chain-of-thought prompting was reported to work only in models above roughly 100B parameters, with flat, near-zero scaling below that §1 §2. Instruction tuning later showed that smaller models can also produce chain-of-thought when trained on such data, but their results stayed below those of large models §1 §2. Distilling from a large teacher was an existing route, and concurrent work used it with raw T5 or GPT-style students §2 §4.1. The paper's stated concern is that a small model with limited capacity may overfit the tuning distribution and generalize poorly when the problem wording changes §2. Assumptions: the teacher's per-step probabilities are only available for its top-5 tokens §3.

Method

Specialization. The student is trained only on maths chain-of-thought data, so that its limited capacity goes to one target ability. The price is a loss of generic ability, measured on BIG-Bench Hard §1 §4.1.

Data. Start from the GSM8K training questions. For each question, ask the teacher for 40 solutions and keep those that reach the correct answer §3 §4.1. Each kept solution is an answer plus a chain of thought. Four formats are built from these §3 fig. 1B:

1. in-context answer-only: four solved examples are prepended, and only the final answer is the target;
2. in-context chain-of-thought: four solved examples are prepended, and the target is the full chain of thought plus answer;
3. zero-shot answer-only: the question alone is the input, and the target is the answer;
4. zero-shot chain-of-thought: the question alone is the input, and the target is the full chain of thought plus answer.

The paper mixes all four formats, because in-context data keeps in-context learning and zero-shot data improves zero-shot answering §3 §4.4. The mixing ratio is not stated in the paper.

Objective. The paper compares two options §3. Sample matching maximizes the likelihood of the teacher's written solutions. Distribution matching minimizes the KL divergence between the student's per-step distribution and the teacher's per-step distribution. The paper uses distribution matching, and notes that it needs the full teacher distribution at every step, which makes stored data about |V| times larger than with sample matching §3. The teacher only reveals its top 5 tokens at each step. The paper keeps those five probabilities and sets all other tokens to zero §3.

Tokenizer alignment. The teacher (GPT tokenizer) and the student (T5 tokenizer) split text differently, so per-step distributions cannot be compared directly. The paper aligns the two token sequences with an edit-distance dynamic program, in the style of sequence alignment, where one token may map to several tokens on the other side §3 eq. 1 eq. 2 eq. 3. If a teacher token and a student token pair one-to-one, the teacher's distribution is used. If the mapping is one-to-many or many-to-one, the student's distribution is set to one-hot §3 fig. 1C.

Model selection. The paper holds out maths datasets (MultiArith, ASDiv, SVAMP, called M-A-S) and selects checkpoints by their accuracy there, instead of by GSM8K validation accuracy. In-distribution and out-of-distribution accuracy can move in opposite directions during training §4.3 tab. 3.

# Build distillation targets for one GSM8K question  [§3] [§4.1]
# Input:  question with its gold answer, teacher code-davinci-002
# Output: correct teacher solutions with top-5 probabilities per step

def build_targets(question, gold_answer, teacher):
    # Ask the teacher for 40 solutions per question  [§4.1]
    solutions = [teacher.generate(question) for _ in range(40)]
    # Keep only solutions that reach the gold answer  [§3] [§4.1]
    kept = [s for s in solutions if final_answer(s) == gold_answer]
    # The teacher exposes its 5 most likely tokens per step  [§3]
    per_step_top5 = [top5_per_step(s) for s in kept]
    return kept, per_step_top5

# Align teacher tokens s[1..m] with student tokens t[1..n]  [§3] [eq. 1-3]
# Input:  s (teacher tokens), t (student tokens)
# Output: cost table f, from which the alignment path is read back

def align(s, t):
    # f(i, j) is the cheapest alignment of the first i and j tokens  [§3]
    for i in range(1, len(s) + 1):
        for j in range(1, len(t) + 1):
            # Substitution cost between the two tokens  [§3] [eq. 1-3]
            cost = edit_distance(s[i - 1], t[j - 1])
            # Recurrence: the minimum of three moves  [eq. 1] [eq. 2] [eq. 3]
            f[i][j] = min(
                f[i - 1][j] + cost,
                f[i][j - 1] + cost,
                f[i - 1][j - 1] + cost,
            )
    # Read the alignment path back from f  [§3]
    return backtrack(f)

# Distribution-matching loss on chain-of-thought positions  [§3]
# teacher_dist: (B, T, V) with the top-5 probabilities, zero elsewhere
# mask:         (B, T), 1 for answer positions

def distribution_matching_loss(student_logits, teacher_dist, mask):
    # Student's per-step distribution over the vocabulary  [§3]
    student_probs = softmax(student_logits)
    # KL(teacher || student) at every step  [code: src/utils.py#L632-L640]
    kl = sum(teacher_dist * log(teacher_dist / student_probs), over V)
    # Average over unmasked positions  [code: train_distill_simple.py#L30-L38]
    return sum(kl * mask) / sum(mask)

Under the distribution option, answer-only batches have no per-step teacher distribution, so the released code trains them with ordinary cross-entropy on the teacher's text instead train_distill_simple.py. The paper does not describe this split.

Hyperparameters

Values marked "code" are not stated in the paper. They come from the released repository at code_commit, and the paper's own run settings may differ. The paper contains no appendix [paper source].

Results

Author-reported. Setup: GSM8K, MultiArith, ASDiv and SVAMP for maths accuracy; BIG-Bench Hard (26 tasks) for generic ability, in answer-only (AO) and chain-of-thought (CoT) prompting §4.1. Baselines are code-davinci-002, LaMDA 137B, PaLM 60B, UL2 20B, and three distillation works §4.1. The paper says its comparisons with the distillation works have several confounders, including the size and format of the tuning data §4.1. Reported values are test-set accuracy tab. 1.

Source: tab. 1. Deltas are the paper's own, against the unspecialized FlanT5 row. The "not stated" entries are shown as "?" in the paper's table.

Further results §4:

• The paper states that, on GSM8K, its 3B specialized model beats the 11B and 6B distillation models it compares with §4.1 tab. 1.
• Starting from an instruction-tuned checkpoint (FlanT5) gave higher GSM8K validation accuracy than starting from raw T5 at 250M, 760M and 3B tab. 2. The paper names this as a major reason its FlanT5 results beat the concurrent distillation work, which used raw T5 §4.2.
• GSM8K validation accuracy for specialized models, by start point: FlanT5 3B 23.8 vs T5 3B 20.6; FlanT5 760M 21.8 vs T5 760M 16.2; FlanT5 250M 15.2 vs T5 250M 14.2 tab. 2.
• Selecting on M-A-S dev instead of GSM8K dev lowers in-distribution accuracy by 2.0 to 2.6 points and raises out-of-distribution accuracy by 0.3 to 1.8 points for FlanT5 3B, Large and Base tab. 3.
• Scaling: across model sizes, the paper's Figure 2 plots GSM8K validation accuracy against model size on a log axis, and the authors describe the specialized curves as log-linear rather than flat fig. 2 §4.2.
• Training dynamics: the paper's Figure 3 shows that BBH-CoT accuracy is 0 from the first checkpoints on, and in-distribution and out-of-distribution accuracy rise and then fluctuate fig. 3 §4.3.
• Data format: Figure 4 compares checkpoints tuned on in-context data with those tuned on zero-shot data, on GSM8K validation with in-context and zero-shot prompting. The paper's text says in-context tuning keeps both abilities and zero-shot tuning loses in-context ability fig. 4 §4.4.
• Distribution matching converges faster and to a lower loss than sample matching. The two give similar validation accuracy fig. 3 §4.4.

Paper vs code

Repository searched: the paper and abstract contain no code link, and the arXiv metadata has none. A GitHub repository search for the title returned FranxYao/FlanT5-CoT-Specialization, owned by Yao Fu (a listed author), described as the implementation of this paper, and MIT-licensed. Cloned at commit a8ffd51afa4f88b9a8fe160c093fb891af5a1cc7. Discrepancies:

• Training objective default. The paper's main method is distribution matching §3. The default in the config is loss_type: match_sample, which trains with cross-entropy on the teacher's text config.yaml. The distribution option uses KL only on chain-of-thought batches and cross-entropy otherwise train_distill_simple.py. The code does not say which option produced the reported numbers.
• Tokenizer alignment. The paper describes an edit-distance dynamic program §3 eq. 1-3. In src/utils.py the dynamic-programming function is a stub that returns nothing and is marked "TBC" utils.py. The mapping that is implemented, transform_codex_token_to_t5_token, maps each teacher token to its closest student-vocabulary token by edit distance, and inserts a blank step where no match exists utils.py utils.py. A dynamic-time-warping alignment appears only in a development notebook dev_align_codex_to_flan_t5_dtw.ipynb. The paper does not say which alignment produced its results.
• Unmatched positions. The paper sets the student distribution to one-hot for non-one-to-one matches §3. The code inserts a zero-probability blank step and records a mask value of 0 utils.py. That mask file is named in the data loader but never read data_utils.py. Blank steps therefore stay in the averaging denominator with zero KL data_utils.py.
• Pre-processing not included. The loader reads pickles from processed_data/ data_utils.py. That folder is not in the repository, so the teacher generation, filtering and alignment cannot be rerun from the repository alone.
• Top-5 probabilities. The paper zeroes the other tokens §3. The code copies the five probabilities without renormalizing them data_utils.py.
• Gradient clipping. The config sets 5.0 config.yaml, but only the Lightning script applies it train_distill_lightning.py. The simple training script does not clip.
• Schedule length. The cosine schedule is given 10 × micro-batches as its total train_distill_simple.py, but it steps once per optimizer step, that is once every 15 micro-batches train_distill_simple.py. See Pitfalls.
• Epoch count. The paper's runs used 2 and 3 epochs fig. 3. The default config uses 10 config.yaml.
• Checkpoint unit. The paper says a checkpoint every 10K instances or updates §4.3. The code saves every 10,000 micro-batches config.yaml train_distill_simple.py.
• Evaluation decoding. The paper does not describe it. The test script passes only a maximum length test_distill.py, while a separate decoding script samples with do_sample=True and 256 new tokens flan_t5_decode_gsm8k.py. The predicted answer is the last number in the output utils.py.
• Negative samples. The unlikelihood path for negative chains is commented out train_distill_simple.py and is not used by default.

Pitfalls

• Learning-rate schedule barely decays. From the code (not stated in the paper): the cosine total is 10 × the number of micro-batches, but the scheduler advances once per 15 micro-batches. With 10 epochs, the schedule reaches only about 1/15 of its span, so the learning rate at the last step is roughly 0.99 × its peak (5e-4). If you want a full decay, set the total to the actual number of optimizer steps.
• Top-5 rows may not sum to 1. The paper says the five probabilities usually sum to close to 1 §3. Nothing in the repository checks this, and the code does not renormalize, so the target is not guaranteed to be a distribution. Decide explicitly whether to renormalize.
• KL direction. The paper says the student is trained to match the teacher's distribution §3 without fixing the direction. The code computes KL(teacher ‖ student) utils.py, which is the forward direction used in standard distillation.
• Memory. The code builds a dense float tensor of shape (micro-batch, steps, vocabulary) for every chain-of-thought batch, with vocabulary 32128 data_utils.py data_utils.py. As an illustration (my arithmetic, not from the paper), a micro-batch of 2 with 256 target steps needs about 66 MB for this tensor in float32 alone.
• Losing generic ability is near-total. Every specialized model reaches 0.0 on BBH-CoT tab. 1. If the target model must keep general reasoning, this method does not fit.
• Model selection trade-off. Better GSM8K validation accuracy does not imply better out-of-distribution accuracy tab. 3 §4.3. Choose the selection set to match the goal.
• Non-standard splits. The test sets are random remainders of the data (800 GSM8K items, for example) §4.1. Numbers are not comparable with official leaderboards.
• Answer extraction. The last number in the output is the prediction utils.py. This is a simple rule that can mark a correct answer wrong when the chain ends with another number.
• Confounded comparisons. The paper says its comparison with concurrent distillation work has multiple confounders, including the size and format of the tuning data §4.1. Treat the gaps as indicative only.
• Unstated costs. The teacher API cost, the number of kept solutions, and training time are not stated in the paper.

Lineage

• Builds on: chain-of-thought prompting, arXiv:2201.11903 §1; instruction-tuned FlanT5, arXiv:2210.11416 §2; the GSM8K dataset, arXiv:2110.14168 §3.
• Compared against: Magister et al., arXiv:2212.08410; Shridhar et al., arXiv:2212.00193; Ho et al., arXiv:2212.10071 §4.1.
• Superseded by: not stated in this wiki.