Search

Try a topic

Notebook

GELU / SwiGLU from Scratch

Code note Deep learning

import torch.nn as nn
import torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up   = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

Why it matters: The gate is a content-dependent multiplier on each hidden unit — the FFN learns which features to let through, not just how to mix them. Trap: vanilla FFN uses d_ff = 4 * d_model (2 matrices); SwiGLU has 3, so to match parameter count use d_ff ≈ (8/3) * d_model — Llama, Qwen and Mistral all do this. Copying 4x into a SwiGLU silently inflates the block by ~50%.

Full implementation in rlvr-from-scratch/model/ffn.py.