You are looking at the engine room of the transformer. Not the attention heads -- the feed-forward network. The visualization above will let you trace a single input vector as it passes through the two-stage FFN pipeline: the up-projection W1, the nonlinear activation gate, and the down-projection W2. When it arrives, you will be able to see exactly how a 4096-dimensional vector gets expanded to 16384 dimensions, selectively gated, and compressed back -- and why this matters more than most researchers realize. ## The Attention Monopoly on Attention Since [Vaswani et al. (2017)](https://arxiv.org/abs/1706.03762) introduced the transformer, the research community has been overwhelmingly focused on the attention mechanism. Efficient attention, sparse attention, linear attention, flash attention -- the list of attention variants fills entire survey papers. Meanwhile, the feed-forward network (FFN) sub-layers, which typically hold **two-thirds** of a transformer's parameters, receive comparatively little scrutiny. Consider the numbers for a standard 12-layer, d=768 transformer: | Component | Parameters per layer | % of layer total | |-----------|---------------------|------------------| | Q, K, V projections | 3 x 768 x 768 = 1.77M | 24% | | Output projection | 768 x 768 = 0.59M | 8% | | FFN W1 (up-proj) | 768 x 3072 = 2.36M | 32% | | FFN W2 (down-proj) | 3072 x 768 = 2.36M | 32% | | Layer norm + bias | ~6K | <1% | The FFN holds 64% of the parameters. In larger models with 4x expansion ratio, this fraction grows further. The parameters live in FFN, but the papers live in attention. ## FFN as Key-Value Memory

The landmark paper by [Geva et al. (2021)](https://arxiv.org/abs/2012.14913) (arXiv:2012.14913) reframed FFN layers as key-value memories. Look at the computation step by step: ``` FFN(x) = W2 * ReLU(W1 * x + b1) + b2 ``` Each **row** of W1 acts as a "key" -- a pattern detector that fires when the input matches. Each **column** of W2 acts as a "value" -- the information that gets written to the residual stream when that key fires. The ReLU (or GELU) activation between them is the gating mechanism: it decides which keys are active. This interpretation has concrete consequences. Individual FFN neurons encode specific knowledge -- facts, syntactic patterns, semantic categories. [Dai et al. (2022)](https://arxiv.org/abs/2104.08696) (arXiv:2104.08696) demonstrated that you can locate and edit specific factual associations by modifying individual rows of W1 and columns of W2. The ROME and MEMIT methods for model editing ([Meng et al., 2022](https://arxiv.org/abs/2202.05262), arXiv:2202.05262) are built directly on this insight. In the interactive visualization above, you can select any neuron index and see its key pattern (what inputs it responds to) and its value pattern (what it writes to the output). The visual makes the key-value structure immediately tangible. ## What Does Each Layer Compute? If FFN layers are memories, what does each layer remember? The evidence points to a clear progression from shallow to deep: **Early layers** (1-4 in a 12-layer model) store low-level linguistic patterns: n-gram statistics, positional information, basic syntactic categories. Ablation studies show that removing these layers degrades grammaticality more than factual accuracy. **Middle layers** (5-8) handle compositional semantics -- combining word meanings into phrase and clause representations. This is where most attention heads perform their relational work, and the FFN layers consolidate the results. **Late layers** (9-12) store higher-level, more abstract patterns: factual associations, task-specific features, and output formatting. This is consistent with the "residual stream" view (Elhage et al., 2021): the residual connections create a shared communication channel, and each FFN layer writes specialized information into it. Visualize a 12-layer transformer as a 12-floor building. The ground floors handle grammar and syntax. The middle floors handle meaning. The top floors handle facts and formatting. The attention mechanism on each floor decides *which* information to route. The FFN on each floor decides *what* to write. ## The Block-Diagonal Structure of FFN

Recent work has explored whether imposing structure on FFN weight matrices can improve efficiency without sacrificing quality. The visualization above lets you compare different factorization strategies side by side, with real parameter counts and accuracy trade-offs. The Monarch matrix framework ([Dao et al., 2022](https://arxiv.org/abs/2204.00595), arXiv:2204.00595) factors dense matrices into products of block-diagonal and permutation matrices: ``` M = P * diag(B_1, ..., B_k) * P^T * diag(C_1, ..., C_k) ``` This achieves O(n * sqrt(n)) complexity instead of O(n^2). The permutation P between the two block-diagonal factors enables cross-block communication -- the same class of transformations used in FFT algorithms. At NeurIPS 2024, several papers extended this idea with Kronecker-factored FFN layers, showing that structured sparsity can match dense FFN performance at 40-60% of the parameter count. The key insight is that FFN weight matrices are empirically low-rank. [Sharma and Kaplan (2020)](https://arxiv.org/abs/2006.12156) (arXiv:2006.12156) showed that most of the singular values of trained FFN matrices are near zero, suggesting massive redundancy. If the effective rank is low, then structured factorizations that exploit this redundancy should work -- and they do. Here is the parameter savings for a d=4096 FFN layer with 4x expansion: | Method | Parameters | % of Dense | Quality Loss | |--------|-----------|------------|--------------| | Dense | 67.1M | 100% | baseline | | Monarch (k=64) | 8.4M | 12.5% | <0.5% perplexity | | Low-rank (r=256) | 4.2M | 6.3% | ~1% perplexity | | Block-diag (k=16) | 4.2M | 6.3% | 1-3% perplexity | | 2:4 Sparsity | 33.6M | 50% | <0.3% perplexity | ## Why This Matters for the Future of Transformers Understanding FFN layers matters for three practical reasons, each with significant implications: **Model compression.** If FFN layers are the primary parameter cost, then compressing them efficiently (via pruning, quantization, or factorization) is the highest-leverage optimization target. Structured pruning of FFN layers consistently yields better compression ratios than pruning attention heads. The block-diagonal visualization above makes clear *why*: most of the weight matrix is near-zero, and the active regions cluster into blocks. **Knowledge editing.** The key-value memory interpretation enables surgical modification of model behavior without retraining. This is the foundation of model editing, knowledge injection, and fine-grained alignment techniques. When you edit row 4721 of W1 in layer 17, you are editing a single "memory slot" in the transformer's knowledge base. **Architecture design.** If FFN layers are doing most of the "thinking," then we should be spending more architecture-search effort on FFN design and less on attention variants. Mixture-of-experts (MoE) models like Switch Transformer ([Fedus et al., 2022](https://arxiv.org/abs/2101.03961), arXiv:2101.03961) are already exploring this direction, replacing dense FFN with sparse expert selection. The next generation of efficient transformers will likely focus on structured FFN as the primary innovation axis. ## The FFN Computation Pipeline: A Closer Look

Let us walk through the FFN computation for a concrete example. Suppose the input to an FFN layer is the residual stream vector for the token "Paris" in the sentence "The capital of France is Paris." **Step 1: Up-projection.** W1 (shape 4096 x 16384) projects the 4096-dim input to 16384 dimensions. This is the "key matching" step: each of the 16384 neurons computes a dot product with the input. Neurons that detect "is-a-capital" patterns, "European-city" patterns, and "sentence-final-entity" patterns will produce large activations. **Step 2: Activation gate.** GELU zeros out roughly 50-70% of the 16384 activations. Only neurons whose keys matched the input survive. This is the sparsity mechanism -- and it is why FFN matrices are empirically low-rank in their effective computation, even when the weight matrices themselves are dense. **Step 3: Down-projection.** W2 (shape 16384 x 4096) maps the surviving activations back to 4096 dimensions. Each surviving neuron's column in W2 contributes its "value" to the output. The result is a weighted sum of the value vectors of all active neurons. The output gets added to the residual stream via the skip connection. The next layer's attention mechanism can now read the information that this FFN layer wrote. The transformer is not just an attention mechanism with some linear layers bolted on. The FFN sub-layers are where most parameters live, where most knowledge is stored, and where most of the computational work happens. They deserve more of our attention -- no pun intended. ## References - [Vaswani et al. (2017)](https://arxiv.org/abs/1706.03762). Attention Is All You Need. arXiv:1706.03762 - [Geva et al. (2021)](https://arxiv.org/abs/2012.14913). Transformer Feed-Forward Layers Are Key-Value Memories. arXiv:2012.14913 - [Dai et al. (2022)](https://arxiv.org/abs/2104.08696). Knowledge Neurons in Pretrained Transformers. arXiv:2104.08696 - [Meng et al. (2022)](https://arxiv.org/abs/2202.05262). Locating and Editing Factual Associations in GPT. arXiv:2202.05262 - [Dao et al. (2022)](https://arxiv.org/abs/2204.00595). Monarch: Expressive Structured Matrices for Efficient and Accurate Training. arXiv:2204.00595 - [Fedus et al. (2022)](https://arxiv.org/abs/2101.03961). Switch Transformers: Scaling to Trillion Parameter Models. arXiv:2101.03961 - [Sharma & Kaplan (2020)](https://arxiv.org/abs/2006.12156). A Neural Scaling Law from the Dimension of the Data Manifold. arXiv:2006.12156