字节跳动虚拟宽度网络：不扩张主干也能拥有更强大脑的秘密 原创

当前大语言模型的发展如火如荼,研究者们一直在探索如何让模型变得更强。传统观念认为,要提升模型能力,就得增加隐藏层的宽度(类似于给模型的"大脑"增加更多神经元),但这会带来计算成本的平方级增长——成本翻倍可不是闹着玩的。

就在本周，字节跳动Seed团队发表了一篇题为《Virtual Width Networks》的论文,提出了一种巧妙的解决方案:虚拟宽度网络(VWN)。这项技术能在几乎不增加计算成本的情况下,让模型拥有更宽广的表示能力,相当于给模型装上了"虚拟增容"的外挂。

研究团队在大规模实验中发现了令人振奋的结果。他们在一个33亿参数激活量的混合专家(MoE)模型上测试,将虚拟宽度扩大8倍后,模型达到相同性能水平所需的训练数据量大幅减少:下一个词预测任务减少了2.5倍,下两个词预测任务更是减少了3.5倍。更有趣的是,随着训练的进行,这种优势还在持续扩大,就像滚雪球一样越滚越大。研究团队还发现了一个重要规律:虚拟宽度因子与模型损失之间存在近似对数线性关系,这为探索虚拟宽度作为大模型效率提升的新维度提供了初步的实证依据。

问题的缘起:宽度扩张的两难困境

要理解虚拟宽度网络的价值,需要先搞清楚传统模型面临的难题。根据缩放法则,增加模型参数或训练数据都能让模型更强大。具体来说,增加模型宽度(隐藏维度)能让模型在每个向量中装入更多信息,从而表示更复杂的函数,大幅提升性能。

然而直接增加隐藏维度会导致参数和计算量呈平方级增长。举个例子,如果把隐藏维度从D增加到2D,参数量和计算量会变成原来的4倍,资源消耗成本难以承受。这就好比你想让房间变大一倍,结果发现不仅地板面积要翻倍,连墙壁、天花板的材料都要按平方倍数增加,最后总成本远超预期。

为了突破这个瓶颈,研究者们开发了条件计算策略,其中最著名的就是混合专家(MoE)架构。MoE的思路是为每个输入动态激活专门的子网络,只使用一部分参数进行计算,从而在不成比例增加每个样本计算成本的情况下扩大模型规模。

但MoE也有局限。传统MoE架构可以看作只扩展了前馈网络的内部维度,而主干的隐藏维度依然固定,这就像房子虽然有了更多储物空间,但主要起居空间还是那么大,表示能力仍然受到隐藏维度的瓶颈限制,与真正加宽隐藏层的模型相比还是有性能差距。虽然直接增加隐藏维度能弥补这个差距,但又会导致参数和计算量平方级增长。这引出了核心问题:能否在避免平方级成本增长的同时,获得更宽表示的好处?

虚拟宽度网络正是为解决这个问题而生。研究团队的核心洞察是,更宽的表示可以通过扩展嵌入层实现,而不必拓宽隐藏层——后者才是平方级计算成本的主要来源。嵌入层的计算开销相对较小,扩展它不会显著增加总体计算负担。从这个角度看,之前的超连接(Hyper-Connections)和AltUp等方法可以视为VWN框架的简化实例。通过改进VWN的设计,研究团队进一步提升了其表示能力,并发现了虚拟宽度的良好缩放特性——具体而言,发现了虚拟宽度因子与损失之间的缩放关系,这为社区探索虚拟宽度缩放作为大模型效率提升的新维度提供了动力。

虚拟宽度网络的核心机制

虚拟宽度网络的输入是一个拓宽的词嵌入,研究团队称之为"超宽嵌入"。在VWN内部,中间表示相应地称为"超宽隐藏状态"。为了处理这些状态,VWN用"广义超连接"(GHC)替代了标准的残差连接。GHC是一个更通用的表述,统一了超连接(HC)和分数连接(FC)的思想,引入了一种灵活的机制,能以轻量级计算将超宽隐藏状态压缩到主干宽度,再输入注意力或前馈模块,然后将模块输出扩展回超宽宽度,更新超宽隐藏状态供下一层使用。在最后一层,一个归约算子(比如线性投影)将最终的超宽隐藏状态映射回原始隐藏宽度,然后通过解嵌入层产生输出logits。

为了更好地利用拓宽的表示,研究团队将VWN与多词预测(MTP)结合,同时优化标准的下一词目标和辅助的n-gram损失。直观地说,更密集的MTP监督信号能充分调动扩展的虚拟空间,而VWN带来的额外表示自由度改善了短程组合建模,产生协同效应。

在标准Transformer中,L层模型的宽度为D,初始词表示通过嵌入查找得到h?∈R^D。这个表示随后通过Transformer层处理,每层由注意力块和前馈网络(FFN)块组成。在第l层,中间隐藏向量h^l∈R^D由h^(l-1)计算得出。最后一层输出词表示h^L∈R^D,再通过线性头投影到|V|维词表空间。Transformer的计算复杂度为O(D²),说明缩放模型宽度D会导致计算成本平方级增长。

嵌入查找操作只占整体计算成本的一小部分。利用这个洞察,研究团队将嵌入维度与隐藏层维度解耦,让嵌入维度可以大幅扩展,同时中间层计算保持原始隐藏维度。这种方法几乎保持了原始计算成本,同时显著增强了词嵌入的表示能力。

为了增加嵌入维度,研究团队提出了超宽嵌入技术。给定固定的隐藏尺寸D,将嵌入维度扩大到更宽的维度D',从而产生更丰富的词嵌入,而计算开销增加不多。形式上,设h^l∈R^D表示第l层的隐藏状态,将这个隐藏向量均匀划分成m个不相交的片段。接下来定义扩展的嵌入向量e∈R^(D'),其中D'=(n/m)D,n和m都是整数且n>m。这里每个片段大小为D'/n。在输入层,设h'?=e,从而使用更宽的词嵌入。

当扩展比率n/m较大时,可以选择使用单个线性投影将原始1倍嵌入映射到更宽维度,类似于对很宽的嵌入表应用低秩分解。另外还可以采用输入增强策略,为每个输入注入比单个孤立词嵌入更多的信息,进一步丰富拓宽的表示。

广义超连接是有效利用更宽词嵌入同时保持中间层计算原始隐藏维度的关键方法。在每层l,GHC引入一个轻量级变换矩阵,编码原始隐藏表示片段与扩展词嵌入之间的加权关系。这个矩阵由多个子矩阵组成,包括B^l和A^l等。

考虑第l个网络层T^l,它整合了Transformer中的自注意力层或前馈网络。GHC的输出表示为H'l(超宽隐藏状态的重塑形式),可以表示为:H'l = B^l转置 · T^l(○A^l转置 · H'(l-1)) + Â^l转置 · H'(l-1)。这里○A^l负责将超宽状态压缩到主干宽度并输入T^l,B^l将层输出写回超宽空间,而Â^l实现超宽状态在层间的直接传递(类似于残差连接的广义形式)。

为了进一步增强前向过程的适应性,研究团队引入了动态GHC(DGHC),其中变换矩阵根据输入表示H'自适应调整。实践中采用静态和动态参数混合的策略,动态参数通过轻量级线性投影网络生成。为确保训练稳定性,首先对输入特征进行归一化,然后应用线性变换和tanh激活函数,输出再由小的可学习矩阵缩放并与相应的静态矩阵相结合。

动态参数W_β和W_α初始化为0,而静态矩阵按特定模式初始化。静态矩阵B采用循环模式初始化,静态矩阵A初始化为块矩阵。静态成分B和A不使用权重衰减,而动态成分使用。

多词预测的协同增效

在输出层,之前的研究表明多词预测(MTP)可以作为k-gram解码的近似。基于这个洞察,研究团队利用MTP在主干模型之上引入额外的VWN层,构建增强的预测头,从而提供细粒度的监督信号。具体而言,遵循DeepSeek的做法,将下一个词的嵌入与前一个词的最后一层嵌入拼接,应用线性投影生成logits。

采用DeepSeek式的单个密集线性来混合隐藏状态和嵌入(即2D→D投影)在VWN下变得过于昂贵,因为宽度扩展了r倍。朴素的密集混合会扩展到2rD→rD;对于r=8,参数和浮点运算量都大幅增长,难以承受。为解决这个问题,研究团队使用块级线性进行混合。将rD维向量划分成n=r×m个大小为D/m的片段,对每个片段应用相同的小型线性,形状为(2D/m)→(D/m)。换句话说,在每个片段内局部融合隐藏状态和嵌入特征,在所有块之间共享线性投影器。这保留了更宽VWN表示的好处,同时将混合成本保持在与r=1情况相当的水平。

从连接性视角理解虚拟宽度网络

从连接性角度重新诠释VWN,可以将其视为沿深度轴的注意力机制。考虑层堆叠为"深度序列",每个层索引类似于位置,隐藏状态充当"垂直KV缓存"。在这种视角下,常见的连接模式映射为前面层上的类注意力窗口:普通的无残差前馈堆叠对应大小为1的滑动窗口(每层仅处理当前输入并遗忘前一个);残差连接实现大小为2的窗口(当前输入加上紧邻的前一个);密集连接将窗口大小扩展到包括所有前面的层,允许每层重用所有先前的表示。带有广义超连接的VWN介于两者之间:它实现了一种学习的、固定成本的、类似线性注意力的机制,在深度上缩放可访问的深度上下文。

形式上,设第l层的拓宽状态为槽矩阵H'l∈R^((D/m)×n),有n个大小为D/m的槽,令r:=n/m为以D单位度量的宽度扩展。明确表示主干映射的GHC递归为:H'l = B^l转置 · T^l(○A^l转置 · H'(l-1)) + Â^l转置 · H'(l-1)。这里Â^l转置传输/衰减存储在槽中的信息(学习的携带/遗忘算子),B^l转置将当前层的主干摘要写入选定的槽。展开这个递归显式地产生:H'l = Σ[从t=0到l-1] (∏[从i=0到t-1] Â^(l-i)转置) · B^(l-t)转置 · T^(l-t)(○A^(l-t)转置, H'(l-t-1)) + (∏[从i=0到l-1] Â^(l-i)转置) · H'_0。这表明H'_l线性聚合来自较早层的主干变换特征,通过"携带算子"Â传播并在每步通过B写入,捕捉了深度上压缩深度缓存的线性注意力精神。

参数m的选择决定了存储深度信息的内存预算(以D单位度量)为r=n/m。GHC在每层保真度和记住的层数之间分配这个预算。当m=1时,模型以完整的D维保真度存储最多r层(层数少,每层带宽高)。当m>1时,模型存储最多n=rm层,每层压缩到D/m维(层数多,每层带宽低)。因此m控制每层压缩,n控制名义深度窗口,r固定总内存预算。学习的、依赖输入的路由通过衰减而非硬截断提供超出名义窗口的软扩展。

直观地说,更大的m以较低的每层保真度为代价扩展了有效记住的层数。对于更宽的模型,增加的表示能力提供了足够的带宽来容纳更大的m。同样,更深的网络受益于更大的m,因为让每层访问长程、浅层信息可以缓解优化困难并改善梯度流。

在硬路由和软路由之间也有区别。如果Â^l和B^l接近置换/二进制门,更新表现得像深度上的固定大小滑动窗口。当m=1时,有r=n个维度D的槽,模型可以保留最后r层的完整保真度。当m>1时,有n=rm个大小为D/m的槽;每层的D维状态被压缩到D/m并写入一个槽,给出压缩形式大小为n的硬窗口。而在软路由下,使用实值的、可能依赖输入的Â^l和B^l(动态GHC),信息被部分保留并在步骤间混合。当Â^l转置的谱半径小于1时,展开的递归式意味着来自前面层的贡献呈指数衰减。有效的深度感受野可以超过名义硬窗口(对m=1大于r或对m>1大于n),尽管信息逐渐衰减和混合。

考虑一个具体配置:(m,n)=(8,64),因此r=8。模型维护n=64个宽度为D/8的槽。在硬路由下,当前层可以利用最近的64层,每层以原始维度的1/8表示。在软路由下,来自早于64层的贡献可能随衰减持续存在,有效地扩大"深度感受野"。

需要注意的是,这里的注意力类比主要借用了沿深度的KV缓存视角,并不意味着层间连接是通过相似性分数或成对相关性构建的,就像标准自注意力那样。GHC使用学习的(静态或输入条件的)路由矩阵以固定成本在层间携带、压缩和写入信息,而不是计算点积分数或对层索引进行softmax。

实验结果:从小规模到大规模的验证

为了检验分数虚拟拓宽下的VWN有效性,研究团队使用1.5倍配置作为代表性案例,在大规模语言模型预训练中联合评估VWN和多词预测(MTP),并在Collection A上测量下游性能,定义为表2所列基准的平均分数。在1.5倍设置中,省略了归约算子之前的组归一化。

主要评估在内部混合专家(MoE)模型的多个规模上进行综合实验,包括MoE 0.4B/4B和MoE 2.5B/30B,均在大规模内部数据集上训练。每个VWN变体采用(m,n)=(2,3)配置,相对于主干隐藏尺寸实现1.5倍虚拟拓宽,从而将扩展的嵌入空间与固定宽度的主干解耦,在近乎恒定的计算下实现受控评估。

在0.4B/4B模型上,训练目标方面,VWN相对于基线持续降低下一词预测损失,而MTP略微增加下一词预测损失。VWN与MTP的组合在增强变体中达到最低损失,但当包含MTP时仍显示出与基线约0.016的差距。在Collection A的下游评估中,单独的MTP与基线相当,而VWN+MTP在整个训练过程中提供最高的平均准确率增益。

在2.5B/25B模型上,训练目标方面,VWN相对于基线降低了下一词损失,在此规模上在VWN之上添加MTP不会降低优化性能,VWN和VWN+MTP都达到类似的低最终损失,在训练结束时各自比基线低约0.015。在下游评估方面,两个变体都优于基线,VWN+MTP在整个训练过程中始终产生最佳平均准确率。

研究团队还在更强的内部基线上研究虚拟宽度缩放。所有模型默认包含多词预测(MTP)头,联合优化标准的下一词和MTP目标。首先在0.8B激活量MoE(MoE-A0.8B)上进行消融,以区分在固定r下增加m(固定r下更细的隐藏分区)和增加r(固定m下更大的虚拟宽度)之间的效果。然后扩展到3.3B激活量MoE(MoE-A3.3B),评估配置(m,n)=(8,64),对应r=8,在保持主干宽度的同时实现嵌入空间的8倍虚拟拓宽。

消融实验表明,在MoE-A0.8B上,在不同虚拟宽度因子r下对分数率m进行消融。每个图显示下一词训练损失与已见词数(十亿)的关系。从左到右:r=2、4和8。在r=2时,将m从2增加到4略微改善收敛,产生明显但适度的差距。在r=4时,m=8和m=16的变体几乎重叠,表明对分数率不敏感。在r=8时,m=4和m=8曲线同样接近,m=8略有优势。总体而言,一旦m>4,m的效果就会减弱,表明在此规模下,分区粒度超过4提供的益处有限。与第4节的讨论一致,研究团队假设在固定r下,更大的模型倾向于需要更高的m来保持足够的虚拟容量,而较小的模型在相对较低的m值下饱和。

关于虚拟宽度因子的缩放规律,研究团队在固定分数率m=8的情况下,通过设置r∈{2,4,8}和n=r·m={16,32,64}来改变虚拟宽度因子,评估VWN在MoE-A0.8B上的表现,分析r的缩放如何影响损失和准确率。在500B词的训练期间,VWN产生一致的、单调的增益,随着r的增大而增大。在500B词时,VWN×2、VWN×4和VWN×8将下一词损失分别减少Δ=0.020、0.028和0.035,将下两词损失分别减少0.030、0.045和0.058,并将下游准确率分别提高+3.2、+3.5和+4.16个百分点。VWN×8>VWN×4>VWN×2>基线的排序在整个训练过程中保持一致,表明在固定m的情况下扩大超宽嵌入系统地增强了模型容量。

观察到的损失减少与虚拟宽度因子r呈对数线性关系。拟合系数为-0.0069,表明虚拟宽度每翻倍对应约0.0069的损失减少。虽然效应大小适中,但它表明虚拟拓宽带来系统性的效率增益。研究团队假设更具表现力的主干和更有效利用虚拟宽度隐藏表示的改进机制可以进一步放大VWN观察到的效率增益。

在大规模模型上,研究团队在3.3B激活量MoE(MoE-A3.3B)上使用(m,n)=(8,64)进一步评估虚拟宽度缩放,其中隐藏维度被划分为m=8个分区,实现8倍虚拟宽度扩展。为灵活控制训练长度,在整个训练过程中保持学习率恒定。VWN显著加速了优化。在MoE-A3.3B上,它用2.5倍更少的词达到基线的下一词损失,用3.5倍更少的词达到下两词损失。同时,相对于基线的下一词损失差距从早期阶段的Δ=0.025增加到3T词时的约Δ=0.032,下两词损失差距从Δ=0.049增长到Δ=0.056。这些趋势表明VWN的优势随着训练的进行而放大——其相对效率不仅早期出现,而且随时间增强。在多词目标上的更大增益进一步突出了虚拟宽度与MTP监督之间的强协同作用:超宽嵌入为短程组合目标提供了更丰富的表示自由度,而广义超连接在虚拟宽度空间和主干之间传输梯度,无需扩展中间层宽度。在Collection B的下游评估中,VWN实现的峰值平均准确率比基线高+2.16个百分点,证实了性能差距持续存在并随着扩展训练继续扩大。

VWN的意义与未来展望

说到底,虚拟宽度网络为我们提供了一种新的思路来提升大模型效率。传统上,研究者们主要关注深度、宽度和数据规模这三个维度的缩放,而VWN开辟了第四个维度——虚拟宽度缩放。这种方法的巧妙之处在于,它将表示能力的提升从主干宽度中解耦出来,让模型能在几乎不增加计算成本的情况下享受更宽表示的好处。

实验结果表明,适度的1.5倍扩展就能带来稳定的改进,而扩展到8倍虚拟宽度时,优化加速更加明显:下一词预测损失的收敛速度提高了2倍以上,多词预测损失的收敛速度提高了3倍以上。更重要的是,研究团队发现了虚拟宽度因子r与损失减少之间的近似对数线性关系,r每翻倍对应约0.0069的平均损失减少。虽然这个增益的幅度适中,但它表明虚拟宽度可以被视为缩放模型效率的一个新的、可预测的维度,补充了现有文献中的深度、宽度和数据缩放法则。

VWN与标准Transformer堆叠和训练方法集成良好,为研究容量/计算权衡以及探索受控宽度扩展如何有效提高质量提供了具体的参考点。与此同时,将这些算法增益转化为生产效率取决于系统现实。尽管质量/计算权衡很有前景,但VWN面临实际约束:随着隐藏宽度的增长,通信和内存访问开销变得不可忽略,当代硬件对非常宽的激活和跨设备路由并不特别友好。目前对极宽配置的工程支持仍然有限,这限制了可部署性。在实践中,1.5倍到4倍范围内的虚拟宽度扩展在当今的技术栈上更可行,而更大的扩展可能需要软件、内存布局和互连策略的协同设计才能充分实现其潜力。

未来的研究可以在几个方向上扩展这项工作。研究更具表现力的主干架构可能会放大虚拟宽度的增益。开发更有效利用虚拟宽度隐藏表示的机制可以进一步提高效率。探索虚拟宽度与其他缩放维度(如模型深度、数据规模)之间的相互作用也很有价值。从系统角度看,针对VWN优化硬件和软件堆栈可以使更大的虚拟宽度扩展在生产环境中变得实用。

对于普通人来说,这项研究的意义在于它为AI模型的发展提供了一条新路径。过去我们总觉得要让AI变得更聪明,就得投入更多的计算资源,导致训练成本水涨船高。VWN告诉我们,通过巧妙的设计,我们可以在不大幅增加成本的情况下提升模型性能。这意味着未来的AI应用可能会变得更加高效和普及,让更多人能够从AI技术中受益。

https://arxiv.org/pdf/2511.11238

Q1:虚拟宽度网络(VWN)的核心创新是什么? 

A:VWN的核心创新在于将表示宽度与主干宽度解耦。传统上增加模型宽度会导致计算成本平方级增长,而VWN通过扩展嵌入空间而非隐藏层宽度,实现了在几乎不增加主干计算成本的情况下获得更宽表示的好处。在大规模实验中,8倍虚拟宽度扩展使下一词预测加速2.5倍,下两词预测加速3.5倍,且优势随训练持续扩大。

Q2:广义超连接(GHC)在VWN中起什么作用? 

A:广义超连接是VWN处理超宽隐藏状态的关键机制。它通过轻量级计算将超宽隐藏状态压缩到主干宽度输入注意力或前馈模块,然后将输出扩展回超宽宽度。GHC统一了超连接和分数连接的思想,提供了灵活的路由和混合机制,既能在层间携带信息,又能让模型访问更长范围的深度信息,改善梯度流和优化效果。

Q3:VWN的虚拟宽度缩放规律是怎样的? 

A:研究发现虚拟宽度因子与损失之间存在近似对数线性关系,拟合系数为-0.0069,意味着虚拟宽度每翻倍对应约0.0069的损失减少。在固定主干的情况下,实验验证了VWN×2、VWN×4和VWN×8相对基线的改进呈单调递增趋势,在下游任务上分别带来+3.2、+3.5和+4.16个百分点的准确率提升,为虚拟宽度作为大模型效率提升的新维度提供了实证依据。