Variable-Width Transformers

AK's Picks LLM PTSAD ML TAVR
扩大模型规模——特别是增加深度和宽度——推动了基于Transformer的语言模型取得显著进展。然而,目前大多数架构均在所有层中保持恒定的宽度,即以固定且均匀的方式分配参数量与计算资源,尽管不同层级可能承担着差异化的计算功能。本文通过提出一种呈“×”形(即“><”形)的Transformer架构,从实证角度探究了在网络深度方向上进行非均匀容量分配的效果。该设计使模型的早期层与末期层更宽,而中间层则相对变窄,并借助一种无需额外参数的残差重缩放机制实现层间宽度的动态调整。在参数量从2亿到20亿(稠密模型)以及30亿(混合专家模型,MoE)的各类纯解码器语言模型上,我们的“><”形Transformer在语言建模损失方面始终优于参数量相匹配的均匀宽度基线模型。此外,由于降低了各层平均宽度,该架构整体所需浮点运算量(FLOPs)更少(在拟合的损失匹配缩放曲线下可减少22%),同时对键值(KV)缓存的内存占用与I/O开销也更小(降低15%)。进一步分析表明,这种“瓶颈式”结构会导致残差流(residual streams)中的表征呈现出质的差异。总体而言,我们的结果表明:采用非均匀的层宽度分配策略,有助于实现语言模型更具资源效益的规模化扩展。
Scaling model size, specifically depth and width, has driven significant progress in transformer-based language models. However, most architectures maintain a constant width across all layers, allocating a fixed parameter and computation budget evenly despite different layers potentially playing distinct computational roles. In this work, we empirically investigate nonuniform capacity allocation across network depth by proposing a $\times$-shaped > <former architecture. This design maintains wider early and late layers while narrowing the middle layers, utilizing a parameter-free residual resizing mechanism. Across decoder-only language models ranging from 200M to 2B parameters (dense) and 3B parameters (MoE), our > <former consistently outperforms parameter-matched uniform baselines on language modeling loss. By reducing the average layer width, this architecture also requires fewer overall FLOPs (22% reduction under fitted loss-matched scaling curves) and smaller KV cache memory and I/O cost (15% reduction). In analysis, we show that this bottleneck structure results in qualitatively different representations in residual streams. Overall, our results demonstrate that nonuniform width allocation can result in more resource-optimal scaling of language models.
许愿