Tutorial 08: Feed-Forward Networks and SiLU Gating

After attention gathers information across tokens, the feed-forward network (FFN) processes each token independently. The FFN is where most of the model's knowledge is stored. This tutorial covers the standard FFN, the gated variant with SiLU activation, and how they fit together.

---

1. The Feed-Forward Network

Every transformer block contains an FFN after the attention sublayer. The basic structure is:

FFN(x) = W_down @ activation(W_up @ x)

where: - W_up projects from hidden_size to intermediate_size (typically 4× larger) - W_down projects from intermediate_size back to hidden_size - activation introduces nonlinearity (without it, two linear layers collapse to one: W2 @ (W1 @ x) = (W2 @ W1) @ x)

Dimensions

For Qwen3-8B with hidden_size=4096 and intermediate_size=14336:

W_up:   (14336, 4096)   — expand 4×
W_down: (4096, 14336)   — compress back

The expansion ratio of 4× is a design choice. Larger ratios give the model more capacity but cost more memory and compute.

---

2. Why Gating?

Modern LLMs (GPT-3, Llama, Qwen, Mistral) do not use a plain FFN. They use a gated FFN with two parallel branches:

GatedFFN(x) = W_down @ (silu(W_gate @ x) * W_up @ x)

The key idea is the * (elementwise multiply):

• W_gate @ x produces the gate — a set of values that control which features pass through.
• W_up @ x produces the up-projection — the actual content.
• silu(gate) * up multiplies them together: the gate controls how much of each feature flows through.

This is like a learned switch for each neuron in the intermediate layer.

---

3. The SiLU Activation

SiLU (Sigmoid Linear Unit), also known as Swish, is defined as:

silu(x) = x * σ(x) = x / (1 + exp(-x))

where σ(x) is the sigmoid function.

Properties

• For large positive x: silu(x) ≈ x (passes through)
• For large negative x: silu(x) ≈ 0 (suppresses)
• For x = 0: silu(0) = 0.5 * 0 = 0
• Smooth and differentiable everywhere

Numerically Stable Implementation

The naive x * (1 / (1 + exp(-x))) can overflow when x is a large negative number (exp(-(-1000)) = exp(1000) = inf). The stable version uses a trick:

def silu(x: mx.array) -> mx.array:
    """
    SiLU activation: x * sigmoid(x), numerically stable.

    For x >= 0: sigmoid(x) = 1 / (1 + exp(-x))  — standard form
    For x < 0:  sigmoid(x) = exp(x) / (1 + exp(x))  — avoids exp(large positive)

    We compute sigmoid using abs(x) to avoid overflow in both directions,
    then use mx.where to select the numerically stable branch.
    """
    z = mx.exp(-mx.abs(x))

    # For x >= 0: 1 / (1 + exp(-x)) is fine (exp(-positive) is small)
    # For x < 0:  exp(x) / (1 + exp(x)) avoids exp(positive) overflow
    sigmoid = mx.where(x < 0, z / (1 + z), 1 / (1 + z))

    return x * sigmoid

Why SiLU Instead of ReLU?

ReLU: max(0, x) — simple but has a "dead zone" where gradients are zero.

SiLU: x * σ(x) — smooth, never exactly zero, and the gradient flows throughout. Empirically, SiLU produces better language models.

---

4. SwiGLU: The Fused Kernel

The gated FFN computes silu(gate) * up. When the Metal extension is available, we can fuse this into a single GPU kernel instead of doing it in two separate operations:

# Pure Python path (two operations):
result = silu(gate) * up

# Fused Metal kernel (one operation):
result = swiglu(gate, up)

The fused kernel is faster because: 1. It reads gate and up from memory once (not twice). 2. The elementwise multiply happens in registers, not in global memory. 3. One GPU dispatch instead of two.

In the MLP layer:

class MLP:
    def __init__(self, layer_prefix, config, weights, features):
        self.w_gate = weights.linear(f"{layer_prefix}.mlp.gate_proj")
        self.w_up = weights.linear(f"{layer_prefix}.mlp.up_proj")
        self.w_down = weights.linear(f"{layer_prefix}.mlp.down_proj")
        self.use_fast_swiglu = features.kernels  # True if Metal extension is available

    def __call__(self, x):
        gate = linear(x, self.w_gate)
        up = linear(x, self.w_up)

        # Choose between fused and unfused paths
        if self.use_fast_swiglu:
            hidden = swiglu(gate, up)      # single Metal kernel
        else:
            hidden = silu(gate) * up        # two MLX operations

        return linear(hidden, self.w_down)

---

5. Why Three Weight Matrices?

The gated FFN has three weight matrices (gate_proj, up_proj, down_proj) instead of the two in a standard FFN. This seems wasteful, but the math shows why it is worth it:

Standard FFN: 2 × hidden × intermediate parameters Gated FFN: 3 × hidden × intermediate parameters

The extra matrix buys nonlinearity in the gating mechanism. Without it, the gating is just a linear transform, which collapses:

W_down @ ((W_gate @ x) * (W_up @ x))  ← nonlinear (gated)
W_down @ (W_gate @ x * W_up @ x)       ← still nonlinear

If we tried to merge W_gate and W_up into one matrix, the elementwise multiply would produce a quadratic function of x — but only in a very specific way that limits expressiveness. Three separate matrices give the model full control over what to gate and what to project.

---

6. Complete FFN Code

import mlx.core as mx

def linear(x, w, bias=None):
    """The foundation of every layer: y = x @ w^T (+ bias)."""
    if bias is not None:
        return mx.matmul(x, w.T) + bias
    return mx.matmul(x, w.T)


def silu(x):
    """SiLU activation: x * sigmoid(x), numerically stable."""
    z = mx.exp(-mx.abs(x))
    sigmoid = mx.where(x < 0, z / (1 + z), 1 / (1 + z))
    return x * sigmoid


class MLP:
    """
    Gated feed-forward network with SiLU activation.

    Architecture: down(silu(gate(x)) * up(x))

    This is the standard FFN in Llama, Qwen, Mistral, and most modern LLMs.
    """

    def __init__(self, config, gate_weight, up_weight, down_weight):
        """
        Args:
            config: Model config (for dimensions).
            gate_weight: Weight for the gate projection. Shape: (intermediate, hidden)
            up_weight: Weight for the up projection. Shape: (intermediate, hidden)
            down_weight: Weight for the down projection. Shape: (hidden, intermediate)
        """
        self.w_gate = gate_weight
        self.w_up = up_weight
        self.w_down = down_weight

    def __call__(self, x):
        """
        Forward pass.

        Args:
            x: Input tensor. Shape: (batch, seq_len, hidden_size)

        Returns:
            Output tensor. Shape: (batch, seq_len, hidden_size)
        """
        # Step 1: Compute gate and up projections (both are linear layers)
        gate = linear(x, self.w_gate)   # (B, L, intermediate_size)
        up = linear(x, self.w_up)       # (B, L, intermediate_size)

        # Step 2: Gate the features
        # silu(gate) produces values in [0, ...] that control feature flow
        # Multiplying by up selects and scales features
        hidden = silu(gate) * up         # (B, L, intermediate_size)

        # Step 3: Project back to hidden dimension
        output = linear(hidden, self.w_down)  # (B, L, hidden_size)
        return output

---

7. Where FFN Fits in the Transformer Block

x ──→ RMSNorm ──→ Attention ──→ (+) ──→ h
                                        │
                                        └──→ RMSNorm ──→ FFN ──→ (+) ──→ output

The FFN processes each token independently — there is no interaction between tokens in the FFN. This is what makes it different from attention. The attention layer mixes information across tokens; the FFN processes each token's representation in isolation.

---

8. Summary

Component	Purpose	Parameters
gate_proj	Controls which features pass through	hidden × intermediate
up_proj	Provides the content to be gated	hidden × intermediate
down_proj	Compresses back to model dimension	intermediate × hidden
SiLU	Nonlinear gating activation	0 (it is an operation)

The gated FFN is where most of the model's "knowledge" lives. The weights encode factual knowledge, linguistic patterns, and reasoning capabilities. The attention layer merely determines which tokens to combine — the FFN is where the actual processing happens.

Next: Tutorial 09 — Transformer Block and Full Model Assembly