FlashMHF: Advanced Multi-Head FFN

• FlashMHF introduces multi-head decomposition in feed-forward networks, enhancing expressivity and reducing memory bottlenecks compared to standard FFNs.
• It employs an I/O-aware fused kernel that streams parameter tiles, achieving a 3–5× reduction in activation memory usage for large-scale models.
• Parallel sub-networks with optimal dimensional balancing enable improved perplexity and downstream accuracy, making FlashMHF a superior alternative in Transformer architectures.

Flash Multi-Head Feed-Forward Network (FlashMHF) represents an advancement in Transformer architecture, specifically targeting the limitations of conventional feed-forward networks (FFNs) by introducing multi-head decomposition and parallel sub-network structures with an I/O-aware computation kernel. FlashMHF is motivated by the structural analogy between single-head attention and FFN layers, aiming to enhance expressivity, scalability, and computational efficiency while mitigating activation memory bloat and performance degradation at scale. Empirical evidence on models ranging from 128 million to 1.3 billion parameters establishes FlashMHF as an efficient and powerful drop-in alternative to traditional FFNs in Transformers, outperforming SwiGLU FFNs in perplexity, downstream accuracy, memory usage, and inference speed (Zhang et al., 7 Dec 2025).

1. Architectural Foundations and Motivations

Transformer FFNs and single-head attention share a mathematical symmetry: attention computes $$\text{softmax}\bigl(\frac{QK^T}{\sqrt{d_k}}\bigr) V$$, while FFN computes $\phi\bigl(X W_1^T\bigr) W_2$, with $W_1$ and $W_2$ respectively analogous to attention keys and values and $\phi$ serving as an element-wise nonlinearity replacing the row-wise softmax. This observation motivates introducing multi-head decomposition in the FFN module, in close analogy to the expressivity achieved by multi-head attention.

Naïvely extending FFNs to a multi-head structure (MH-FFN) faces two major obstacles:

• Memory Consumption: Activation memory scales as $O(L \cdot H \cdot d_{ff})$, where $H$ is the number of heads, $L$ the sequence length, and $d_{ff}$ the intermediate width.
• Scaling Imbalance: As Transformer models grow, practitioners typically hold the per-head width $d_h$ fixed and allow $d_{ff}$ to scale, causing the ratio $R = d_{ff} / d_h$ to balloon (from 16 to 45 in tested scales), empirically degrading performance away from the optimal regime ($R \approx 8/3$).

2. Mathematical Formulation and Kernel Design

FlashMHF overcomes these challenges via two central innovations.

2.1 Parallel Sub-Networks per Head

Each FFN head is further split into $E$ parallel SwiGLU sub-networks, with each sub-network width $d_e \approx (8/3)d_h$, restoring the optimal $d_{ff}/d_h$ ratio within each head. For input $X$, projection $$Q = \text{split}_H(X W_{\text{in}}) \in \mathbb{R}^{L \times H \times d_h}$$, weights $K, U, V$ are assigned per sub-network and head.

The per-token output for head $h$ is: $$S_{\ell, h, :} = \sum_{e=1}^{E} R^h_{\ell, e} T^h_{\ell, e}$$ where

$$T^h_{\ell, e} = (\text{SiLU}(Q_{\ell, h, :} K_e^{h T}) \odot (Q_{\ell, h, :} U_e^{h T})) V_e^h$$

and dynamic gating weights $R^h$ are computed via softmax-normalized logits.

2.2 I/O-Aware Fused Kernel

Unlike conventional FFN implementations that materialize the entire $L \times d_{ff}$ activation in high-bandwidth memory (HBM), FlashMHF streams small tiles of sub-network parameters in/out of memory and computes outputs directly in on-chip SRAM: $$O \leftarrow 0, \quad \text{for } m = 1, \ldots, M: \quad O += (\text{SiLU}(Q K_m^T) \odot (Q U_m^T)) V_m$$ This blockwise streaming reduces peak memory usage from $O(L \cdot d_{ff})$ to $O(L \cdot d_{model})$, achieving a 3–5× reduction in activation memory.

3. Implementation Strategies: Tiling and Dimensional Balancing

FlashMHF employs Triton kernels that parallelize computation across batch, head, and sequence blocks. For each block (B, h, seq_block), an accumulator $O_{\text{acc}}$ of size $BLOCK_{SEQ} \times d_h$ is allocated in SRAM, and parameter tiles $$(K_{\text{tile}}, U_{\text{tile}}, V_{\text{tile}})$$ are prefetched to minimize latency. The fused GEMMs and gating weights are computed on-the-fly, with outputs written to global memory.

On NVIDIA Hopper architectures, ThunderKittens kernels further exploit warp groups and multi-stage ring buffers for efficient prefetch and latency hiding.

Dimensional balancing is achieved by setting $d_e \simeq (8/3)d_h$ and choosing the number of sub-networks $E$ such that $d_{ff} = E \cdot d_e$ matches the desired total intermediate width. As model scale grows, both $H$ and $d_h$ can be increased with $E$ recalibrated to preserve the internal ratio.

4. Complexity and Memory Analysis

The computational complexity (FLOPs) of FlashMHF matches standard SwiGLU FFNs, as both amount to $O(3 L d_{model} d_{ff})$.

Space Complexity Comparison

FFN Variant	Peak Memory Usage	Reduction Factor
Standard FFN	$O(L \cdot d_{ff})$	1×
Naïve MH-FFN	$O(L \cdot H \cdot d_{ff})$	$\sim$H×
FlashMHF	$O(L \cdot d_{model})$	3–5×

FlashMHF requires only $O(L d_{model})$ for activations, enabling efficient scaling to long sequences and large models.

5. Empirical Evaluation and Benchmarking

FlashMHF was evaluated on The Pile using models of 128M, 370M, and 1.3B parameters. Key experimental settings include context length 4096, batch size 64, and training on up to 100B tokens.

Perplexity and Downstream Accuracy

Model ± Scale	370M Loss	1.3B Loss
Baseline SwiGLU	3.030	2.843
FlashMHF (128)	3.014	2.793

For downstream zero-shot tasks (HellaSwag, SIQA, PIQA, OBQA, WinoGrande, RACE):

Model ± Scale	370M Avg. Acc.	1.3B Avg. Acc.
Baseline	39.92	41.75
FlashMHF (128)	40.48	43.35

FlashMHF-128 exhibits up to 0.85 perplexity point improvement and 1.6% absolute accuracy gain at 1.3B scale compared to baseline.

Latency and Memory

• Peak activation memory reduced by 3–5× across sequence lengths 192–16k.
• Latency improvements of $\sim$1.05–1.08× over SwiGLU.

6. Implications, Limitations, and Future Directions

FlashMHF advances Transformer FFN design by enabling multi-head feature disentanglement and implicit beam search over reasoning paths through parallel sub-networks. The scale-balanced configuration preserves optimal representational ratios, supporting superior expressivity and scalability. The I/O-aware kernel paradigm eliminates activation memory bottlenecks without increasing FLOPs.

Limitations include modest latency gains compared to FlashAttention—potentially addressable by deeper hardware-level optimization. Dense per-token gating is used; incorporating sparse expert selection may further reduce computational overhead. Automatic architectural tuning (e.g., determining optimal $H$, $d_h$, $E$ per layer) and extending the blockwise kernel approach to additional Transformer components represent promising avenues.

A plausible implication is enhanced feasibility of deploying larger sequence lengths and model scales within fixed hardware budgets. Extending these principles may benefit other parameter-heavy substructures in neural architectures.

7. Summary and Positioning Within Transformer Research

FlashMHF synthesizes the insight that FFNs structurally resemble attention and generalizes multi-head modularity to the feed-forward context. By uniting scale-balanced parallel SwiGLU sub-networks with blockwise I/O-aware fused kernels, it provides a memory-efficient, scalable, and empirically superior alternative to conventional FFNs. As validated on Transformer models from 128M to 1.3B parameters, FlashMHF consistently delivers improved perplexity, accuracy, inference speed, and memory usage, representing a notable architectural innovation in the Transformer family (Zhang et al., 7 Dec 2025).