ToDi: Token-wise Distillation for LLMs

• ToDi is token-wise distillation that adaptively combines Forward KL and Reverse KL based on the teacher–student token probability ratio.
• It employs a sigmoid-based weight to dynamically adjust loss emphasis, boosting underestimated tokens and suppressing overestimated ones.
• Empirical evaluations demonstrate that ToDi outperforms uniform divergence methods, achieving higher ROUGE-L scores and more faithful probability alignment.

Searching arXiv for the primary ToDi paper and closely related usages of the acronym. ToDi most commonly denotes Token-wise Distillation, a knowledge distillation objective for autoregressive LLMs that replaces the conventional uniform divergence choice with fine-grained divergence control at the token level. In the formulation introduced in 2025, ToDi adaptively combines Forward KL and Reverse KL for each vocabulary token using a sigmoid-based weight derived from the teacher–student probability log-ratio, so that underestimated tokens are preferentially boosted and overestimated tokens are preferentially suppressed. The method was proposed to address a central limitation of KL-based distillation for LLMs: standard objectives apply a single divergence, or a single global mixture, across the whole vocabulary at a timestep, despite token-wise mismatches having distinct optimization requirements (Jung et al., 22 May 2025).

1. Problem formulation and motivation

In autoregressive LLM distillation, a teacher model defines a token distribution $p(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$ and a student model defines $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$ over a shared vocabulary $\mathcal{V}$. Writing

$$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$

the distillation objective seeks to make $q_\theta$ approximate $p$.

The motivating claim behind ToDi is that the dominant KL-based formulations impose a uniform-loss assumption. At each timestep, they compute loss over the whole vocabulary using either Forward KL everywhere, Reverse KL everywhere, a symmetric or fixed mixture such as Jensen–Shannon, or a timestep-level adaptive mixture such as Adaptive KL. In all such cases, the same divergence behavior is applied uniformly across vocabulary items, even though some tokens are severely underestimated by the student, some are overestimated, and some are already well matched.

ToDi treats this mismatch structure as the central design problem. A token satisfying $p_i \gg q_i$ should receive a strong corrective signal that raises the student probability, whereas a token satisfying $q_i \gg p_i$ should receive a strong corrective signal that lowers it. The method’s core contribution is therefore not a new global divergence, but a per-token allocation of divergence type, conditioned on the current teacher–student ratio.

2. Gradient complementarity of Forward KL and Reverse KL

The method is derived from a gradient analysis of the token-level contributions of Forward KL and Reverse KL. For a token $v_i$ at timestep $t$,

$q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$0

Their partial derivatives with respect to the student probability $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$1 are

$q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$2

Letting $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$3, the Forward KL gradient is $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$4, while the Reverse KL gradient is $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$5. This yields the complementary behavior that motivates ToDi. When the student underestimates a token, so that $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$6, Forward KL gives a strong negative gradient and gradient descent pushes $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$7 upward strongly; Reverse KL is weaker. When the student overestimates a token, so that $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$8, Reverse KL gives a strong positive gradient and gradient descent pushes $q_\theta(\cdot \mid \mathbf{y}_{<t}, \mathbf{x})$9 downward strongly; Forward KL is weak.

The paper summarizes this complementarity as follows: when $\mathcal{V}$0, Forward KL provides a strong push up and Reverse KL a weak push up; when $\mathcal{V}$1, Forward KL provides a weak pull down and Reverse KL a strong pull down. The significance of this analysis is methodological: ToDi is not motivated by an abstract preference for symmetry, but by the observation that the two KL directions are useful in different token-level regimes.

3. Token-wise objective and fine-grained divergence control

ToDi implements this insight through a token-wise interpolation between Forward KL and Reverse KL. For each timestep $\mathcal{V}$2 and token $\mathcal{V}$3, it defines a weight

$\mathcal{V}$4

where $\mathcal{V}$5 is the logistic sigmoid and $\mathcal{V}$6 is a stop-gradient operator. The token-level ToDi loss is then

$\mathcal{V}$7

and the full sequence loss sums this quantity over timesteps and vocabulary items.

This weighting rule has an immediate interpretation. If $\mathcal{V}$8, then $\mathcal{V}$9, so $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$0, emphasizing Forward KL. If $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$1, then $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$2, so $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$3, emphasizing Reverse KL. If $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$4, the weight remains near $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$5, yielding approximately symmetric behavior. The control is therefore both continuous and ratio-dependent, rather than a hard switch.

Because $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$6 is detached, differentiation treats it as a constant. The per-token ToDi gradient is

$$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$7

In the regime $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$8, the weight approaches $$p_i := p(v_i \mid \mathbf{y}_{<t}, \mathbf{x}), \qquad q_i := q_\theta(v_i \mid \mathbf{y}_{<t}, \mathbf{x}),$$9, so ToDi behaves approximately like Forward KL; in the regime $q_\theta$0, the weight approaches $q_\theta$1, so it behaves approximately like Reverse KL.

The paper also relates ToDi to Jeffreys divergence. An algebraic rewriting shows that Jeffreys can be expressed using the same sigmoid-of-log-ratio functional form, but without detaching the weight. ToDi’s stop-gradient therefore changes the optimization dynamics. This distinction is important because one common misinterpretation is that ToDi is merely a re-expression of a symmetric KL mixture. The reported empirical behavior argues against that interpretation: the detached, ratio-conditioned mixture performs differently from plain Jeffreys (Jung et al., 22 May 2025).

4. Training configuration, model setups, and benchmarks

The reported experiments use two main teacher–student settings. In the GPT-2 family, the teacher is GPT2-1.5B and the student is GPT2-120M, with full fine-tuning of the student. In the LLaMA family, the teacher is LLaMA2-7B and the student is TinyLLaMA-1.1B, with LoRA fine-tuning of the student. Appendix experiments additionally include OLMo2-7B→1B, Qwen2.5-1.5B→0.5B, and Gemma3-4B→1B (Jung et al., 22 May 2025).

For each training example, teacher and student are run with the same input and teacher-forced target sequence. At each timestep and vocabulary token, the method computes $q_\theta$2, $q_\theta$3, the log-ratio, the sigmoid weight, and the token-level ToDi loss. The total distillation loss is combined with standard cross-entropy on the ground-truth token using

$q_\theta$4

The default weighting function is the sigmoid of the log-ratio with no scaling. A generalized form introduces a parameter $q_\theta$5: $q_\theta$6 The paper identifies $q_\theta$7 as the default and best-performing setting, $q_\theta$8 as the case where $q_\theta$9 everywhere and the objective becomes Jeffreys divergence, and $p$0 as a hard switch between Forward KL and Reverse KL. No temperature scaling is explicitly used for logits.

Training uses the instruction-following dataset databricks/dolly-15k, with approximately 11K train, 1K validation, and 500 test examples. Evaluation is conducted on DollyEval, S-NI, UnNI, SelfInst, and VicunaEval, with instruction following assessed by ROUGE-L and each evaluation repeated with five random seeds and averaged. The student training recipes reported in the appendix are 20 epochs, learning rate $p$1, batch size 32, and full fine-tuning for GPT2-120M; and 10 epochs, learning rate $p$2, batch size 32, and LoRA fine-tuning with rank 256, alpha 8, dropout 0.1 for TinyLLaMA-1.1B. All experiments are run on single RTX A6000 GPUs.

5. Empirical results, ablations, and efficiency

On the main ROUGE-L evaluation, ToDi improves over supervised fine-tuning and over KD baselines using Forward KL, Reverse KL, Jensen–Shannon, total variation distance, skewed KL variants, and Adaptive KL. For GPT2-1.5B → GPT2-120M, the average ROUGE-L over five datasets is 16.61 for SFT, 18.12 for FKL, 18.38 for RKL, and 18.66 for ToDi. For LLaMA2-7B → TinyLLaMA-1.1B, the average ROUGE-L is 21.58 for SFT, 24.00 for FKL, 24.09 for RKL, 24.15 for the best non-ToDi baseline AKL, and 24.83 for ToDi. The method is reported as best or second best on most individual datasets and best on the averages (Jung et al., 22 May 2025).

The paper also reports a GPT-4 preference evaluation for TinyLLaMA student models on 5,000 UnNI examples. In pairwise comparison against FKL, RKL, JS, TVD, SKL, SRKL, and AKL using GPT-4 as judge, ToDi has higher win rates versus every baseline, with most differences statistically significant at $p$3. This suggests that the measured gains are not confined to ROUGE-L.

Ablations emphasize that the gains are attributable to token-wise adaptation, not merely to blending divergences. In the generalized ToDi study, $p$4 performs best, $p$5 is worse, and $p$6, which reverses the weighting logic, causes clear degradation. The appendix further reports that too small a $p$7 leads to almost uniform weights and weaker responsiveness, while very large $p$8 approximates a step function and yields training instability and worse performance. A direct comparison against fixed global mixtures $p$9 shows that token-wise weighting consistently outperforms all fixed uniform ratios. Pearson similarity analyses further show that ToDi models have higher similarity to teacher output distributions than AKL models, suggesting more faithful distribution alignment.

The computational analysis is equally central to the method’s positioning. ToDi has per-step complexity $p_i \gg q_i$0, like standard FKL and RKL, because it requires only per-token log-ratios and sigmoid evaluations and no sorting or ranking. AKL, by contrast, requires sorting tokens into head and tail regions and therefore incurs $p_i \gg q_i$1 overhead. The reported learning curves for TinyLLaMA show that ToDi dominates from the first epoch, remains smooth and stable, and converges to the best final validation ROUGE-L among the compared methods.

6. Scope, limitations, and disambiguation

Within LLM distillation, ToDi occupies a specific position among KL-based objectives. Earlier approaches use Forward KL, Reverse KL, symmetric $p_i \gg q_i$2-divergences such as Jensen–Shannon and total variation distance, skewed variants such as SKL and SRKL, or dynamic mixtures such as AKL that still apply a uniform ratio across the vocabulary at a timestep. ToDi can be understood as a token-wise, ratio-based generalization of these ideas: it retains the same basic building blocks, Forward KL and Reverse KL, but localizes the divergence decision to individual vocabulary items.

The paper also states several practical constraints. ToDi assumes that teacher and student share the same vocabulary, since its per-token comparison depends on aligned token probabilities. It requires access to full teacher logits or probabilities, making it a white-box method rather than a black-box API method. The experiments do not cover the very largest LLMs, and although the default $p_i \gg q_i$3 is reported as robust, other regimes may require tuning. These limitations delimit the method’s immediate applicability, even though the implementation recipe is otherwise straightforward.

The acronym ToDi is overloaded elsewhere in the arXiv literature. It is used for Task-Oriented Diffusion Inversion (“TODInv, ‘ToDi’”), a framework for high-fidelity text-based editing of real images in diffusion models (Xu et al., 2024). It also appears in discussion of task-oriented dialogue benchmarks, where ToD is sometimes abbreviated “ToDi” in connection with the Indonesian benchmark IndoToD (Kautsar et al., 2023). A further 2025 usage denotes TextOmics-guided Diffusion for hit-like molecular generation conditioned on omics expressions and molecular textual descriptions (Yuan et al., 14 Jul 2025). In current LLM distillation research, however, ToDi refers specifically to Token-wise Distillation via Fine-Grained Divergence Control (Jung et al., 22 May 2025).