---
title: 'FlashMHF: Multi-Head FFN in Transformers'
url: https://www.emergentmind.com/topics/multi-head-ffn-flashmhf
type: topic
---

# FlashMHF: Multi-Head FFN in Transformers

Multi-Head Feed-Forward Networks (MH-FFN), and in particular the Flash Multi-Head FFN (FlashMHF), represent an evolution of the feed-forward architecture within Transformer models. FlashMHF is designed as a direct replacement for conventional point-wise SwiGLU FFNs in Transformer blocks, introducing a multi-head structure and blockwise fused computation inspired by the symmetry with multi-head self-attention. The central innovations are an I/O-aware kernel that computes outputs online in on-chip SRAM, thereby reducing memory pressure, and the use of dynamically gated parallel sub-networks to preserve the optimal intermediate-to-head dimensionality ratio as models scale—addressing critical scalability and expressivity bottlenecks encountered by naïve multi-head FFN designs. FlashMHF demonstrates improvements in perplexity and downstream task accuracy, significant memory reduction, and inference speedup while maintaining architectural compatibility with standard Transformer compositions [2512.06989].

## 1. Integration into Transformer Architectures

FlashMHF is implemented as a drop-in replacement for traditional FFNs in Transformer blocks. The baseline structure,
$$
\mathrm{LayerNorm} \rightarrow \text{Multi-Head Self-Attention} \rightarrow \mathrm{Add}~\&~\mathrm{Norm} \rightarrow \text{FFN} \rightarrow \mathrm{Add}~\&~\mathrm{Norm},
$$
originally computes the FFN as
$$
\mathrm{FFN}(X) = \left(X W_{\text{up}} \odot \mathrm{SiLU}(X W_{\text{gate}})\right) W_{\text{down}},
$$
with input $X \in \mathbb{R}^{B \times N \times d_{\text{model}}}$. FlashMHF replaces this component, applying $O = \mathrm{FlashMHF}(X)$ with the same input-output dimensionality and identical placement of normalization and residual connections. This ensures seamless compatibility at all integration points.

The multi-head structure mirrors multi-head attention, assigning $H$ heads each with head dimension $d_h$ such that $d_{\text{model}} = H \cdot d_h$. Where in attention heads receive per-head queries $Q_h$, keys $K_h$, and values $V_h$, a naïve MH-FFN would analogously compute per-head outputs
$$
S_h = \left(\mathrm{SiLU}(Q_h K_h^\top) \odot (Q_h U_h^\top)\right) V_h,
$$
aggregating head outputs via concatenation and output projection. However, without architectural constraints, this approach rapidly incurs prohibitive memory and parameter inefficiencies as model width scales [2512.06989, Sect. 3.1].

## 2. Fused SRAM-Oriented Kernel Design

The FlashMHF core computation is dominated by a blockwise kernel intentionally structured to run within on-chip SRAM, minimizing reading/writing of large intermediate tensors from high-bandwidth memory (HBM). The operational sequence for one head and one of its sub-networks with input $Q \in \mathbb{R}^{L \times d_h}$, learned parameters $K, U, V \in \mathbb{R}^{d_e \times d_h}$ (with $d_e$ the sub-network width), and chunking $d_e$ into $M$ blocks of size $b$, is:
$$
O \leftarrow 0 \in \mathbb{R}^{L \times d_h}\\
\text{for } m=1\ldots M: \\
\quad K_m = K[m \cdot b:(m+1) \cdot b, :]\\
\quad U_m = U[m \cdot b:(m+1) \cdot b, :]\\
\quad V_m = V[m \cdot b:(m+1) \cdot b, :]\\
\quad M = Q K_m^\top\\
\quad N = Q U_m^\top\\
\quad A = \mathrm{SiLU}(M) \odot N\\
\quad O += A V_m.
$$
This structure (see Eq. 9 and Algorithm A.1 in [2512.06989]) ensures no full $L \times d_{\text{ff}}$ intermediate activation is ever materialized off-chip. Rather, intermediate results are accumulated in SRAM registers, with only the final output $O$ written to slower memory, a methodology analogous to FlashAttention’s streaming approach. The forward pass pseudocode manages batch, head, and block dimensions as described, with gating weights dynamically introduced per block.

## 3. Dynamically Weighted Parallel Sub-Networks

To address scaling pathologies inherent in naïve multi-head FFN splits, such as an exploding ratio $d_{\text{ff}}/d_h$ as $d_{\text{model}}$ grows, FlashMHF partitions the intermediate FFN dimension across $E$ parallel sub-networks within each head, giving $d_{\text{ff}}=E d_e$ with $d_e \approx \alpha d_h$ ($\alpha \approx 8/3$ as with SwiGLU). This enforces an optimal and consistent $d_{\text{ff}}:d_h$ ratio independently of head count or model scale.

Each head $h$ learns a small gating matrix $W^h \in \mathbb{R}^{d_h \times E}$, generating per-token, per-head sub-network selection weights:
$$
P^h = Q_{:, h, :}\, W^h \in \mathbb{R}^{L \times E},\\
R^h_{\ell, e} = \frac{\sigma(P^h_{\ell, e})}{\sum_{e'} \sigma(P^h_{\ell, e'}) + \epsilon}
$$
($\sigma$ is the elementwise sigmoid, Eq. 8). For token $\ell$, head $h$, and sub-network $e$, each sub-network computes a gated SwiGLU-like output which is then mixed according to $R^h_{\ell, e}$, and finally all per-head outputs are concatenated and projected back to $d_{\text{model}}$ as standard.

## 4. Theoretical and Practical Resource Analysis

FlashMHF maintains the overall compute complexity of the underlying FFN transformation but dramatically alters peak memory requirements. For batch size $B$ and sequence length $N$:
- **Standard SwiGLU FFN:** peak activation memory is $O(B N (d_{\text{model}} + d_{\text{ff}} + d_{\text{model}}))$ due to storing $A \in \mathbb{R}^{B N \times d_{\text{ff}}}$.
- **FlashMHF:** peak memory reduces to $O(B N d_{\text{model}})$; no activation of size $d_{\text{ff}}$ is stored off-chip, only current block-level inputs and outputs plus gating terms.

Empirically, on H100/Hopper GPUs, FlashMHF reduces peak HBM by 3–5× across a broad range of sequence lengths (192–16K tokens), and achieves up to 1.08× inference speedup on long contexts (average ~1.05×). The speedup magnitude is modest relative to FlashAttention due to the baseline efficiency of cuBLAS-based FFN implementations and FFN output bandwidth constraints, but the reduction in activation memory is substantial [2512.06989, Figs. 7a, 7b].

## 5. Empirical Evaluation

FlashMHF was evaluated across language models of three scales: ~128M, ~370M, and ~1.3B parameters. All models were trained on The Pile with context length 4096, batch size 64, and GPT-NeoX tokenization, using baseline LLaMA-style attention and SwiGLU FFNs for direct comparison. Downstream evaluations employed six common benchmarks (HellaSwag, Social IQA, Physical IQA, OpenBookQA, WinoGrande, RACE).

Key results include:
- **Perplexity (PG19):** At 370M, FlashMHF achieves perplexity 3.014 vs. 3.030 for baseline; at 1.3B, 2.793 vs. 2.843. The reduction (~0.85 ppl at 1.3B) is consistent across scales.
- **Downstream accuracy:** FlashMHF-128hdim attains 40.48% average (370M) and 43.35% (1.3B) compared to baselines of 39.92% and 41.75%, respectively.
- **Efficiency:** Consistent memory reduction (3–5×) and up to 8% inference speedup on long-context, deep models.
- **Robustness:** Naïve MH-FFN variants fail to scale past 128M model size, and ablations (e.g., ParamKV) underperform, underscoring the necessity of both relational structure and parallel sub-network design [2512.06989, Table 2, Table 3, Fig. 6].

## 6. Design Considerations, Limitations, and Extensions

FlashMHF introduces kernel and implementation complexity exceeding that of standard cuBLAS FFN calls, particularly on Triton and Hopper architectures. For small models or short sequence lengths, the overhead may outweigh gains. The selection of $d_h$ (head dimension) is a crucial hyperparameter: small $d_h$ risks under-capacity, large $d_h$ reduces head count and diminishes representational diversity.

Potential extensions include exploring alternative gating functions (such as sparse top-$k$ MoE), fusing upstream operations (e.g., LayerNorm+FlashMHF) into a single kernel, and adapting the architecture to non-GPU hardware (TPU, SW). Joint optimization of $H$, $E$, and $d_h$ is motivated for different scaling regimes.

A plausible implication is that multi-head partitioning is emerging as a broadly superior architectural principle in both attention and feed-forward pathways of large-scale Transformers, as it enables improved parameter and computational efficiency, more stable scaling, and substantial reduction in memory footprints, all while preserving end-to-end training and inference semantics [2512.06989].

## 7. Summary

Flash Multi-Head FFN (FlashMHF) replaces conventional SwiGLU FFNs in Transformer blocks with a multi-head, dynamically gated, blockwise computation framework. Its design maintains the optimal intermediate-to-head dimensional relationship via parallel sub-networks, fuses computation to SRAM to eliminate off-chip intermediate tensors, and empirically improves both language modeling and downstream performance metrics. FlashMHF achieves up to 1.08× inference acceleration and 3–5× activation memory reduction, representing a scalable and efficient alternative for modern Transformer architectures [2512.06989].

Source: https://www.emergentmind.com/topics/multi-head-ffn-flashmhf