让语言模型流动起来：MediaTek Research团队研发的Latent Flow Transformer，通过流匹配压缩Transformer层数

大型语言模型（LLMs）在近年来取得了令人瞩目的进展，但它们的高计算需求和庞大的参数量使得普通用户难以轻松获取和应用这些模型。在这个背景下，来自MediaTek Research的研究团队包括Yen-Chen Wu、Feng-Ting Liao、Meng-Hsi Chen、Pei-Chen Ho、Farhang Nabiei和Da-shan Shiu最近发表了一项突破性研究——Latent Flow Transformer（LFT）。这项研究于2025年5月20日发布在arXiv预印本平台（arXiv:2505.14513v1），旨在解决大型语言模型结构效率问题，为模型压缩提供全新思路。

想象一下，如果你正在搭建一座多层大楼，但发现有些楼层其实并不必要，去掉它们后大楼依然稳固且功能齐全，你会怎么做？这正是LFT尝试解决的问题。当代的大型语言模型就像一座由数十上百层组成的高楼，每一层（transformer层）都会对信息进行处理和传递。研究人员发现，这些层中有许多可能是冗余的，但传统的压缩方法如剪枝和知识蒸馏并未能提供足够规模的减少。

与此同时，在图像生成领域，扩散模型和流匹配（Flow Matching）模型已经证明了其令人印象深刻的效率，最新的方法甚至能够在仅一次采样过程中生成高质量图像。受此启发，研究团队提出了Latent Flow Transformer，这是一种创新的transformer变体，它通过流匹配（Flow Matching）来学习潜在空间的传输映射，实现了将多个transformer层压缩为单一层的目标。

简单来说，LFT就像是在语言模型的多层结构中建立了一条"高速公路"，让信息可以直接从起点跳跃到终点，而不必经过中间的每一个"收费站"（层）。这种方法不仅保持了与原始架构的兼容性，还大大减少了计算需求和参数量。

为了确定哪些层适合被替换，研究团队引入了一个名为"重耦合比率"（Recoupling Ratio）的指标，它能准确预测流匹配的质量。有趣的是，他们的分析发现中间层特别适合压缩，这支持了一个直觉——早期层和后期层在功能上存在质的差异。

在Pythia-410M模型上的实验中，研究团队证明了通过流匹配训练的LFT可以将24层中的6层压缩，并且性能优于直接跳过2层的方法（LM logits的KL散度为0.407 vs. 0.529），这证明了该设计的可行性。更令人兴奋的是，当使用他们新提出的Flow Walking（FW）算法训练时，LFT进一步将12层蒸馏为一层，同时将KL散度降低到0.736，超过了跳过3层的效果（0.932），这显著缩小了自回归和基于流的生成范式之间的差距。

让我们深入了解这项研究的细节，看看MediaTek Research团队是如何实现这一令人印象深刻的突破的。

一、Transformer层的连续时间视角

要理解LFT的工作原理，我们首先需要转变看待Transformer层的方式。想象一下，如果把神经网络的每一层比作一个加工站，那么传统的观点是数据包裹必须在每个站点都停留并接受处理。但研究者提出了一个更流畅的视角：将这些离散的处理站点看作一条连续的流水线。

在这种连续时间的视角下，一个有L个离散层的神经网络，特别是带有残差连接的网络，可以被视为一个将隐藏状态从初始状态逐步改进到最终状态的过程。用数学语言来说，这可以表示为一个常微分方程（ODE）：

``` dht/dt = uθ(ht, t) ```

这里，`θ`是参数化这个方程的参数，当时间`t`与层数`l/L`匹配时，`ht`就对应于第`l`层的隐藏状态。想象成一条河流，水（数据）从起点流向终点，在每一个时刻都有一个微小的变化，这些微小变化的累积形成了整体的转换效果。

然而，学习这种参数θ涉及到昂贵的模拟计算：

``` ht2 = ht1 + ∫(t1到t2) uθ(ht, t) dt ```

这就像是要预测河水从一点流到另一点需要解决复杂的水流方程，计算成本很高。

二、无需模拟的流匹配训练

这就是流匹配（Flow Matching）方法发挥作用的地方。流匹配提供了一种无需模拟的解决方案，让我们可以直接学习从起点到终点的"流动"。

想象你有一个小球从位置x0要移动到位置x1，流匹配就是学习这个球在单位时间内通过预定义路径移动所需的速度场。它通过最小化以下损失函数来学习：

``` LFlowMatching = Et[||uθ(xt, t) - vt||?] ```

其中，xt和vt分别是粒子在时间t的位置和速度。

在实践中，一条直线、匀速轨迹是流匹配的常见选择。在这种情况下，xt和vt可以简单地表示为： ``` xt = (1-t)·x0 + t·x1 vt = x1 - x0 ```

这就像是小球沿着直线匀速移动，从起点到终点。

在离散时间推理过程中，数据点沿着流动轨迹以离散步骤移动。当从时间点t移动到t+d时，一个步骤是：

``` xt+d = xt + d·uθ(xt, t) ```

为了提高流轨迹的稳定性和准确性，可以使用中点估计：

``` xt+d = xt + d·uθ(xt + d/2·uθ(xt, t), t + d/2) ```

这就像是在预测未来位置时，不仅考虑当前位置，还考虑中间过程的状态，从而获得更精确的估计。

三、配对数据流匹配的挑战

在配对样本的流匹配中，至关重要的是保持源分布和目标分布之间的确定性对应关系。这种设置下，当插值轨迹相交时，会对标准流匹配方法带来挑战，因为这些方法在交叉点附近会平均冲突的速度信号，产生有偏估计，无法准确地将源点传输到其配对的目标。

想象两条河流交汇，如果只看交汇点的水流方向，你无法确定哪股水来自哪条上游河道。为了克服这个问题，有研究引入了辅助加速度场来调节每条轨迹上的变化率，从而使对真实配对映射的对齐更加清晰。另一种方法是学习潜在空间投影，在流估计之前解开交叉对，有效地防止轨迹交叉。

四、Latent Flow Transformer架构

有了这些背景知识，我们现在可以深入了解Latent Flow Transformer（LFT）的具体工作原理。LFT是一种新型的、为语言建模设计的Transformer变体，它试图利用流相关概念的优势，正如这些概念在图像生成中所展示的那样。

LFT通过将教师模型中的连续Transformer层块替换为单个学习的传输运算符（称为潜在流层）来减少模型大小。该运算符使用流概念进行训练，以准确地将该块输入的潜在表示映射到其对应的输出。

对于LFT的速度场估计网络，研究团队遵循了前人的方法，通过附加的缩放和移位运算符增强了一个教师Transformer层，以及预测这些因子的MLP网络。他们通过从这个增强网络的输出中减去输入潜在表示来获得速度估计。

五、重耦合比率：预测流匹配质量

选择最佳的层块进行替换对LFT的性能至关重要。通过流匹配学习的可行性受到流路径交叉的限制，而在LFT中这一挑战更为严重，因为保留原始输入-输出配对排除了使用重流方法的可能。

为了指导层选择，研究团队引入了一个称为"重耦合比率"的指标。给定来自层m到n的样本潜在对，这个估计器量化了它们原始配对与由最优传输（OT）指定的配对之间的偏差。最优传输确定了层m和层n的潜在表示之间的最小成本映射。

重耦合比率R被定义为与最优传输矩阵M不一致的配对关系的百分比：

``` R := 1 - E[Tr(M)/OM] ```

这里OM是矩阵M的阶。因为重耦合比率量化了与原始配对的不一致程度，所以较低的R表示更好的一致性，从而预测较少的流交叉问题，学习LFT的可行性更高。

六、学习速度场

要学习速度场，研究团队对一系列x0、x1对应用了标准流匹配算法。对于替换从层m到层n（包括n）的所有层的LFT，他们将给定令牌的层m的输入潜在表示作为x0，将层n的输出潜在表示作为相应的x1。

学习过程如下： 1. 从训练数据集中抽取一个样本d 2. 获取(x0, x1)对，分别是LLM教师模型在层m和层n上处理数据d的潜在表示 3. 随机抽取时间点t 4. 计算插值位置xt = (1-t)x0 + tx1 5. 更新参数θ以最小化uθ(xt, t)和真实速度(x1-x0)之间的平方误差

七、LFT的展开结构与Transformer的相似性

在推理时，将潜在流层展开成一组固定的时间点t0=0 < t1 < t2... < 1，会将潜在演化过程硬化为静态处理图，这有利于可视化数据流和优化硬件实现。对于单步流匹配结合简单重建近似规则的特定情况，潜在流层相当于单个标准Transformer层。如果使用多个步骤，潜在流层相当于具有跨层注意力的Transformer层堆栈。

这种展开的LFT与标准Transformer之间的结构相似性具有重要的实际意义，使研究人员和实践者能够利用为基于Transformer的LLM开发的广泛生态系统和高度优化的基础设施。

八、Flow Walking：提高配对数据流匹配

为了解决交叉轨迹的挑战，同时支持将潜在流层展开回Transformer结构，研究团队引入了Flow Walking (FW)算法来训练和推理LFT的潜在流层。

FW使用数值积分来近似从x0到x1的路径，定义一个步骤为sθ(xt, t, t')=xt+d，其中d=t'-t，xt+d可以由前面介绍的方程确定。关键直觉是通过在交叉点附近轻微分离轨迹来学习非交叉轨迹。学习速度场的目标函数定义为：

``` LFlowWalking(k) = Et1,...,tk-1[||x0 + ∑Δθ,ti - x1||?] ```

其中Δθ,ti = sθ(xti-1, ti-1, ti) - xti-1，t0=0，tk=1，ti从[0,1]中随机采样。

研究团队在实验中发现，选择k=3并随机采样t可以提供一个高效且泛化的训练损失。此外，他们还探索了添加直线性正则化器的可能性，通过标准流匹配来实现：

``` L = LFlowWalking + αLFlowMatching ```

这产生了一个非交叉、直线插值的配对数据，显示出在配对数据应用中的潜力。

九、实验结果与分析

研究团队在Pythia-410M模型上进行了广泛的实验来验证所提出的框架。他们的主要发现包括：

1. **层选择至关重要**：通过重耦合比率选择的层构建的LFT始终优于使用任意层选择的LFT。这验证了该指标在预测流匹配质量方面的有效性。

2. **重耦合比率分析**：分析显示，Transformer模型的中间层特别适合压缩，支持早期层和后期层在质量上有所不同的假设。特别是通过计算层0到层6之间的最优传输矩阵，他们发现重耦合比率为0.53，表明可能会遇到大量与流交叉相关的问题。相比之下，层6到层18之间的最优传输矩阵显示重耦合比率为零，表明对于这批数据，现有的配对已经是最优的。

3. **LFT的蒸馏质量**：在The Pile上训练时，标准流匹配（LFT-SFM）和Flow Walking（LFT-FW）都显示出快速收敛，并且明显优于简单的层跳过方法。当替换Pythia-410m的层6-12（25%的参数）或层6-18（50%的参数）时，LFT-FW在多个离散时间点的推理上始终优于基线。值得注意的是，LFT-FW与k=1达到了与回归模型相当的性能，并且优于具有k=8的LFT-SFM，表明Flow Walking的隐式速度估计更准确地引导模型朝向其目标隐藏状态。

4. **离散时间点对推理性能的影响**：离散时间点的数量k是LFT在推理时的关键超参数。对于LFT-SFM，KL散度和NMSE随着k的减少而减少，这意味着早期速度估计更正确地引导隐藏状态朝向其目标。对于LFT-FW，它在k=3时达到最佳性能，与方程中的三步积分相匹配。当k接近1时，KL散度急剧上升，表明t=0时的隐式速度估计不准确，必须通过多步校正进行改进。有趣的是，尽管从未使用超过三个积分步骤进行训练，LFT-FW对k=8表现出强大的泛化能力，仅显示轻微的性能下降。

5. **性能对比**：在替换层6-18的设置下，LFT-FW（k=3）达到了0.736的KL散度，明显优于跳过三层的基线（0.932），验证了该方法的有效性。这表明Flow Walking成功地克服了主要障碍，即在远距离Transformer层之间对齐潜在传输，有效地桥接了自回归和基于流的建模范式之间的差距。

十、讨论与未来方向

研究团队指出了几个有趣的应用和未来研究方向：

1. **推测解码**：推测解码指的是使用小型草稿模型来加速全尺寸模型的推理。由于直线流Transformer的大小可能比其教师模型小一个数量级，将直线流Transformer用作推测解码中的草稿模型是很自然的。考虑到研究中展示的KL距离结果，这样的配置非常有前景。

2. **动态计算**：传统Transformer具有固定的计算成本和固定的生成质量。相比之下，直线流Transformer允许动态地改变步骤数，即在逐句或甚至逐令牌的基础上动态改变计算投入。这让我们可以考虑如何在使用的步骤数和相应的质量之间达到最佳平衡。

3. **应用于其他模型**：最近，状态空间模型和循环模型因其处理上下文的恒定计算复杂度而成为Transformer的替代候选。研究团队推测，如果应用他们的方法，可以创建直线流RWKV、直线流MAMBA、直线流xLSTM等。

4. **流解缠**：在当前方法中，他们保留了源Transformer的输入和输出层不变。已经确立的是，安排流匹配对以最小化流交叉显著提高了性能。一个提议的优化方向是优化输入和输出层以最小化流交叉，条件是输入嵌入和输出logit。

5. **从头训练流替换的Transformer**：一个开放问题是，一旦流替换并结构简化，是否可以从头训练所得到的浅层Transformer，而无需首先预训练完整深度模型。扩散文献中的先前观察表明，从头训练的一步生成器很少能匹配蒸馏的性能。语言建模中可能也存在类似的模式。

这项研究为提高Transformer模型的效率提供了新的视角和方法，有望在降低计算需求的同时保持性能，使大型语言模型更加普及和易于访问。