---
title: 'Align-TI: Distilling Token Interactions for MLLMs'
url: https://www.emergentmind.com/topics/align-ti
type: topic
---

# Align-TI: Distilling Token Interactions for MLLMs

Searching arXiv for the specified paper and closely related multimodal distillation work to ground the article.
arXiv search query: Align-TI multimodal large language models token interactions distillation
Align-TI is a knowledge distillation framework for multimodal large language models (MLLMs) introduced in “Beyond Next-Token Alignment: Distilling Multimodal Large Language Models via Token Interactions” [2602.09483]. 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 [2602.09483].

## 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 \(\boldsymbol{y}\) response tokens [2602.09483].

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
\[
P(y_k \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}),
\]
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 [2602.09483].

Vanilla knowledge distillation is written as
\[
\mathcal{L}_{\mathrm{kd}}(\theta)= \mathbb{E}_{(\boldsymbol{x},\boldsymbol{y}_{1:L}^{\mathcal{D}}) \sim \mathcal{D}} \left[ \sum_{k=1}^{L} D_{\mathrm{KL}\!\left( P_T(\cdot \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \,\|\,  P_S^\theta(\cdot \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \right) \right].
\]
Align-TI retains this term but argues that MLLM knowledge also resides in the interaction structure that precedes and conditions token generation [2602.09483].

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 [2602.09483].

## 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 [2602.09483].

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
\[
P(y_{k+1} \mid y_k, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}})
\]
rather than only one-step distributions under oracle prefixes [2602.09483].

This token-interaction view leads to the full Align-TI objective:
\[
\mathcal{L}(\theta)= \mathcal{L}_{\mathrm{sft}}(\theta) +\mathcal{L}_{\mathrm{iva}}(\theta) +\mathcal{L}_{\mathrm{kd}}(\theta) +\mathcal{L}_{\mathrm{tpa}}(\theta).
\]
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 [2602.09483].

The paper also relates TPA to an ideal sequence-level KD objective,
\[
\min_\theta \mathbb{E}_{\boldsymbol{x}\sim\mathcal{D}_x} \left[ D_{\mathrm{KL}\!\left( P_T \| P_S^\theta \right)(\boldsymbol{y}_{1:L}\mid \boldsymbol{x}) \right],
\]
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 [2602.09483].

## 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 [2602.09483].

To identify which transformer layer is most semantically grounded, the paper introduces the **Instruction-Relevance Score (IRS)**. Let
\[
\boldsymbol{\alpha}^{(l)}(\boldsymbol{x})
\]
denote the vectorized instruction-to-vision attention weights from layer \(l\). Then
\[
\mathrm{IRS}(l) = 1 - \mathbb{E}_{\substack{\boldsymbol{x}_1, \boldsymbol{x}_2 \\
\text{i.i.d.} \sim \mathcal{D}_x}} \left[ \cos\!\left( \boldsymbol{\alpha}^{(l)}(\boldsymbol{x}_1), \boldsymbol{\alpha}^{(l)}(\boldsymbol{x}_2) \right) \right].
\]
The selected layer is
\[
l^* = \arg\max_l \mathrm{IRS}(l).
\]
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 [2602.09483].

At the selected layer \(l^*\), the instruction-to-vision attention submatrix is
\[
\boldsymbol{A}_{i \rightarrow v} \in \mathbb{R}^{N_i \times N_v},
\]
where \(N_i\) is the number of instruction tokens and \(N_v\) is the number of visual tokens. The importance of visual token \(k\) is defined by averaging over instruction tokens:
\[
w_k = \frac{1}{N_i}\sum_{u=1}^{N_i} \boldsymbol{A}_{i\to v}(u,k).
\]
IVA then uses these teacher-derived weights to modulate visual-token alignment:
\[
\mathcal{L}_{\mathrm{iva}}(\theta) = \sum_{k=1}^{N_v} \left( \frac{1}{N_i}\sum_{u=1}^{N_i}\boldsymbol{A}_{i\to v}(u,k) \right) D_{\mathrm{KL}}(\text{teacher visual token }k \;\|\; \text{student visual token }k).
\]
The paper’s equation is typeset with some corruption, but its intended structure is a weighted per-visual-token KL term [2602.09483].

Empirically, IVA improves average performance over both plain SFT and uniform visual alignment. In the reported ablation, the baseline average is \(64.3\), adding IVA alone yields \(65.1\), and uniform visual alignment yields \(64.5\). Teacher-derived attention weights also outperform student-derived weights, with \(65.1\) versus \(64.7\), which the paper attributes to the student initially lacking robust instruction-aware focus [2602.09483].

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
\[
P(y_k \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}})
\]
under the oracle prefix, whereas inference requires robustness under student-generated histories. TPA therefore aligns transition probabilities of the form
\[
P(y_{k+1} \mid y_k, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}).
\]
This is presented as a way to transfer the teacher’s local generative logic and mitigate exposure bias [2602.09483].

The paper writes the combined objective as
\[
\begin{aligned}
\arg \min_{\theta}\; \mathbb{E}_{(\boldsymbol{x}, \boldsymbol{y}_{1:L}^{\mathcal{D}}) \sim \mathcal{D}} \Biggl[ \sum_{k=1}^L \biggl( 
&D_{\mathrm{KL}\!\left( P_T(\cdot \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \;\|\; P_S^\theta(\cdot \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \right) \\
&+ \mathbb{E}_{y_k \sim P_S^\theta} D_{\mathrm{KL}\!\left( P_T(\cdot \mid y_k, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \;\|\; P_S^\theta(\cdot \mid y_k, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \right)
\biggr) \Biggr].
\end{aligned}
\]
The second term is TPA proper [2602.09483].

Because a full vocabulary-sized transition matrix is infeasible to align exhaustively, the expectation over \(y_k\) is approximated with \(d\) student-sampled candidate tokens:
\[
\{y_k^{(u)}\}_{u=1}^d \sim P_S^\theta(\cdot \mid \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}).
\]
The Monte Carlo estimator is
\[
\mathcal{L}_{\mathrm{tpa}}(\theta) \simeq \mathbb{E}_{(\boldsymbol{x},\boldsymbol{y}_{1:L}^{\mathcal{D}})\sim\mathcal{D}} \left[ \sum_{k=1}^L\frac{1}{d}\sum_{u=1}^d D_{\mathrm{KL}\!\left( P_T(\cdot \mid y_k^{(u)}, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \| P_S^\theta(\cdot \mid y_k^{(u)}, \boldsymbol{x}, \boldsymbol{y}_{<k}^{\mathcal{D}}) \right) \right].
\]

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 [2602.09483].

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 [2602.09483].

The efficiency analysis is explicit. Vanilla KD requires student \(N\) forward passes and teacher \(N\). TPA with parallelization requires student \(2N\) and teacher \(N\), while TPA without parallelization would require student \(dLN+N\) and teacher \(dLN\). The ribbon mask is therefore essential to the practicality of the method [2602.09483].

In ablations, TPA yields the larger gain of the two new modules: starting from \(64.3\), TPA alone reaches \(66.4\), IVA alone reaches \(65.1\), and the full combination reaches \(66.7\) [2602.09483].

## 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 [2602.09483].

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 [2602.09483].

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 \(\mathcal{L}_{\mathrm{sft}}\) and \(\mathcal{L}_{\mathrm{kd}}\). 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
\[
\mathcal{L}_{\mathrm{sft}} + \mathcal{L}_{\mathrm{iva}} + \mathcal{L}_{\mathrm{kd}} + \mathcal{L}_{\mathrm{tpa}}.
\]
Algorithm 1 in the appendix gives this logic directly [2602.09483].

The paper also provides an exposure-bias analysis. Training-time accumulated error is written as
\[
E_{\text{train}}(l) = \sum_{t=1}^l \mathbb{E}_{(\boldsymbol{x}, \boldsymbol{y}) \sim \mathcal{D}^{e}} \left[ D_{\mathrm{KL}}(P_T \| P_S^\theta)(y_t \mid \boldsymbol{x}, \boldsymbol{y}_{<t}) \right],
\]
while test-time accumulated error is
\[
E_{\text{test}}(l) = \sum_{t=1}^l \mathbb{E}_{\boldsymbol{x} \sim \mathcal{D}_x^{e}, \, \boldsymbol{y}_{<t} \sim P_S^\theta} \left[ D_{\mathrm{KL}}(P_T \| P_S^\theta)(y_t \mid \boldsymbol{x}, \boldsymbol{y}_{<t}) \right].
\]
The paper further defines
\[
\%\mathrm{ExAccErr}_{\le}(l) = \frac{R(l)-E(l)}{E(l)} \times 100\%.
\]
The reported curves show TPA keeping this quantity roughly within \((-10\%,10\%)\), versus about \(30\%\) for SFT and vanilla KD, which the paper interprets as evidence that TPA substantially mitigates exposure bias [2602.09483].

## 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 [2602.09483].

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 [2602.09483].

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** [2602.09483].

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 **\(d=4\)** sampled tokens, with nucleus sampling outperforming greedy top-\(d\) in low-\(d\) regimes [2602.09483].

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 [2602.09483].

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 [2602.09483].

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 [2602.09483].

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 [2602.09483].

Source: https://www.emergentmind.com/topics/align-ti