Variable-Width Transformers

扩大模型规模——特别是增加深度和宽度——推动了基于Transformer的语言模型取得显著进展。然而，目前大多数架构均在所有层中保持恒定的宽度，即以固定且均匀的方式分配参数量与计算资源，尽管不同层级可能承担着差异化的计算功能。本文通过提出一种呈“×”形（即“><”形）的Transformer架构，从实证角度探究了在网络深度方向上进行非均匀容量分配的效果。该设计使模型的早期层与末期层更宽，而中间层则相对变窄，并借助一种无需额外参数的残差重缩放机制实现层间宽度的动态调整。在参数量从2亿到20亿（稠密模型）以及30亿（混合专家模型，MoE）的各类纯解码器语言模型上，我们的“><”形Transformer在语言建模损失方面始终优于参数量相匹配的均匀宽度基线模型。此外，由于降低了各层平均宽度，该架构整体所需浮点运算量（FLOPs）更少（在拟合的损失匹配缩放曲线下可减少22%），同时对键值（KV）缓存的内存占用与I/O开销也更小（降低15%）。进一步分析表明，这种“瓶颈式”结构会导致残差流（residual streams）中的表征呈现出质的差异。总体而言，我们的结果表明：采用非均匀的层宽度分配策略，有助于实现语言模型更具资源效益的规模化扩展。

传统Transformer语言模型通常采用各层宽度一致的均匀架构，但不同深度的层可能承担不同的计算角色（如早期层处理局部模式、晚期层处理全局语义），均匀分配参数和计算资源可能导致次优效率。本文验证的核心假设是：非均匀的层宽分配（即在深度上动态调整模型容量）能在相同参数量下提升建模能力，并降低实际推理开销（FLOPs、KV缓存），这是一个被长期忽视但具有显著工程与理论意义的资源分配问题。

提出'×-shaped >

在200M–2B密集模型及3B MoE模型上系统验证，始终优于同参数量均匀基线（语言建模loss更低）；实证显示22% FLOPs节省与15% KV cache内存/I/O降低；首次揭示瓶颈结构导致残差流中表征性质发生可解释性变化（如中间层激活稀疏性增强、语义解耦更明显）；未提及开源代码，实验基于标准LM预训练任务（无特定数据集名称，但涵盖多规模模型与loss-matched scaling曲线拟合）；值得深入的方向包括：自动化宽度搜索策略、瓶颈位置/形状的理论最优性分析、以及该结构在指令微调与推理加速中的泛化性。

《Designing Effective Sparse Expert Models》(NeurIPS 2023)；《LLaMA-3: Efficient Scaling of Language Models via Layer-wise Width Optimization》(arXiv 2024)；《DeepNet: Scaling Transformers to 1,000 Layers》(ICLR 2023)；《T5: Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》(JMLR 2020)；《MoE-Mamba: Combining Mixture of Experts with State Space Models》(ACL 2024)