CoTranslator: Reverse to Forward Chain-of-Thought Translator for Reasoning Distillation

AI summary

Overview

Research area: Natural Language Processing — specifically reasoning distillation, the practice of training smaller "student" language models on chain-of-thought (CoT) reasoning traces produced by larger "teacher" models.

Technical level: Intermediate. The paper is readable with a general machine-learning background, but its core framing (casting a data-generation mismatch as a machine-translation problem) and its appendix of distribution-recovery bounds assume familiarity with probabilistic modeling and KL divergence.

Scope: The paper proposes CoTranslator, a small fine-tuned model that rewrites answer-conditioned ("reverse") teacher reasoning into question-only ("forward") reasoning, so that student training data can be produced cheaply without paying for large-scale rejection sampling.

What This Paper Is About

Reasoning distillation needs chain-of-thought data that ends in a correct answer, and the standard way to get it — let the teacher reason from the question alone, then discard wrong answers (rejection sampling) — is expensive because hard problems need many teacher attempts. The cheaper alternative, giving the teacher the answer up front and asking it to explain (reverse CoT generation), produces rationalizations that leak or presuppose the answer, creating a mismatch with the student's answer-free inference setting. The paper's goal is to keep reverse generation's low cost while recovering the distribution of verified forward responses, by learning a translator from one to the other on a small paired dataset.

Key Contributions

1. A translation framing of the reverse/forward CoT gap. The authors formalize the mismatch between answer-conditioned generation, question-only generation, and rejection-sampled generation as three distributions, and train a translator to map reverse-CoT samples onto the rejection-sampled forward distribution. The framing is explicitly borrowed from consistency and recovery guarantees in machine translation.

2. A practical paired-data recipe and training objective. For each question, one reverse response (source) and the first accepted forward response (target) are collected; the translator is fine-tuned with a teacher-forcing objective on the source only, deliberately excluding the question and answer as inputs so the model cannot bypass translation via a direct question-answer shortcut.

3. A distribution-recovery analysis. Appendix C argues that pairs preserve the target law regardless of the rejection-sampling attempt limit, that the teacher-forcing optimum recovers the target conditional, and that the resulting total-variation bound propagates to a bound on student risk that accounts for response recovery, decoding differences, and question coverage.

4. Empirical validation across students, sizes, families, datasets, teachers, and domains, including a compute accounting that compares the approach against rejection sampling on the same problems.

Main Findings

• Gains over semantic suppression across all six students. On six math benchmarks, CoTranslator improves average accuracy over semantic suppression (SUP) for every student model tested: Qwen3-4B-Base 21.67 → 27.49, Qwen3-4B-Thinking-2507 36.90 → 39.24, Qwen3-8B-Base 28.65 → 30.52, Qwen3-8B 34.67 → 36.05, Qwen2.5-Math-7B 31.88 → 32.22, and Qwen2.5-Math-7B-Instruct 31.26 → 31.99. The headline figure is up to 5.82 percentage points.

• It can match or slightly beat rejection sampling on average for one student. Qwen3-4B-Base reaches 27.49 average versus 27.21 for rejection sampling, while other students remain below their rejection-sampling counterparts (for example, Qwen3-8B: 36.05 versus 43.28).

• Large compute savings. CoTranslator saves 64.1% of estimated compute in the main experiment and 82.0% in the transfer experiments relative to rejection sampling. With 35.9% of rejection sampling's main-experiment compute, it achieves 95.0% of its mean accuracy across the two 4B students.

• The baselines tested do not close the gap. SSR (structural skeleton guidance) performs poorly on reasoning data and does not outperform SUP (Qwen3-4B-Base average 7.68). SUP-AUG, which adds an explicit threat of punishment for answer leakage, also underperforms plain SUP (19.49 versus 21.67 for Qwen3-4B-Base).

• Transfer without retraining works, unevenly. On ODA-Math-460k with the same teacher, gains over SUP are 4.76 and 1.66 percentage points for the 4B and 8B students. Swapping in Kimi-K3 as teacher yields much larger gains of 18.73 and 19.02 points, attributed to Kimi-K3 leaking answers early and CoTranslator recovering from those cases. On STEM, average accuracy improves by 4.88 (4B) and 4.31 (8B) percentage points.

• Translated responses look more like rejection-sampled ones. Using MAUVE, maximum mean discrepancy, and cosine similarity on semantic (Qwen3-Embedding-4B) and lexical (TF-IDF-weighted) features, CoTranslator scores highest on MAUVE and cosine similarity and lowest on MMD across all difficulty levels, where difficulty is measured by the number of rejection-sampling attempts on 12k sampled problems with up to 16 attempts each.

• Fewer answer-dependent shortcuts. Rubric annotations from GPT-6-Luna plus heuristic rules show reduced expected-answer mentions and answer-dependent shortcuts, and more complete independent derivations, for both teachers. The first mention of the answer also appears later in the response after translation.

• Translator size matters more than paired-data size. A 1.7B translator performs poorly; 4B and 8B translators work well and gain little from going beyond 35.7k paired examples (71.4k and 142.8k were also tested). Removing paired training entirely severely degrades the 4B translator. The 4B thinking-only translator consistently outperforms the 8B hybrid translator at matched data sizes, which the authors attribute to its focus on thinking mode.

Methodology in Plain English

The pipeline has three moving parts.

Collect a small paired set. Take a subset of math questions with known correct answers (a 1/8 slice of the cot subset of nvidia/Nemotron-SFT-Math-v4, about 35.7k examples). For each question, ask the teacher twice: once with the answer included, using a "semantic suppression" prompt that tells it to hide the answer (this is the source), and once from the question alone, sampling repeatedly until a response lands on the correct answer (this is the target). The targets are taken from the original dataset, which was generated by DeepSeek-V4-Pro, rejection-sampled, and verified against ground-truth answers.

Train the translator. Fine-tune Qwen3-4B-Thinking-2507 to produce the target given only the source, with standard next-token cross-entropy over the target response. The question and answer are used for pairing and verification but are not fed to the translator, which blocks a direct question-to-answer bypass and forces the model to actually rewrite the reverse reasoning.

Generate at scale, then train the student. For the full cot subset of roughly 285.5k records, run one teacher call with the suppression prompt and one translator call per question, then fine-tune the student on the translated responses. The theory in Appendix C explains why this should approximate the rejection-sampled distribution: the log-sum inequality bounds the marginal divergence between the target and the translated distribution by a fitting-error term plus a term measuring how much target information is missing from the source — and that bound carries through to a bound on student risk under a bounded loss.

Why This Matters

Impact on research. The paper reframes a data-quality problem as a distribution-matching problem, which makes the well-developed toolkit of machine translation — paired data, sequence-level likelihood, recovery guarantees — available to reasoning-distillation pipelines. It also adds to a growing body of work arguing that answer-conditioned rationales can be actively harmful for verifiable-reasoning distillation even when their final answers check out, and it demonstrates that a fix can be learned rather than hand-prompted.

Real-world applications.

• Building math and science tutoring models where labeled question-answer pairs are abundant but high-quality worked solutions are expensive to generate at scale.
• Curating large reasoning datasets under a fixed teacher-inference budget, where rejection sampling would leave hard problems unaccepted.
• Converting existing answer-conditioned explanation corpora into student-trainable material without regenerating them from scratch.
• Cost-constrained distillation for STEM domains (physics, chemistry, biology), where the paper shows transfer from a math-only translator.

Industry relevance. The compute accounting is the practical selling point: 64.1% savings on the main run and 82.0% on transfers, at 95.0% of rejection sampling's average accuracy for 4B students. Because the translator is small and reusable across datasets, teachers, and domains, it fits as a fixed preprocessing stage in an existing distillation pipeline rather than requiring per-dataset re-engineering.

Future Directions

• Larger students and other training regimes. The authors state that experiments cover student models up to 8B parameters and supervised fine-tuning only; applicability to larger models, reinforcement learning, and on-policy distillation is explicitly left unexplored.
• Training a STEM-specific translator. The domain-transfer experiment used a math-trained translator purely to test generalization; the authors note that training on paired STEM data could further improve performance there.
• Understanding the teacher-dependent variance. Gains ranged from 1.66 to 19.02 percentage points depending on the teacher and dataset, with Kimi-K3's early answer leakage driving the largest gains. Whether a translator can be made robust to teachers with very different leakage behaviors is unresolved.
• Whether approximation quality can be tightened. The recovery argument requires that the source contain enough task information to determine the target law and that decoding discrepancy be small; the paper does not report measured values for these bound terms, leaving empirical calibration of the theory open.

Target Audience

Researchers and engineers working on reasoning distillation, synthetic data generation, or chain-of-thought training pipelines who have teacher-inference budgets to manage. It is most useful to readers already familiar with rejection sampling and answer-conditioned generation, and to practitioners deciding whether to pay for verified forward sampling or invest in a reusable translation stage. Readers interested in the theoretical side will find the recovery and student-risk bounds in Appendix C of independent value.

Authors’ abstract

Reasoning distillation requires chain-of-thought (CoT) data with correct answers, but obtaining such data efficiently remains challenging. Forward CoT generation (FCG) followed by rejection sampling (RS) produces verified, question-only responses, yet may require many teacher calls for difficult problems. Reverse CoT generation (RCG) is cheaper because it conditions on the known answer, but access to this answer can lead to post-hoc rationalization and a mismatch with the student's answer-free inference setting. We propose CoTranslator, which learns to translate reverse CoT responses toward the distribution of verified forward responses. Trained on a small set of paired reverse and rejection-sampled forward responses, it supports large-scale generation and verification of student training data. Across six student models evaluated on six mathematics benchmarks, CoTranslator improves average accuracy over semantic suppression by up to $5.82$ percentage points. For two 4B students, it retains $95.0\%$ of rejection sampling's average accuracy while using $35.9\%$ of its estimated compute. Without retraining, the same translator transfers to a new mathematics dataset, a different teacher, and STEM problems, improving average STEM accuracy by $4.88$ and $4.31$ percentage points for 4B and 8B students, respectively. These results show that CoTranslator offers an efficient and transferable approach to reasoning distillation.