提出FlashMHF，让前馈网络更高效可扩展
Flash Multi-Head Feed-Forward Network
- 用动态加权并行子网络平衡中间层与头维度
- 推理加速最高1.08倍，内存降低3-5倍
- 适合大模型部署，提升语言模型精度
我们探索将多头前馈网络（MH-FFN）作为Transformer中传统前馈网络（FFN）的替代方案，源于单头注意力与前馈网络在结构上的相似性。尽管多头机制能增强注意力的表达能力，但直接将其应用于前馈网络面临两大挑战：内存消耗随头数线性增长，以及随着模型规模扩大，中间层尺寸增长与固定头维数之间的比例失衡，导致可扩展性和表达力下降。为此，我们提出Flash Multi-Head FFN（FlashMHF），包含两项核心创新：一是基于I/O感知的融合核，在SRAM中在线计算输出，类似FlashAttention；二是采用动态加权的并行子网络设计，维持中间层与头维度间的平衡比例。在参数量从128M到1.3B的模型上验证，FlashMHF在困惑度和下游任务准确率上均优于SwiGLU FFN，同时峰值内存降低3-5倍，推理速度最高提升1.08倍。本工作确立了多头设计作为前馈网络的优越架构原则，为Transformer中的前馈网络提供了高效、可扩展的全新选择。
原文摘要 · Abstract (English)
We explore Multi-Head FFN (MH-FFN) as a replacement of FFN in the Transformer architecture, motivated by the structural similarity between single-head attention and FFN. While multi-head mechanisms enhance expressivity in attention, naively applying them to FFNs faces two challenges: memory consumption scaling with the head count, and an imbalanced ratio between the growing intermediate size and the fixed head dimension as models scale, which degrades scalability and expressive power. To address these challenges, we propose Flash Multi-Head FFN (FlashMHF), with two key innovations: an I/O-aware fused kernel computing outputs online in SRAM akin to FlashAttention, and a design using dynamically weighted parallel sub-networks to maintain a balanced ratio between intermediate and head dimensions. Validated on models from 128M to 1.3B parameters, FlashMHF consistently improves perplexity and downstream task accuracy over SwiGLU FFNs, while reducing peak memory usage by 3-5x and accelerating inference by up to 1.08x. Our work establishes the multi-head design as a superior architectural principle for FFNs, presenting FlashMHF as a powerful, efficient, and scalable alternative to FFNs in Transformers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。