Flash 多头前馈网络
机器学习
2025-12-09 v1 人工智能
计算与语言
摘要
我们探索了将多头前馈网络 (Multi-Head FFN) 作为变换器架构中前馈网络的替代方案，灵感来源于单头注意力与前馈网络之间的结构相似性。虽然多头机制在注意力中增强了表达能力，但直接将其应用于前馈网络会面临两个挑战：随头数增加而增长的内存消耗，以及在模型规模扩大时，中间维度增长与固定头维度之间比例失衡，这会损害可扩展性和表达能力。为解决这些挑战，我们提出了 Flash 多头前馈网络 (FlashMHF)，其包含两个关键创新：一种类似 FlashAttention 的 I/O 感知融合内核，实时在 SRAM 中计算输出；以及采用动态加权并行子网络设计，维持中间维度与头维度之间的平衡比例。在 1.28 亿至 13 亿参数的模型上验证的 FlashMHF 在困惑度和下游任务准确率上持续优于 SwiGLU 前馈网络，同时将峰值内存使用降低 3-5 倍，推理加速最高可达 1.08 倍。我们的工作确立了多头设计作为前馈网络更优异架构原则，展示了 FlashMHF 作为变换器中前馈网络的强大、高效且可扩展的替代方案。
关键词
引用
@article{arxiv.2512.06989,
title = {Flash Multi-Head Feed-Forward Network},
author = {Minshen Zhang and Xiang Hu and Jianguo Li and Wei Wu and Kewei Tu},
journal= {arXiv preprint arXiv:2512.06989},
year = {2025}
}
备注
17 pages, 8 figures