Flash Multi-Head Feed-Forward Network (2512.06989v1)

Abstract: 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.

• The paper presents FlashMHF, which improves FFN expressivity and computational efficiency by balancing the per-head expansion ratio.
• It introduces a dense-MoE-inspired partition of the intermediate dimension and a custom I/O-aware flash kernel that reduces activation memory by 3–5× and speeds up inference.
• Empirical analysis across Llama-like architectures shows consistent gains in validation loss, perplexity, and downstream task performance at multiple scales.

Flash Multi-Head Feed-Forward Network: Architecture and Empirical Analysis

Introduction and Motivation

The paper "Flash Multi-Head Feed-Forward Network" (2512.06989) systematically investigates the role of feed-forward networks (FFNs) within the Transformer architecture, focusing on their expressive limitations and the scalability challenges inherent to standard and multi-head designs. The authors emphasize the structural symmetry between single-head attention and FFN, recognizing both as parameterized transformations over input queries, but with attention using a row-wise softmax and the FFN applying an elementwise nonlinearity. This observation motivates the central premise: that multi-head designs, which are critical for expressivity in self-attention, may also be beneficial for FFNs.

However, simple multi-head adaptations to FFNs (MH-FFNs) present two acute challenges. First, activation memory grows linearly with the head count, quickly becoming a bottleneck at scale. Second, a scaling imbalance emerges because the intermediate FFN dimension ($d_\text{ff}$) expands with model size while the per-head dimension ($d_h$) remains fixed, leading to performance degradation. The paper sets out to address these via the Flash Multi-Head FFN (FlashMHF), introducing a scale-balanced architecture and a custom I/O-aware kernel that yields both quality and efficiency gains.

Architectural Formulation: FlashMHF

FlashMHF is designed around two critical advances: (1) a scale-balanced parallel FFN sub-network architecture and (2) an I/O-aware blockwise computation kernel.

In FlashMHF, the FFN's intermediate dimension is partitioned into $E$ parallel sub-networks for each head. This dense-MoE-inspired structure allows $d_e$, the dimension per sub-network, to be set so that the expansion ratio per head ($d_e/d_h$) remains optimal across scales, preserving representational expressivity. Unlike prior MH-MoE architectures, FlashMHF enables independent parameterization of all heads, avoiding expressivity loss.

A learned gating matrix produces a token-dependent soft selection over sub-networks for each head, enabling dynamic aggregation and further increasing expressive capacity.

The critical computational innovation is a blockwise flash kernel analogous to FlashAttention, which computes FFN activations in SRAM, never materializing large intermediate tensors in HBM. This significantly lowers memory requirements and reduces inference latency, allowing both longer sequences and larger models to be run on constrained hardware.

Empirical Performance and Ablation Analysis

FlashMHF was experimentally validated on standard Llama-like architectures, comparing against strong SwiGLU FFN baselines, parametric key-value (PKV) attention, dense MoE (with $H=1$), and naïve MH-FFN variants. Experiments spanned model sizes from 128M to 1.3B parameters.

Language Modeling and Scaling

In both 128M and 370M regimes (Figure 1), FlashMHF surpasses the baseline on final validation loss for PG19, with the following key findings:

Figure 1: Comparative evaluation loss of Baseline, Parametric KV, FlashMHF, and MH-FFN for 128M and 370M scales, demonstrating FlashMHF's consistent gains in both regimes.

• FlashMHF maintains advantages as model scale increases, whereas naïve MH-FFN loses competitiveness and Dense-MoE underperforms beyond small scale, reinforcing the necessity of ratio balancing.
• FlashMHF with $d_h = 128$ exhibits the best train and evaluation loss curves in the 370M setting, confirming the balance between per-head expressivity and subspace diversity.

Head dimension ablations reveal a clear tradeoff: small $d_h$ values lead to representational bottlenecks, while larger $d_h$ reduce the number of heads and specialization. Moderate $d_h$, such as $128$, yields the best empirical results (Figures 3 and 4).

Figure 2: Training loss curves for various head dimensions across model scales, illustrating FlashMHF's superior convergence behavior.

Figure 3: Training loss for FlashMHF at increasing model scales, showcasing stable and efficient optimization.

At the 1.3B scale, FlashMHF maintains its margin in perplexity and convergence speed, with a consistent advantage across scaling regimes.

Downstream Task Generalization

FlashMHF variants outperformed baseline FFN models across a suite of reasoning and reading comprehension datasets (HellaSwag, SIQA, PIQA, OpenBookQA, WinoGrande, RACE). The best results were consistently achieved by FlashMHF with $d_h = 128$, demonstrating improved general capabilities.

Memory and Latency Benchmarks

A central empirical claim is the drastic reduction in peak memory by 3-5x and inference speedups up to 1.08x, especially as sequence length increases. Figure 4 presents the memory consumption profile as a function of input sequence length.

Figure 4: Memory usage versus sequence length for SwiGLU FFN, MH-FFN, and FlashMHF, highlighting FlashMHF's superior scalability and efficiency.

These gains directly facilitate deployment of longer context windows and larger parameter regimes on the same hardware.

Theoretical and Practical Implications

The results challenge the standard FFN architectural paradigm in Transformers, indicating that multi-head, scale-balanced FFNs can offer both higher expressivity and superior efficiency. The separation between activation functions in attention (softmax) and FFN (elementwise) is bridged via key-value reinterpretations of FFN computations. FlashMHF's efficiency benefits have direct consequences for LLM deployment, enabling longer sequences, faster inference, and reduced hardware demands.

Furthermore, the dense-MoE style parallelism and token-wise dynamic routing connect FFN design with recent expert models, but avoid their data movement and routing overhead. The I/O-aware kernel represents a generalizable optimization direction for other memory-intensive architectural components.

FlashMHF's architecture demonstrates strong transferability across scales, encouraging reevaluation of model design for multi-modal and generative tasks, where memory and latency bottlenecks remain prominent.

Future Directions

Potential avenues for extending FlashMHF include:

• Exploring token-wise or context-dependent expert activation patterns and sparse gating for further efficiency.
• Integrating FlashMHF with modular architectures beyond language modeling (e.g., vision, multi-modal).
• Investigating theoretical properties of multi-head FFN routing in the context of calibration and sample efficiency.
• Applying the blockwise flash kernel paradigm to other bottleneck layers in large models.

Conclusion

FlashMHF establishes a new architectural principle for FFN layers within Transformers, resolving scaling and memory bottlenecks while providing consistent accuracy improvements over established baselines. The integration of scale-balanced parallel sub-networks and an I/O-aware kernel demonstrably enhances both model quality and computational efficiency. FlashMHF represents a compelling successor for dense FFN modules and signals an inflection point in the design of future LLMs and large-scale neural architectures.