让Transformer的层数有宽有窄,更省资源还更好用。
Variable-Width Transformers

- 早层和晚层变宽,中层变窄,用无参数方法自动调节。
- 在200M到3B参数模型上均优于同等规模均匀结构,少22%计算量。
- 适合追求高效训练与部署的开发者,尤其关注推理成本者。
扩大模型规模(特别是深度和宽度)推动了基于Transformer的语言模型发展。然而,大多数架构在各层保持固定宽度,即使不同层可能承担不同计算角色,也平均分配参数和计算预算。本文通过实证研究提出一种×形变体(> <former),在深层网络中采用非均匀容量分配:早期和晚期层较宽,中间层较窄,并使用无参数残差重缩放机制实现。在从200M到2B参数(密集)及3B参数(MoE)的解码器类语言模型上,该架构在语言建模损失上持续优于参数匹配的均匀基线。通过降低平均层宽,该设计整体减少22%的浮点运算量(在损失匹配缩放曲线下),并降低15%的KV缓存内存与输入输出开销。分析表明,这种瓶颈结构使残差流产生定性不同的表征。结果表明,非均匀宽度分配可实现更优的资源利用率与模型扩展策略。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。