🤖 machine learning

Toeplitz MLP Mixers are Low Complexity, Information-Rich Sequence Models

本文介绍了 Toeplitz MLP 混合器（TMM），这是一种类 Transformer 架构，它用三角掩码 Toeplitz 矩阵乘法替代注意力机制，在实现次二次方复杂度的同时，相较于其他次二次方模型展现出更优的训练效率、信息保留能力和上下文学习性能。

原作者： Benjamin L. Badger, Ethan Roland

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Benjamin L. Badger, Ethan Roland

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是论文《Toeplitz MLP 混合器是低复杂度、信息丰富的序列模型》的通俗解释，辅以日常类比。

核心难题：“巴别图书馆”瓶颈

想象一下，你正在写一个故事，但每当你想要添加一个新句子时，都必须从头开始重读整个故事。如果故事很短，这很容易。但如果故事是一整本小说，为每一个新词重读全部内容将耗费无穷无尽的时间。

这正是当前顶尖人工智能模型（称为Transformer）所面临的问题。它们使用一种称为“注意力”的机制，使其能够查看每一个之前的词来决定下一个词。虽然这种机制功能强大，但随着文本变长，它会变得极其缓慢且消耗大量内存。这就像试图在一个图书馆里找一本特定的书，但在拿起书之前，你必须检查每一层架子上的每一本书。

新解决方案：“Toeplitz 混合器”

作者介绍了一种名为**Toeplitz MLP 混合器（TMM）**的新模型。你可以将其视为组织图书馆的一种新方法。

TMM 不是逐个检查每一本书，而是利用一种特殊的、重复的模式（数学上称为Toeplitz 矩阵）来组织信息。

类比：想象工厂里的传送带。在标准模型中，工人必须走完整条生产线才能抓取特定的零件。而在 TMM 中，零件按照重复且可预测的模式排列。工人可以使用捷径（一种称为快速傅里叶变换的数学技巧）瞬间抓取所需之物，无论生产线有多长。
结果：这使得模型速度更快，且占用内存更少，尤其是在首次读取长提示（如预填充文档）时。

意外发现：速度并不意味着“愚蠢”

通常，当你通过简化模型查看数据的方式来加快计算机模型的速度时，它会变得“更笨”。它开始遗忘事物。

“状态空间”模型：其他快速模型（如 Mamba）试图通过保存一个关于它们到目前为止所读内容的单一“摘要”来实现高效。这就像试图只记住最后一句话来回忆一本 500 页的书。这很快，但模型往往会遗忘细节。
TMM 的优势：TMM 不依赖单一摘要。它保持整个文本的“模式”可访问。
- 复制测试：作者通过要求模型复制一长串数字或单词来测试这一点。
- 结果：当“摘要”模型（Mamba）难以记住列表时，TMM 却能几乎完美地复制它，就像那些缓慢且庞大的 Transformer 一样。它比其他快速模型更好地保留了信息。

为什么这行得通？（“无偏差”理论）

论文指出，其他快速模型具有内置的“偏差”或习惯。例如，它们可能被编程为特别关注最近的词，而忽略较旧的词。

类比：想象一个学生只学习教科书的最后一章，因为他认为那最重要。他可能通过最后一章的测验，但如果被问及开头部分就会不及格。
TMM：TMM 没有这种偏差。它均匀地对待文本模式。因为它不强迫自己遗忘旧信息，所以自然地保留了更多细节，从而在从长文档中查找特定事实或遵循复杂指令等任务上表现更好。

“可逆性”的魔力

作者进行了一些深入的数学分析（使用一种称为“算子指数理论”的方法），发现了一个反直觉的事实：

尽管 TMM 被设计为以一种应该丢失部分细节的方式处理信息（因为它是一个只向前看的“因果”模型），但数学表明其内部层实际上非常接近可逆。
类比：想象一台碎纸机。通常，你无法把纸恢复原状。但 TMM 就像一台将纸切成条带的碎纸机，这些条带仍然完美对齐。如果你拥有正确的工具，你可以将它们几乎完美地重新粘贴在一起。这表明模型在处理信息的同时，非常完好地保留了原始信息的“形状”。

局限性（不足之处）

论文诚实地指出了 TMM 目前无法做到的事情：

规模：他们测试的模型相对较小（2 千万到 3 亿参数）。我们尚不知道这是否适用于当今大型科技公司使用的巨型模型。
逐字生成：虽然 TMM 在读取长提示（“预填充”阶段）时速度极快，但一旦开始逐字生成文本，它的速度就会减慢到与旧版 Transformer 相同。它在开始时很快，但不一定在终点线快。

总结

Toeplitz MLP 混合器是一种新型 AI 架构，它利用巧妙的数学捷径快速阅读长文本而不遗忘细节。它证明了你不必为了获得“速度”而牺牲“记忆”。它就像一位使用智能归档系统瞬间找到书籍的图书管理员，而不是一位试图在脑海中 memorize 整个图书馆的图书管理员。

技术摘要：作为低复杂度、高信息丰富度序列模型的托普利茨 MLP 混合器

1. 问题陈述

基于 Transformer 的大语言模型（LLM）从根本上受到自注意力机制二次时间复杂度（ $O(n^2)$ ）和空间复杂度的限制。尽管人们已尝试通过多种方法将这种复杂度降低至次二次或线性水平（例如线性化注意力、状态空间模型如 Mamba），但这些方法往往引入了显著的权衡。具体而言，线性复杂度模型经常遭受训练效率降低、数值不稳定以及明显无法保留输入信息的问题。这种信息保留缺陷表现为在需要精确复制、上下文学习和长上下文检索的任务中表现不佳，而这些正是现代 LLM 能力的标志。核心挑战在于降低计算复杂度，同时不对令牌信息访问施加严重限制，以免损害这些涌现能力。

2. 方法论：托普利茨 MLP 混合器（TMM）

作者引入了托普利茨 MLP 混合器（TMM），这是一种类似 Transformer 的架构，用序列维度上的三角掩码托普利茨矩阵乘法替代了自注意力。

核心架构

托普利茨参数化： 与标准 MLP 混合器中令牌混合权重为任意矩阵不同，TMM 将这些权重参数化为托普利茨矩阵。在托普利茨矩阵中，给定对角线上的所有值都是相同的，这使得整个矩阵可以由单个向量（第一行和第一列）表示。
因果掩码： 为了强制实施下一个令牌预测的因果性，材料化的权重矩阵被施加了三角掩码。在实现中，激活值左乘该掩码后的托普利茨矩阵（ $T'$ ），该矩阵为上三角矩阵。
通过 FFT 计算： 通过将托普利茨矩阵嵌入循环矩阵，作者利用快速傅里叶变换（FFT）来计算矩阵 - 向量乘积。这将训练和推理预填充阶段的计算复杂度从 $O(n^2)$ 降低至 $O(dn \log n)$ ，空间复杂度为 $O(dn)$。
多头和多核变体： 该架构支持多头配置，其中独立的托普利茨矩阵应用于投影的头维度。此外，“核”变体允许在有限数量的隐藏维度元素上进行混合，类似于 1D 卷积。

理论基础

本文通过算子指数理论的视角分析 TMM。具体而言，它考察了训练后的托普利茨层的 Fredholm 指数。作者指出，虽然 TMM 并非对所有可能的输入严格可逆，但训练后的因果 TMM 层通常表现出类似于移位算子的指数（具体为 -1），这表明了一种不同于其他架构偏见的序列信息移位机制。

3. 主要贡献与主张

本文提出了关于 TMM 与其他次二次架构（如 Hyena、Mamba 和窗口化注意力 Transformer）相比有效性的三个主要主张：

训练效率： 在恒定计算条件下，TMM 实现了与其他次二次模型相似或更低的评估损失。与 Hyena 和 Mamba 等线性复杂度基线相比，它们表现出更优越的收敛速率。
信息保留： TMM 保留了比可比的次二次模型更多的输入信息。这体现在：
- 信息提取： 从最终隐藏状态再生输入序列的保真度更高。
- 复制能力： TMM 在复制任务上的训练效率接近或超过了二次复杂度 Transformer，而 Mamba 和 Hyena 模型通常在此任务上表现挣扎。
- 信息容量： 当明确训练以最大化信息保留时，即使需要显著更少的计算资源（1/2 到 1/8），TMM 也表现出比 Hyena 或 Mamba 模型更高的信息容量。
基准测试性能： 由于卓越的信息保留能力，TMM 在需要精确输入表示的基准测试（如信息检索 IR 和上下文学习 ICL）中优于其他低复杂度模型，同时在通用语言理解任务上保持同等水平。

4. 实验结果

作者在 FineWeb-edu 的 100 亿令牌子集上训练了模型，并将其与 Transformer、Mamba 2 和 Hyena 模型进行了评估。

训练效率： 在匹配计算和内存的情况下，TMM 的收敛速度快于 Hyena 和 Mamba 模型。作者指出，线性注意力实现存在数值不稳定和收敛缓慢的问题，因此他们专注于 TMM。
复制任务： 在要求模型将 512 个令牌从输入复制到输出的任务中，TMM 实现了与 Transformer 相当的训练效率。值得注意的是，一个模型维度（ $d_m$ ）仅为 Hyena 模型 1/32 的 TMM，实现了几乎相同的复制效率。研究发现，可训练的托普利茨参数至关重要；冻结这些矩阵会导致完全无法复制，而仅训练可学习的头投影则不足以完成任务。
基准测试：
- IR/ICL： TMM 在 SQuAD、SQuAD v2、LongBench 和 IFEval 上优于 Hyena。
- 通用任务： 在 ARC-Easy、GLUE 和 HellaSwag 等任务上，TMM 的表现与 Hyena 和 Transformer 相当。
消融研究： 作者发现，使用两个由非线性分隔的连续托普利茨层（标准 MLP 混合器风格）比单层效率更低。此外，增加头数改善了渐近收敛，而将核大小增加到超过一定限度则会降低效率。

5. 意义与局限性

意义

本文认为，TMM 提供了一条通往具有竞争力的训练效率和快速解码的“简单、无注意力路径”。其主要意义在于证明低复杂度模型并不 inherently 需要会损害信息保留的严重架构偏见（如近期偏见或状态空间选择性）。通过使用显式的托普利茨参数化，该模型避免了 Mamba 和 Hyena 中出现的“复制差距”，表明结构化访问模式（托普利茨）可以在预填充阶段无需二次成本的情况下，保留 Transformer 的高信息丰富度能力。

作者还强调了一个反直觉的理论发现：训练后的因果 TMM 层在其完整输入空间上不可逆，但其 Fredholm 指数往往比那些显式训练为可逆的模型（如自编码器）更接近零（或具体为 -1，充当移位算子）。这表明因果模型中层的“可逆性”更多是关于通过移位保留信息流，而非数学上的可逆性。

局限性

作者明确承认了几个局限性：

规模： 实验是在 2000 万到 3 亿参数、在约 130 亿至 260 亿令牌上训练的模型上进行的。无法保证这些发现能扩展到在高出几个数量级的计算资源上训练的模型。
自回归推理： 虽然 TMM 在预填充阶段为 $O(dn \log n)$ ，但在计算第一个令牌后，自回归生成会恢复到 $O(n^2)$ 复杂度（因为输入矩阵的行必须被缓存并顺序处理）。因此，TMM 为长提示提供了更快的首令牌时间，但并未加速后续令牌的生成，与 Transformer 相比没有优势。
配置空间： 由于超参数组合数量呈指数级增长，该研究仅测试了有限数量的架构配置。

总之，TMM 代表了序列建模的一个有前景的方向，它在计算效率与复杂推理和检索任务所需的高信息保留之间取得了平衡，前提是能够管理自回归生成的推理瓶颈。