Align-TI: Distilling Token Interactions for MLLMs

• Align-TI is a knowledge distillation framework for multimodal LLMs that transfers both instruction-aware visual grounding and token interaction dynamics.
• The IVA module improves performance by selectively aligning visual tokens based on teacher-derived attention, enhancing instruction relevance.
• The TPA module aligns token transition probabilities via Monte Carlo sampling, effectively mitigating exposure bias during autoregressive decoding.

Searching arXiv for the specified paper and closely related multimodal distillation work to ground the article. arXiv search query: Align-TI multimodal LLMs token interactions distillation Align-TI is a knowledge distillation framework for multimodal LLMs (MLLMs) introduced in “Beyond Next-Token Alignment: Distilling Multimodal LLMs via Token Interactions” (Chen et al., 10 Feb 2026). It is formulated around the claim that standard MLLM distillation is too narrow when it aligns teacher and student mainly through static next-token distributions. The method instead treats distillation through the perspective of token interactions, distinguishing two primary interaction types: vision-instruction token interactions during prefilling and intra-response token interactions during decoding. On that basis, Align-TI adds two modules—Instruction-aware Vision Alignment (IVA) and Transition Probability Alignment (TPA)—on top of ordinary supervised fine-tuning and standard token-level knowledge distillation (Chen et al., 10 Feb 2026).

1. Problem formulation and conceptual basis

Align-TI studies distillation from a large MLLM teacher to a smaller MLLM student. The multimodal input is written as

$\boldsymbol{x} = (X_v, X_q),$

where $X_v$ is an image and $X_q$ is a text instruction or question. A typical MLLM is written as

$$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$

with $\mathrm{Vis}_{\psi}$ as the vision encoder, $\mathrm{Proj}_{\omega}$ as the vision-language projector, and $\mathrm{LLM}_{\phi}$ as the language backbone. The output sequence is

$$\mathcal{O} = \{\boldsymbol{v}, \boldsymbol{q}, \boldsymbol{y}\},$$

where $\boldsymbol{v}$ denotes visual tokens, $\boldsymbol{q}$ instruction tokens, and $X_v$0 response tokens (Chen et al., 10 Feb 2026).

The paper positions standard distillation as insufficient for two reasons. First, it does not explicitly teach the student what visual evidence the teacher grounded on. Second, it aligns only

$X_v$1

that is, next-token distributions under the oracle prefix, while inference depends on the student’s own generated history rather than ground truth. This creates exposure bias and autoregressive error accumulation (Chen et al., 10 Feb 2026).

Vanilla knowledge distillation is written as

$X_v$2

Align-TI retains this term but argues that MLLM knowledge also resides in the interaction structure that precedes and conditions token generation (Chen et al., 10 Feb 2026).

A common misconception is to read Align-TI as merely another next-token KD variant. The method is explicitly presented as moving beyond next-token alignment by transferring both instruction-conditioned visual extraction and token-to-token generative logic (Chen et al., 10 Feb 2026).

2. Token interactions as the organizing principle

The framework is organized around two interaction types that the paper treats as primary in MLLMs.

The first is vision-instruction token interactions. The paper argues that the same image can trigger different attended regions for different instructions, and that only a few visual tokens are highly relevant while many have low utility. A student with limited capacity should therefore not align all visual tokens uniformly; it should imitate the teacher’s instruction-relevant visual focus (Chen et al., 10 Feb 2026).

The second is intra-response token interactions. During decoding, what matters is not only the teacher’s output distribution at one step, but how one generated token changes the distribution over the next. This motivates alignment of

$X_v$3

rather than only one-step distributions under oracle prefixes (Chen et al., 10 Feb 2026).

This token-interaction view leads to the full Align-TI objective: $X_v$4 The formulation indicates that Align-TI is not a replacement for supervised fine-tuning or vanilla KD, but an augmentation of them with two interaction-aware losses (Chen et al., 10 Feb 2026).

The paper also relates TPA to an ideal sequence-level KD objective,

$X_v$5

and interprets vanilla KD as a limited approximation to that objective. This suggests that TPA is intended as a tractable local expansion of sequence-level alignment rather than a full replacement for sequence-level KL (Chen et al., 10 Feb 2026).

3. Instruction-aware Vision Alignment

IVA is the component responsible for distilling instruction-aware visual grounding. The paper argues that prior visual-token alignment methods often align visual tokens uniformly, despite the fact that instruction relevance is sparse and instruction-dependent (Chen et al., 10 Feb 2026).

To identify which transformer layer is most semantically grounded, the paper introduces the Instruction-Relevance Score (IRS). Let

$X_v$6

denote the vectorized instruction-to-vision attention weights from layer $X_v$7. Then

$X_v$8

The selected layer is

$X_v$9

The interpretation is straightforward: a high IRS indicates that the layer’s instruction-to-vision attention varies with the instruction rather than remaining nearly invariant across inputs (Chen et al., 10 Feb 2026).

At the selected layer $X_q$0, the instruction-to-vision attention submatrix is

$X_q$1

where $X_q$2 is the number of instruction tokens and $X_q$3 is the number of visual tokens. The importance of visual token $X_q$4 is defined by averaging over instruction tokens: $X_q$5 IVA then uses these teacher-derived weights to modulate visual-token alignment: $X_q$6 The paper’s equation is typeset with some corruption, but its intended structure is a weighted per-visual-token KL term (Chen et al., 10 Feb 2026).

Empirically, IVA improves average performance over both plain SFT and uniform visual alignment. In the reported ablation, the baseline average is $X_q$7, adding IVA alone yields $X_q$8, and uniform visual alignment yields $X_q$9. Teacher-derived attention weights also outperform student-derived weights, with $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$0 versus $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$1, which the paper attributes to the student initially lacking robust instruction-aware focus (Chen et al., 10 Feb 2026).

A plausible implication is that IVA is less about generic intermediate-feature matching and more about selective transfer of the teacher’s multimodal evidence-selection policy.

4. Transition Probability Alignment

TPA addresses the second interaction type, intra-response token interactions, and is the more distinctive of the two Align-TI components. The paper argues that vanilla KD only aligns

$$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$2

under the oracle prefix, whereas inference requires robustness under student-generated histories. TPA therefore aligns transition probabilities of the form

$$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$3

This is presented as a way to transfer the teacher’s local generative logic and mitigate exposure bias (Chen et al., 10 Feb 2026).

The paper writes the combined objective as

$$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$4

The second term is TPA proper (Chen et al., 10 Feb 2026).

Because a full vocabulary-sized transition matrix is infeasible to align exhaustively, the expectation over $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$5 is approximated with $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$6 student-sampled candidate tokens: $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$7 The Monte Carlo estimator is

$$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$8

A notable design choice is that the sampled states are student-conditioned rather than teacher-conditioned. The paper argues that this better covers the student’s actual predictive space, so the supervision targets states the student is likely to visit at test time (Chen et al., 10 Feb 2026).

To avoid the prohibitive cost of separate forward passes for every sampled branch, the method constructs an augmented sequence and applies a ribbon attention mask. Each sampled candidate token can attend to the shared prefix but not to other candidates for the same step. This allows teacher and student to evaluate multiple transition branches in parallel within a single masked pass (Chen et al., 10 Feb 2026).

The efficiency analysis is explicit. Vanilla KD requires student $$\mathcal{O} = \mathrm{LLM}_{\phi}\left( \mathrm{Proj}_{\omega} \left( \mathrm{Vis}_{\psi}(X_v) \right), X_q \right),$$9 forward passes and teacher $\mathrm{Vis}_{\psi}$0. TPA with parallelization requires student $\mathrm{Vis}_{\psi}$1 and teacher $\mathrm{Vis}_{\psi}$2, while TPA without parallelization would require student $\mathrm{Vis}_{\psi}$3 and teacher $\mathrm{Vis}_{\psi}$4. The ribbon mask is therefore essential to the practicality of the method (Chen et al., 10 Feb 2026).

In ablations, TPA yields the larger gain of the two new modules: starting from $\mathrm{Vis}_{\psi}$5, TPA alone reaches $\mathrm{Vis}_{\psi}$6, IVA alone reaches $\mathrm{Vis}_{\psi}$7, and the full combination reaches $\mathrm{Vis}_{\psi}$8 (Chen et al., 10 Feb 2026).

5. Training pipeline, model configuration, and workflow

Align-TI is designed for standard LLaVA-style MLLMs with a vision encoder, a vision-language projector, and an LLM decoder. In the main experiments, the teacher LLMs are Qwen2-7B and Qwen3-8B, while the student LLMs are Qwen2-0.5B / 1.5B and Qwen3-0.6B / 1.7B. The appendix specifies SigLIP-B/14 as the vision encoder and a two-layer MLP with GeLU as the projector for teacher construction. The teacher is frozen; the student is fully fine-tuned (Chen et al., 10 Feb 2026).

The data organization follows MobileVLM V2. Pretraining uses 1.2M caption samples, and fine-tuning uses 2.4M mixed captioning and VQA samples, for 3.6M total student-training examples. Knowledge distillation is applied only in the fine-tuning stage rather than pretraining, because the paper argues that learnable capacity is limited and KD is more useful in the fine-tuning regime (Chen et al., 10 Feb 2026).

A single training example proceeds through a specific sequence of operations. The teacher is first run on the teacher-forced multimodal sequence to obtain next-token distributions, visual-token outputs, and instruction-to-vision attention maps. IRS identifies the layer used for IVA. The student is then run on the oracle prefix to compute $\mathrm{Vis}_{\psi}$9 and $\mathrm{Proj}_{\omega}$0. For TPA, candidate tokens are sampled from the student at each timestep, the augmented sequence and ribbon mask are constructed, and teacher and student are run again under that masked layout to compute transition KLs. The final loss is the sum of

$\mathrm{Proj}_{\omega}$1

Algorithm 1 in the appendix gives this logic directly (Chen et al., 10 Feb 2026).

The paper also provides an exposure-bias analysis. Training-time accumulated error is written as

$\mathrm{Proj}_{\omega}$2

while test-time accumulated error is

$\mathrm{Proj}_{\omega}$3

The paper further defines

$\mathrm{Proj}_{\omega}$4

The reported curves show TPA keeping this quantity roughly within $\mathrm{Proj}_{\omega}$5, versus about $\mathrm{Proj}_{\omega}$6 for SFT and vanilla KD, which the paper interprets as evidence that TPA substantially mitigates exposure bias (Chen et al., 10 Feb 2026).

6. Empirical results, efficiency, and limitations

Align-TI is evaluated on six benchmarks: GQA, SQA, TextVQA, POPE, MME, and MMBench. Teacher performance is reported as 75.1 average for Qwen2-7B and 76.0 for Qwen3-8B (Chen et al., 10 Feb 2026).

At roughly 2B scale, the distilled student achieves notably strong results. The Qwen2-1.5B student reaches an average of 72.2, and the Qwen3-1.7B student reaches 73.6. The paper emphasizes that this surpasses LLaVA-1.5-7B, which scores 68.8, yielding a 7.0\% relative improvement under the six-benchmark average used in the paper (Chen et al., 10 Feb 2026).

At smaller scale, the Qwen2-0.5B student reaches 66.7, and the Qwen3-0.6B student reaches 69.7. Against vanilla KD on Qwen2-0.5B, the average improves from 65.0 to 66.7, which corresponds to the reported 2.6\% relative improvement (Chen et al., 10 Feb 2026).

The ablation evidence is unusually clear. Starting from an SFT baseline of 64.3, adding IVA produces 65.1, adding TPA produces 66.4, and adding both gives 66.7. Uniform visual alignment yields 64.5, so teacher-weighted IVA is materially better than unweighted token matching. The best IVA layer in the reported layer-selection ablation is layer 21, and the final TPA setup uses $\mathrm{Proj}_{\omega}$7 sampled tokens, with nucleus sampling outperforming greedy top-$\mathrm{Proj}_{\omega}$8 in low-$\mathrm{Proj}_{\omega}$9 regimes (Chen et al., 10 Feb 2026).

Efficiency is also quantified. Training-time cost is 355 hours for vanilla KD, 509 hours for full Align-TI, and 962 hours for GKD. Memory usage is 70.6 GiB for vanilla KD, 75.6 GiB for full Align-TI, and 76.8 GiB for GKD. This indicates that TPA is the main source of overhead, but the full system remains substantially cheaper than GKD while yielding stronger results (Chen et al., 10 Feb 2026).

Inference efficiency is part of the paper’s practical claim. For example, LLaVA-1.5-7B uses 14.0 GiB peak memory, 90 ms first-token latency, and 33.8 tok/s, whereas Align-TI-2B uses 4.8 GiB, 57 ms, and 64.8 tok/s, while still achieving the stronger six-benchmark average (Chen et al., 10 Feb 2026).

The paper nevertheless leaves several boundaries explicit. Experiments are restricted to image-text benchmarks; extension to video or other modalities is not empirically established. Distillation of the vision-language projector is identified as future work. A plausible implication is that Align-TI’s conceptual framework—token interactions rather than static next-token alignment—could extend beyond image-text MLLMs, but the paper does not claim such generalization as an empirical result (Chen et al., 10 Feb 2026).

Overall, Align-TI is best understood as a distillation framework that transfers not only output distributions but the interaction structure underlying perception and generation. IVA distills the teacher’s instruction-aware visual evidence selection, TPA distills the teacher’s local autoregressive transition logic, and the combined system establishes a stronger distillation baseline for parameter-efficient MLLMs than vanilla next-token KD alone (Chen et al., 10 Feb 2026).