Flash Multi-Head Feed-Forward Network
This work addresses a bottleneck in scaling Transformer models for AI researchers and practitioners by offering a more efficient and scalable FFN alternative.
The paper tackles the inefficiency and scalability issues of applying multi-head mechanisms to feed-forward networks (FFNs) in Transformers by proposing Flash Multi-Head FFN (FlashMHF), which improves perplexity and downstream task accuracy while reducing peak memory usage by 3-5x and accelerating inference by up to 1.08x in models from 128M to 1.3B parameters.
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.