Flash Multi-Head Feed-Forward Network

Fuente: arXiv

Saved in:

Bibliographic Details
Main Authors:	Zhang, Minshen, Hu, Xiang, Li, Jianguo, Wu, Wei, Tu, Kewei
Format:	Preprint
Published:	2025
Subjects:	Machine Learning Artificial Intelligence Computation and Language I.2.7
Online Access:	Acceder al recurso
Tags:	Add Tag No Tags, Be the first to tag this record!

_version_	1866914185915400192
author	Zhang, Minshen Hu, Xiang Li, Jianguo Wu, Wei Tu, Kewei
author_facet	Zhang, Minshen Hu, Xiang Li, Jianguo Wu, Wei Tu, Kewei
contents	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.
format	Preprint
id	arxiv_https___arxiv_org_abs_2512_06989
institution	arXiv
publishDate	2025
record_format	arxiv
spellingShingle	Flash Multi-Head Feed-Forward Network Zhang, Minshen Hu, Xiang Li, Jianguo Wu, Wei Tu, Kewei Machine Learning Artificial Intelligence Computation and Language I.2.7 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.
title	Flash Multi-Head Feed-Forward Network
topic	Machine Learning Artificial Intelligence Computation and Language I.2.7
url	https://arxiv.org/abs/2512.06989