Soft Teacher Forcing

Mohammad Fahes; Cordelia Schmid; Anurag Arnab

cs.CLOct 12, 2026

Soft Teacher Forcing

Authors: Mohammad Fahes, Cordelia Schmid, Anurag Arnab

Organizations: Google Deepmind

Abstract

Soft tokens have so far been used to replace discrete chain-of-thought in latent reasoning methods. These approaches typically require multi-stage training and, in some cases, auxiliary modules, yet they still underperform supervised fine-tuning (SFT) on <input, chain-of-thought, output> triples. In this paper, we ask whether soft tokens can be used to improve SFT itself, instead of as an avenue for latent reasoning. We introduce Soft Teacher Forcing (Soft-TF), a simple training algorithm that combines teacher-forced embeddings with soft tokens through random convex combinations. The soft tokens are computed in a single parallel forward pass over the ground-truth sequence. Soft-TF adds no parameters, no external model and no curriculum, and leaves inference unchanged. Experiments show that Soft-TF improves over SFT on GSM8K, MATH, MathQA and AQuA, by +2.1 to +4.1 points on GSM8K across seven Gemma, Llama and Mistral models.

Figures & tables

Figure 1: ( Left ) Teacher-forced soft sequence computation. Given an initial input sequence, we forward it through the LLM to obtain output probability vectors. These are used as weights for convex combinations of vocabulary tokens, forming soft tokens. ( Right ) Training with Soft-TF. We train the LLM with the next-token prediction objective on convex combinations of the original sequence embeddings and soft sequence embeddings. We call this framework Soft Teacher Forcing (Soft-TF).

Method	GSM8K	MATH	MathQA	AQuA
TF (no CoT)	14.00 $\pm$ 0.28	12.40 $\pm$ 0.32	49.92 $\pm$ 0.49	40.71 $\pm$ 1.94
TF (CoT)	44.70 $\pm$ 0.89	14.44 $\pm$ 0.91	59.86 $\pm$ 0.80	38.19 $\pm$ 4.30
CoCoNut (Stage 1)	26.03 $\pm$ 0.69	9.32 $\pm$ 1.07	46.55 $\pm$ 1.01	32.13 $\pm$ 2.06
CoCoNut (Stage 2)	19.36 $\pm$ 0.78	6.80 $\pm$ 0.47	48.17 $\pm$ 0.88	31.18 $\pm$ 0.85
CoCoNut ( Hao et al., 2025 )	13.59 $\pm$ 0.56	10.08 $\pm$ 1.01	45.16 $\pm$ 1.12	34.96 $\pm$ 1.98
CoCoNut † ( Hao et al., 2025 )	21.00 $\pm$ 1.25	0.92 $\pm$ 0.77	12.14 $\pm$ 2.28	12.05 $\pm$ 2.58

Table 1: Performance (accuracy %) of Soft-TF compared against baselines on GSM8K, MATH, MathQA and AQuA datasets. All methods fine-tune Gemma2-2B pretrained checkpoint, except SoftThinking, which is an inference-only method using the TF (CoT) checkpoint. For CoCoNut † , we report the best decoding strategy per dataset (see Appendix B.2 ).

Method	TF (no CoT)	TF (CoT)	Soft-TF (Ours)
Gemma3-1B	6.82 $\pm$ 0.42	16.77 $\pm$ 0.50	19.38 $\pm$ 0.60 $\uparrow+2.61$
Gemma2-2B	14.00 $\pm$ 0.28	44.70 $\pm$ 0.89	48.52 $\pm$ 1.24 $\uparrow+3.82$
Gemma3-4B	18.54 $\pm$ 0.48	54.86 $\pm$ 0.50	57.95 $\pm$ 1.23 $\uparrow+3.09$
Gemma2-9B	30.37 $\pm$ 0.44	73.89 $\pm$ 0.77	77.07 $\pm$ 1.19 $\uparrow+3.18$
Llama3.2-1B	5.90 $\pm$ 0.26	28.23 $\pm$ 0.88	30.33 $\pm$ 0.59 $\uparrow+2.10$
Llama3.2-3B	13.18 $\pm$ 0.26	48.58 $\pm$ 1.20	52.66 $\pm$ 1.12 $\uparrow+4.08$

Table 2: Accuracy (%) on the GSM8K dataset using different models.

Method	GSM8K	MATH
TF (CoT)	44.70 $\pm$ 0.89	14.44 $\pm$ 0.91
hidden	46.81 $\pm$ 1.11	14.40 $\pm$ 0.94
argmax	48.35 $\pm$ 1.06	13.76 $\pm$ 1.11
convex-full	48.52 $\pm$ 1.24	15.48 $\pm$ 1.50
convex-top-5	49.28 $\pm$ 0.59	16.20 $\pm$ 1.33

Table 3: Ablation study on soft tokens in Soft-TF.

Figure 2: Ablation on the number of forward passes

K

for soft sequence computation. The x-axis is on a log scale.

Panel of Table 4 — Table 4: Choice of $\alpha$ . ( 4(a) ) Constant vs. random $\alpha$ : Soft-TF with constant $\alpha\in]0,1]$ . The dashed horizontal line shows accuracy (%) with $\alpha\sim\text{Beta}(0.1,0.1)$ as a reference. Shaded regions show the standard deviation over 5 runs. ( 4(b) ) Sampling distribution of $\alpha$ : Beta $(0.1,0.1)$ vs. Uniform $[0,1]$ vs. Bernoulli $(0.5)$ with $\alpha\in\{0,1\}$ , with the TF (CoT) baseline as reference.

Method	TF (no CoT)	TF (CoT)	Soft-TF (Ours)
Standard ( $\gamma=0$ )	14.00 $\pm$ 0.28	44.70 $\pm$ 0.89	48.52 $\pm$ 1.24
Soft ( $\gamma=1$ )	13.62 $\pm$ 0.41	41.62 $\pm$ 1.21	43.31 $\pm$ 0.84
Mixed ( $\gamma=0.5$ )	14.10 $\pm$ 0.26	43.99 $\pm$ 1.05	46.50 $\pm$ 1.00
Hybrid	14.06 $\pm$ 0.27	44.49 $\pm$ 1.29	48.72 $\pm$ 1.46

Table 5: Accuracy (%) on the GSM8K dataset for different sampling strategies.

Method		GSM8K	MATH
TF (CoT)		44.70 $\pm$ 0.89	14.44 $\pm$ 0.91
Noisy TF (CoT)	Noise level ( $\sigma$ )
	0.1	44.50 $\pm$ 1.13	14.04 $\pm$ 1.02
	0.3	45.23 $\pm$ 0.30	14.44 $\pm$ 0.41
	0.5	45.44 $\pm$ 0.65	14.12 $\pm$ 0.84
	0.7	45.55 $\pm$ 1.03	13.72 $\pm$ 1.79

Table 6: Comparison to Noisy TF. Accuracy (%) is reported across different noise levels

\sigma

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Iterations	TF (no CoT)	TF (CoT)	Soft-TF
5k	6.82 $\pm$ 0.42	16.77 $\pm$ 0.50	16.24 $\pm$ 0.70
10k	6.07 $\pm$ 0.46	15.63 $\pm$ 0.74	19.38 $\pm$ 0.60
15k	5.70 $\pm$ 0.66	16.41 $\pm$ 0.83	17.30 $\pm$ 1.00

Appendix

Table 7: Gemma3-1B on GSM8K: accuracy (%) at different total iteration budgets (

5

seeds). Each method is trained separately at each budget using its corresponding learning-rate schedule and evaluated at the final iteration. Bold marks the budget selected for each method.

Method	Sec/step	Peak mem/device	Throughput
TF (CoT)	0.882	21.42 GiB	9.07 samp/s
Soft-TF	1.255	24.47 GiB	6.38 samp/s
Overhead	$+42\%$	$+14\%$	$0.70\times$

Appendix

Table 8: Measured training cost of Soft-TF against TF (CoT) on Gemma2-2B, GSM8K.

Dataset	TF (CoT)	Soft-TF	$\Delta$	$p$	95% CI
GSM8K	44.70 $\pm$ 0.89	49.28 $\pm$ 0.59	$+4.58$	0.0002	[ $+3.56$ , $+5.60$ ]
MATH	14.44 $\pm$ 0.91	16.20 $\pm$ 1.33	$+1.76$	0.0053	[ $+0.88$ , $+2.64$ ]
MathQA	59.86 $\pm$ 0.80	61.49 $\pm$ 0.96	$+1.63$	0.058	[ $-0.09$ , $+3.34$ ]
AQuA	38.19 $\pm$ 4.30	42.06 $\pm$ 1.95	$+3.87$	0.0405	[ $+0.27$ , $+7.47$ ]

Appendix

Table 9: Two-sided paired

t

-tests over

5

seeds for Soft-TF against TF (CoT) on Gemma2-2B.

Method	GSM8K	MATH	MathQA	AQuA
TF (no CoT)	14.00 $\pm$ 0.28	12.40 $\pm$ 0.32	49.92 $\pm$ 0.49	40.71 $\pm$ 1.94
TF (CoT)	44.70 $\pm$ 0.89	14.44 $\pm$ 0.91	59.86 $\pm$ 0.80	38.19 $\pm$ 4.30
CoCoNut † (hybrid)	20.68 $\pm$ 1.71	0.52 $\pm$ 0.44	0.96 $\pm$ 0.51	11.89 $\pm$ 4.48
CoCoNut † (standard ( $\gamma=0$ ))	21.00 $\pm$ 1.25	0.92 $\pm$ 0.77	0.94 $\pm$ 0.69	11.18 $\pm$ 3.81
CoCoNut † (Mixed ( $\gamma=0.5$ ))	1.00 $\pm$ 0.38	0.04 $\pm$ 0.09	12.14 $\pm$ 2.28	12.05 $\pm$ 2.58
CoCoNut † (soft ( $\gamma=1$ ))	1.00 $\pm$ 0.38	0.04 $\pm$ 0.09	12.14 $\pm$ 2.28	12.05 $\pm$ 2.58

Appendix

Table 10: Performance (accuracy %) of CoCoNut † with different sampling strategies on GSM8K, MATH, MathQA, and AQuA datasets. Mixed and soft decoding yield identical results on all four datasets.

Figure 3: Sequential soft sequence computation with

K=3

. In contrast to the teacher-forced soft sequence computation (

K=1

, cf. Figure 1 ) adopted in this paper, sequential soft sequence computation requires multiple forward passes through the transformer.

K

reflects the number of soft tokens exhibiting sequential dependency.

Panel of Figure 4 — Figure 4: Inference strategies described in Section 2.4 . (a) Standard discrete inference, (b) Soft inference, (c) Mixed inference, (d) Hybrid inference. Note that, for illustration only, we use $K=2$ for hybrid inference, while the ablation in Table 5 uses $K=1$ .

Figure 5: The CoCoNut † variant, which does not perform compression. Hidden states are used to construct the teacher-forced soft sequence, which is then used for training without resorting to the original sequence.

Figure 6: The SoftThinking † variant extends the original SoftThinking ( Zhang et al., 2025 ) to training. The teacher-forced soft sequence is used for training without resorting to the original sequence.

Explore similar work

Which Tokens Should SFT Actually Learn? A Token-Trimming Perspective on Mathematical Reasoning
Sep 9, 2026Yaning Jia, Chunhui Zhang, Wenxuan Xu +3Numerical Reasoning in Language Models Supervised Fine-Tuning
Dartmouth College · Carnegie Mellon University
A Model with No Head and Many Thoughts
Aug 31, 2026Nikita Koriagin, Yaroslav Aksenov, George Bredis +3Efficient Language Model Reasoning LLM Inference Acceleration
1Yandex · 2T-Tech
Memorize Theorems, Not Instances: Probing SFT Generalization through Mathematical Reasoning
May 10, 2026Ruiying Peng, Mengyu Yang, Jing Lei +3Supervised Fine-Tuning Memorization in Language Models
Tsinghua Shenzhen International Graduate School · Huawei Technologies

Sep 9, 2026J/K move · Enter open · S save

Which Tokens Should SFT Actually Learn? A Token-Trimming Perspective on Mathematical Reasoning

Yaning Jia, Chunhui Zhang, Wenxuan Xu, Xingjian Diao+2

Dartmouth College · Carnegie Mellon University

Numerical Reasoning in Language Models Supervised Fine-Tuning Adaptive Loss Weighting

Open paper AskPDF arXiv

Supervised fine-tuning (SFT) applies a uniform cross-entropy loss to all target tokens, even though different tokens provide unequal learning signals for mathematical reasoning. This uniform treatment can over-sharpen already mastered tokens while amplifying learning pressure on uncertain, low-confidence tokens, leading to suboptimal training dynamics. We propose Trimmed Logit-Gap SFT (TrimSFT), a simple token-level reweighting method that scales the SFT loss according to the logit gap between the gold token and its strongest competitor. TrimSFT trims supervision away from both extremes: tokens already mastered (large logit gap) and tokens weakly supported by the current model (small or negative logit gap), concentrating learning within an intermediate logit-gap region between them. We instantiate this principle with a Gaussian weight centered at margin m with bandwidth τ, requiring no reference model or additional forward pass. We evaluate TrimSFT on six base models from the Llama, Qwen, and DeepMath families across five mathematical reasoning benchmarks. TrimSFT consistently improves over standard SFT, achieving the best average performance on five out of six models, with gains of up to +26.9 points over SFT on MATH500. Further analyses show that the bandwidth τ matters more than the exact margin location, and that half-trim variants that remove supervision pressure from only one side yield inferior trade-offs. A token-level logit-gap distribution analysis suggests that TrimSFT reshapes model confidence in a more balanced way than uniform SFT or monotonic reweighting methods. These results suggest that reasoning SFT can benefit from trimming both extremes rather than treating all tokens uniformly.

Abstract

Explore similar work

Which Tokens Should SFT Actually Learn? A Token-Trimming Perspective on Mathematical Reasoning

A Model with No Head and Many Thoughts

Memorize Theorems, Not Instances: Probing SFT Generalization through Mathematical Reasoning