可变宽度Transformer

TL;DR 提出 × 形 Transformer（> <former），在 decoder-only 模型中采用宽-窄-宽的非均匀宽度分配，结合无参数残差塑形，以更少计算量和 KV 缓存实现超越等参数均匀基线的语言建模性能。

问题

问题背景

当前大语言模型研究中，模型规模扩展（scaling）主要依赖深度和宽度的同步增长。大部分 Transformer 架构（如 GPT 系列）在所有层保持恒定的模型宽度，将参数和计算预算均匀分配。

现有方法的局限

恒定宽度设计假设每层具有相同的表达能力需求，但实际不同层承担的计算角色不同：浅层提取局部特征，深层整合全局语义，而中间层可能处于过渡状态，存在表示冗余。均匀分配会导致：

• 中间层参数效率低下，激活稀疏性高（MLP 激活利用率低）；
• 中间层残差流趋于表示坍缩，不同 token 的表示相似度过早升高；
• 固定宽度带来固定的 KV 缓存与计算成本，难以针对硬件资源动态优化。

为什么这个问题难且重要

• 设计空间复杂：非均匀宽度分配涉及各层宽度的连续决策，搜索空间巨大，且需保证残差连接维度匹配。论文提出 × 形架构（宽-窄-宽）并配合无参数残差缩放机制，简洁地解决了维度变化问题，避免引入额外参数。
• 实用性高：在参数量匹配的前提下，该设计不仅降低验证损失，还能显著减少 FLOPs（最高 22%）与 KV 缓存 I/O（15%），这对推理成本和部署效率至关重要，尤其适合 MoE 模型降低通信开销。
• 认知价值：揭示了固定宽度假设的局限性，推动理解不同深度对容量的需求差异，可能启发更细粒度的 动态宽度分配 或 条件计算。

行业类比

类似于在移动端部署视觉模型时，通过 层间通道剪枝 削减中间层卷积核数量，在不牺牲准确率的前提下大幅降低推理延迟与内存占用。

核心洞察

• **宽度非均匀分配打破了 Transformer 层的同构假设**：绝大多数 Transformer 架构对所有层保持恒定宽度，隐含假设每层计算角色相同。本文通过系统性实验发现“两头宽、中间窄”的 × 形宽度调度（><former）在同等参数量下持续优于均匀基线，表明中间层的表示可被压缩而不损害整体性能。这与单纯增加深度或宽度的传统缩放思路形成对比，揭示了按层功能差异分配容量的有效性，为更细粒度的模型缩放提供了新维度。
• **无参数的残差调整机制使得可变宽度架构可以直接嵌入标准 Transformer**：不同宽度的层之间传递残差流时，作者引入了一种无参数（parameter-free）的残差大小调整方法，无需引入额外的可学习参数或投影矩阵。相比需要学习线性投影的异构架构，该设计避免了额外的参数和计算开销，并且可以直接替换均匀宽度的 Transformer 层，简化了工程实现。这在保持简洁性的同时，实现了性能与效率（FLOPs 减少22%，KV 缓存减少15%）的同步提升。

方法

核心设计

Variable-Width Transformer(变宽 Transformer)以**×形宽度分配**作为网络的基本结构，对应实现为 > <former 架构。整体流程仍遵循标准自回归语言模型：输入 token 嵌入 → 多层 Transformer 块 → 语言建模头输出 logits，但每一层的隐维度不再恒定。

层宽度调度与残差流

宽度按 × 形 schedule 分配：首尾若干层保持较宽的隐维度 d_max，中间层则压缩至较窄的 d_min（瓶颈），形成“两端宽、中间窄”的变体。不同层之间通过无参数残差尺寸调整(parameter-free residual resizing)衔接，使残差流可跨不同维度传递而不引入额外可训参数：

• 下采样(宽→窄)：直接截取前 d_min 维，或通过固定的平均池化，无需学习投影。
• 上采样(窄→宽)：尾部补零或复制前端维度，保持无参数性质。

这保证了与等参数量的恒定宽度基线相比，不因维度变换而增加参数量，实现公平的容量重分配。

训练与推断特性

训练目标为标准下一 token 预测的交叉熵损失，兼容 dense 及 MoE 变体。因平均层宽度降低，该架构天然减少总 FLOPs（实验测得在拟合的等 loss 曲线下减少 22%）及 KV 缓存与 I/O 开销（减少 15%），且始终优于参数匹配的均匀宽度 Transformer。

与同类方法的差异

不同于 DynaBERT 或 LayerDrop 等需动态选择宽度/深度的自适应方法，> <former 采用静态、结构性的瓶颈设计，不需要训练门控网络或额外的搜索阶段，具有低实现复杂度。与通过 LoRA 等适配器改变等效宽度的方案相比，该工作直接在预训练阶段内重分配尺寸，而非微调阶段。

实验

实验设计

• 在解码器语言模型 (decoder-only LM) 上对比不同宽度分配策略, 参数量覆盖 200M 至 2B 稠密模型 及 3B MoE 模型。
• 核心架构为 × 形状 (> <former), 即早期和晚期层保持较宽, 中间层收窄, 通过无参数残差缩放 (parameter-free residual resizing) 连接。
• 均匀宽度基线保持所有层宽度相同, 总参数量与 × 形状模型严格匹配。训练设置 (数据、优化器、训练步数) 保持一致。

关键发现

• × 形状架构在语言建模损失 (LM loss) 上始终优于均匀宽度基线, 且优势随模型规模增大而保持。
• 平均层宽度降低带来计算与存储效率提升: 在拟合的 loss-matched 缩放曲线上, FLOPs 减少 22%; KV 缓存内存及 I/O 成本降低 15%。
• 分析表明, 瓶颈结构改善了 MLP 激活利用率, 缓解中间层表示坍缩 (representation collapse), 残差流中的表示与均匀模型有质性差异。

与基线对比解读

• 标准 Transformer 每层同等预算分配未考虑不同层功能差异; × 形状通过显式容量重分配, 将更多参数集中于输入编码的早期层与输出预测的晚期层, 中间层以窄瓶颈实现压缩与抽象。
• 这种非均匀设计不仅提升语言建模精度, 还直接转化为推理效率增益 (FLOPs 与 KV 缓存均下降), 说明模型缩放不仅依赖“更深更宽”, 层间宽度分布的优化同样关键, 为资源受限部署与大规模模型设计提供了新思路。

行业影响

落地场景

该架构可广泛应用于基于Transformer的生成式AI产品，尤其适合对推理延迟和显存敏感的场景：

• 云端大模型API服务：提供文本生成、代码补全等API的平台，可通过降低单次请求FLOPs提升吞吐量。
• 边缘与端侧部署：手机、IoT设备上的轻量模型，受益于KV cache缩减，能在有限内存下运行更大模型。
• 实时交互系统：智能客服、虚拟助手、实时翻译等，减少延迟改善用户体验。

商业价值

• 降本增效：同等性能下FLOPs减少22%，KV cache/IO成本降低15%，直接转化为GPU使用时长和硬件数量的下降，大幅节约推理成本。
• 体验提升：在延迟敏感业务中，更快的响应速度可提高用户留存与转化率；在批处理场景，更高吞吐量能支撑更大规模业务。
• 灵活扩容：能基于现有硬件服务更多并发请求，或在不增加硬件的情况下支持更长上下文，扩展业务边界。

与现有产品/工作流的接口

• 即插即用：> <former 仅改变各层宽度，无需额外参数或复杂适配，可直接替换现有Transformer block，沿用主流训练框架（如Megatron、HuggingFace）和推理引擎（vLLM、TensorRT-LLM）。
• 渐进式集成：可先对部分业务模型进行试验性替换，因其训练流程、数据、优化器几乎不变，迁移成本极低。
• 兼容MoE：论文验证了在MoE模型（3B参数）上的效果，因此可扩展到当前流行的混合专家架构。

具体落地 Use Case

1. 电商智能客服：某全球电商平台部署LLM作为自动客服，需在促销高峰期处理海量并发对话。采用 > <former 架构的模型，可在相同硬件上承载更高QPS，降低p99延迟，确保会话流畅度，同时节省30%以上的GPU租赁费用。
2. 内容平台AI写作助手：为创作者提供实时续写、摘要生成功能的SaaS工具，在浏览器端或移动端运行模型。缩减的KV cache使得更大模型能在消费级GPU甚至NPU上运行，避免上传数据至云端，实现低延迟本地推理，提升用户隐私和产品差异化。

局限

• **实验规模与通用性有限**：论文仅验证了最大 3B 参数的 decoder-only 语言模型，且主要基于语言建模损失。对于更大规模模型（如数十亿参数）以及不同模态（图像、语音）或任务（生成、推理）的表现仍不明确。文中也承认未探索持续预训练或微调场景下的宽度调度适应性。
• **宽度调度依赖先验搜索**：× 形宽度分配需通过独立实验确定具体的宽度比例和瓶颈位置，不同模型尺寸、数据集和计算预算可能需要重新搜索，缺乏自适应的动态调整机制。此外，瓶颈结构可能对长序列中的信息传递产生负面影响，论文未充分分析在长上下文任务上的表现。
• **与异构计算分配方法对比不足**：本文核心思路是层间非均匀宽度，但未与 Mixture-of-Depths、动态层 skipping 或 conditional computation 等方法直接比较。这些方法同样旨在差异化层间计算，而本文仅与均匀宽度基线对比，难以证明 × 形分配在整体资源效率上的绝对优势。无参数残差调整虽然轻量，却可能限制表示灵活性。