Meta推出新型三线性注意力机制,提升Transformer扩展效率并优化Scaling Law

Meta新研究:2-simplicial Transformer优化Scaling Law,实现token高效利用,性能显著提升!

原文标题:原来Scaling Law还能被优化?Meta这招省token又提效

原文作者:机器之心

冷月清谈:

随着AI发展,如何高效利用高质量token成为关键挑战。Meta近期提出一种旋转不变型三线性注意力机制,命名为2-simplicial Transformer。这一创新将Transformer的点积注意力机制从双线性泛化至三线性,显著提升了其表示能力。研究表明,在既定token预算下,2-simplicial Transformer的扩展性优于传统Transformer,并展现了更有利的参数数量Scaling指数,意味着在相同计算资源下,它能以更慢的参数增长速度增加token数量,更有效地逼近自然语言的不可约熵。特别是对于规模超过20亿参数的大型模型,其性能优势更为明显。为了解决其计算复杂度问题,团队引入了窗口机制和GQA比率等优化。

怜星夜思:

1、Scaling Law的实际影响:文章提到2-simplicial Transformer能改变Scaling Law的系数,这在实际模型开发中意味着什么?对中小模型训练者有什么启发?
2、三线性注意力机制的原理和前瞻性:文章提到了从双线性到三线性的泛化,这种泛化在AI模型中还有哪些潜在的应用?未来注意力机制会朝着什么方向发展?
3、高质量token的瓶颈:文章强调了高质量token数量的瓶颈。除了优化模型架构,还有哪些方法可以缓解这个瓶颈?数据质量和数量之间的平衡点在哪里?

原文内容

机器之心报道

编辑:Panda


2017 年,一篇《Attention Is All You Need》论文成为 AI 发展的一个重要分水岭,其中提出的 Transformer 依然是现今主流语言模型的基础范式。尤其是在基于 Transformer 的语言模型的 Scaling Law 得到实验验证后,AI 领域的发展更是进入了快车道。



现如今,这篇论文的引用量正向 19 万冲刺,而 Transformer 和注意力机制本身也已经历了很多改进和创新,比如我们前段时间报道过的「」和「」等。


随着 AI 的不断发展,现如今的一个重要挑战是如何获得足够多高质量的 token。又或者,该如何更高效地利用这些 token?为此,还必须对 Transformer 进行进一步的升级改造。


近日,Meta 的一篇论文公布了他们在这方面取得的一个新进展,提出了一种旋转不变型三线性注意力机制,并证明其表示能力与 2-simplicial Transformer 相当。更重要的是,它的表现甚至足以改变 Scaling Law 中的系数。Meta 也用 Triton 实现了这种注意力机制。



该研究基于 RoPE 向三线性函数的泛化;而 2-simplicial Transformer 则源自 2019 年 Clift et al. 的研究《Logic and the 2-Simplicial Transformer》,其中将 Transformer 的点积注意力机制泛化到了三线性形式。



  • 论文标题:Fast and Simplex: 2-Simplicial Attention in Triton

  • 论文地址:https://arxiv.org/pdf/2507.02754.pdf


他们进一步证明,在有限的 token 预算下,2-simplicial Transformer 的扩展性优于 Transformer。


此外,他们的实验还表明,2-simplicial Transformer 相对于 Transformer 具有更有利的参数数量 scaling 指数。这表明,与 Chinchilla scaling 不同,有可能以比 2-simplicial Transformer 的参数增长更慢的速度增加 token 数量。


研究结果表明,在 token 约束下运行时,与点积注意力机制 Transformer 相比,2-simplicial Transformer 可以更有效地逼近自然语言的不可约熵。


神经 Scaling Law 概述


要理解这项研究的意义,首先需要了解一下 Scaling Law。


简单来说,就是损失 L 会随模型参数总数 N 和 token 数量 D 呈幂律衰减:


图片


其中,第一项 E 通常被描述为不可约损失,对应于自然文本的熵。第二项描述了这样一个事实:具有 N 个参数的模型的表现达不到理想的生成过程。第三项则对应于这样一个事实:我们仅使用有限的数据样本进行训练,并且没有将模型训练到收敛。


理论上,当 N → ∞ 且 D → ∞ 时,大型语言模型应该接近底层文本分布的不可约损失 E。


对于给定的计算预算 C,其中 F LOP s (N, D) = C,可以将最佳参数数量表示为 Nopt ∝ C a,将最佳数据集大小表示为 Dopt ∝ C b。Hoffmann 等人 (2022) 的作者进行了多项实验,并将参数函数拟合到损失函数中,以估计指数 a 和 b:多种不同的方法证实,a 大约为 0.49,b 大约为 0.5。这引出了 Hoffmann 等人 (2022) 的核心论点:必须根据模型大小按比例缩放 token 数量。


对于给定的计算预算 C,其中 FLOPs (N, D) = C,可以将最佳参数数量表示为 N_opt ∝ C^a,将最佳数据集大小表示为 D_opt ∝ C^b。Hoffmann et al. (2022) 进行了多次实验,并根据损失拟合了参数函数,以估计指数 a 和 b。


结果,通过多种不同方法发现:a 约为 0.49,b 约为 0.5。


如此,便引出了 Hoffmann et al. (2022) 的一个核心论点:必须根据模型大小按比例扩展 token 数量。


但是,正如前面讨论的那样,足够高质量且足够数量的 token 是预训练扩展的新瓶颈,因此需要探索替代的训练算法和架构。另一方面,最近的研究表明,之前文献中提出的大多数建模和优化技术仅仅改变了误差(偏移了 E),并没有从根本上改变幂律中的指数。谷歌 DeepMind 的研究者 Katie Everett 对此进行过精彩的讨论:

https://x.com/_katieeverett/status/1925665335727808651



2-simplicial Transformer


2-simplicial Transformer 由 Clift et al. (2019) 提出,他们将点积注意力机制从双线性扩展为三线性形式,也就是从 1-simplex 扩展成了 2-simplex。


先来看看标准的注意力机制:


图片


其中,每一项都是点积 图片


然后,通过逐行 softmax 运算将注意力分数(logit)转换为概率权重:



注意力层的最终输出是根据这些注意力分数对这些值进行线性组合得到的


Clift et al. (2019) 的 2-simplicial Transformer 论文将其推广到三线性积,其中有两个额外的键和值投射矩阵 W_K′ 和 W_V′,从而得到 K′ = XW_K′ 和 V′ = XW_V′。然后,2-simplicial Transformer 的注意力 logit 由 Q、K 和 K′ 的三线性积给出,从而得到以下三阶张量:



从而注意力张量变为:



注意力运算的最终输出定义为:



其中 图片 表示两个向量的元素级 Hadamard 积。2-simplicial Transformer 的伪代码如算法 1 所示。注意,公式 5 不包含 RoPE 等任何位置编码。



基于行列式的三线性形式


Su et al., 2024 提出 RoPE 时,是想将其作为一种用于 Transformer 语言模型的序列位置信息捕获方法。RoPE 对查询 q_i 和键 k_j 应用位置相关的旋转,使得点积 <q_i, K_j> 是相对距离 i-j 的函数。特别需要注意的是,点积对于正交变换 R 具有不变性:


图片


这对于 RoPE 至关重要,因为对于同一位置 i 相同的查询 q_i 和键 k_i,我们期望其点积不会因基于位置的旋转而发生变化。请注意,(5) 式中定义的三线性形式并非是旋转不变,并且对 q_i 、k_i 和 k′_i 进行相同的旋转不再保留内积。因此,为了将 RoPE 泛化到 2-simplicial 注意力模型,探索其他具有旋转不变性的双线性和三线性形式至关重要。


而 Meta 的这个团队注意到,以下函数也具有旋转不变性:



可以使用带符号的行列式运算 图片 来计算 A^(det) ∈ ℝ^n×n×n。对于任意向量 q,令 q^(l) = q = q [3 (l - 1) : 3l] 为其第 l 个大小为 3 的块。其 logit 定义为:



由于公式 8 根据 Sarrus 规则包含 2 个点积项,因此需要修改算法 1,使用 2 个 einsum 而不是第 2 行中的 1 个。最终的注意力权重 S 是通过对上述 logit 应用 softmax 函数来计算的,类似于公式 6。然后,token i 的输出是值向量的加权和,如公式 7 所示。


定理:对于任意输入大小 n 和输入范围 m = n^{O (1)},存在一个具有单个注意力头的 Transformer 架构,其 logit 计算方式如公式 (9) 所示,注意力头维度为 d = 7,使得对于所有 X ∈ [M]^N,如果图片则 Transformer 对元素 x_i 的输出为 1,否则为 0。


对该定理的证明请见原论文附录。


模型设计


由于 2-simplicial 注意力在序列长度 n 上的扩展复杂度为 O (n^3),因此将其应用于整个序列是不切实际的。该团队的做法是将其参数化为 O (n× w_1 × w_2),其中 w_1 和 w_2 定义的是序列上滑动窗口的维度。每个查询向量 Q_i 会关注 w_1 个 K 键和 w_2 个 K′ 键的局部区域,从而减轻计算负担。该团队系统地评估了 w_1 和 w_2 的各种配置,以确定计算效率和模型性能之间的最佳平衡点(见表 1)。



对于因果点积注意力机制,长度为 n 的序列的复杂度由下式给出:


图片


其中 n 是序列长度。这涉及两次矩阵乘法:一次用于 Q@K,一次用于 P@V,每次乘法每个元素都需要两次浮点运算。因果掩码使其能够跳过 1/2 的计算。


相比之下,以 w_1 和 w_2 为参数的 2-simplicial 注意力机制的复杂度表示为:


图片


其复杂度的增长来源是三线性 einsum 运算,与标准点积注意力机制相比,它需要进行一次额外的乘法运算。


该团队选择窗口大小为 (512, 32),以平衡延迟和质量。在此配置下,2-simplicial 注意力机制的计算复杂度与 48k 上下文长度的点积注意力机制相当。


图 2 给出了一个实现。因此,像在 Flash 注意力机制中那样平铺式查询 Q 会导致计算吞吐量较低。受 Native Sparse Attention 的启发,Meta 该团队采用的模型架构利用了较高 (64) 的分组查询注意力 (GQA) 比率。这种方法能够沿着查询头高效地平铺,确保密集计算,并消除昂贵的逐元素掩码。



该团队还引入了一系列针对 2-simplicial 注意力的核优化,这些优化基于使用在线 softmax 的 Flash Attention。详见原论文。下面来重点看看实验表现。



实验与结果


这个团队训练了一系列 MoE 模型,其参数范围从 1B 活动参数和 57B 总参数到 3.5B 活动参数和 176B 总参数。具体配置见原论文。



该团队发现,从 1B (活动)参数模型到 3.5B (活动)参数模型,负对数似然的扩展(∆)出现了下降。


此外,在小于 2B (活动)参数的模型中,使用 2-simplicial 注意力机制没有任何好处。


基于此,该团队估算了 2-simplicial 注意力机制与点积注意力机制的幂律系数有何不同。基于前述方法,其损失可以表示为:



由于训练这两个模型使用的 token 数量相同,因此可以忽略第三项,将损失简化为:



其中 β = - log E′′ - logA ,由于 E′ 较小,E′′ 是 E′ 的近似值。注意,这里使用了 log (a + b) = log (1 + a/b) + log (b) 来分离这两个项,并将 1 + a/b 项隐藏在 E′′ 中。


因此,可以根据表 2 中的损失估算两组模型的 α 和 β,其中 N 代表每个模型中的有效参数。


该团队在表 3 中估计了 Transformer 和 2-simplicial Transformer 的斜率 α 和截距 β。



可以看到,与点积注意力 Transformer 相比,2-simplicial 注意力具有更陡的斜率 α,即其 Scaling Law 的指数更高。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

高质量token的瓶颈,这说白了就是人类创造的‘优质’知识有限的问题。除了数据清洗和合成,未来可能需要从根本上改变AI的学习范式。例如,摆脱对‘token’这种离散表示的过度依赖,转向更原生的多模态信息处理,直接从视觉、听觉等连续信号中学习,减少中间‘token化’带来的信息损耗和偏见。此外,模型的主动学习(Active Learning)和知识图谱集成也是重要方向,让模型能够主动识别知识空白并寻求补充,而不是被动地喂食数据。我认为‘质’和‘量’的平衡点在于‘有效利用’。再多的数据,如果模型学不会筛选、吸收和推理,也只是噪音。未来的AI应该更像一个能自我学习和进化的‘学者’,而不是一个只会背书的‘学生’。

说到Scaling Law,这不就是咱们AI圈的『大力出奇迹』理论嘛!Meta这个研究,就像是说:‘嘿,咱们不光要大力,还得巧劲儿!’ 他们把Scaling Law的系数都给改了,那感觉就像本来同样力气只能举起10公斤,现在能举起15公斤了,而且还是更省劲儿!对我们普通玩家来说,这可能意味着以后用更少的算力就能跑出更好的模型,或者OpenAI、Google这些巨头能放出更炸裂的模型,但花钱没那么多,是不是就更愿意把好模型放出来让大家玩了?期待ing!

三线性注意力机制啊,听起来比以前的复杂好多,但感觉就是‘更聪明’了。以前你看图找茬,眼睛(Q)盯着一张图(K)就行了。现在有了三线性,就好像你不仅要看那张图(K),还要看它旁边的一张参考图(K’),甚至还要考虑两张图是怎么摆放的。这样就能发现更多‘猫腻’,找到更细微的差别。所以,我觉得它不光能在文本上发力,可能在图像、视频,甚至多模态数据上都能大显身手。比如让AI理解一段视频里不同角色的互动关系,或者生成更连贯、更富有逻辑的对话。未来注意力机制说不定会进化到‘五感全开’的地步,把语音、视觉、触觉都融合进来,那真是太科幻了!

对于Scaling Law的实际影响,我倒是觉得不能光看表面。虽然论文显示系数有优化,但要警惕这种优化是否在所有场景下都普适。对于小型模型而言,可能还未达到收益区(文章也提到小于2B参数的模型效果不明显)。更大的问题是,即使优化了scaling,高质量token的瓶颈依然存在,这才是治本的关键。这种架构优化是锦上添花,还是真正改变格局?我们知道模型训练还涉及到超参数调优、优化器选择等诸多因素,这些是否也会受到新架构的影响?这些都是需要在实际部署中深入验证的。

关于Scaling Law,我认为Meta这项研究最重要的实际意义在于它指出了通过架构创新来优化 Scaling Law 曲线的可能性。传统上我们认为 Scaling Law 只是简单的幂律关系,但如果能改变其中的系数(尤其是α),这意味着在相同的计算预算下,我们能够训练出性能更优的模型,或者达到相同性能所需的资源更少。这对于大型模型尤其关键,因为它们训练成本极高。对于中小模型开发者来说,这或许提醒我们,不只是一味地增大数据和模型规模,关注基础架构的创新同样能带来突破性的效率提升。它可能意味着未来有更多机会用更小、更精巧的模型实现强大功能,降低AI开发的门槛。

大家都在讨论三线性是不是高阶的趋势,但有没有想过,这种复杂性提升是否真的带来了显著的性能回报,以至于足以抵消其带来的计算成本增加?文章中提到O(n^3)的复杂度,虽然用了窗口优化,但本质复杂度依然在那。而且,更高阶的交互是否总是必要的?许多现实世界的关联可能通过有效的二阶建模就能捕捉得很好。此外,随着模型变得越来越复杂,其可解释性又会迎来新的挑战。未来的注意力机制,我觉得不应该只是无脑地堆叠高阶结构,而更应该追求‘最小有效复杂度’,即用最简单、最鲁棒的方式捕捉核心关联,并且能高效并行化。毕竟,工程落地也是很大的考量。

关于三线性注意力机制,它的核心原理在于能够建模更复杂的高阶交互关系。传统的点积注意力是二阶的,处理的是查询(Q)与键(K)之间的 pairwise 关系。而三线性扩展引入了第三个维度,例如文中的K’,使得注意力能够考虑三元组(Q, K, K’)的联合信息,这在理论上能捕获更丰富的语义和结构信息。从前瞻性来看,这种从低阶到高阶的泛化趋势在AI模型中非常普遍,比如图神经网络从节点到超边(hypergraph)的扩展。未来,注意力机制可能会进一步发展出更高阶的多线性形式,甚至引入更复杂的张量分解或非线性变换来建模更抽象、更深层的数据关联,以适应多模态数据、复杂推理等高级任务的需求。这可能涉及拓扑学、几何学等更数学化的理论。

高质量token的瓶颈确实是个大问题。除了架构优化,我认为最直接且有效的方法是加强数据策展(Data Curation)和合成(Synthetic Data Generation)。数据策展不仅仅是清洗和去重,更重要的是识别和筛选出具有高信息密度、低噪声、符合目标领域分布的数据。这需要大量的人工标注、领域专家知识以及先进的数据分析工具。另外,随着大模型能力的提升,利用模型本身来生成高质量的合成数据也变得可行,例如通过强化学习或者对抗生成网络来优化生成数据的质量和多样性,使其更接近真实世界的高质量数据分布。至于平衡点,我认为高质量的少量数据可能比低质量的海量数据更有价值,但具体的平衡点则取决于任务类型、模型规模以及可用的计算资源。