Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Gated Linear Units › Implementation: A Complete SwiGLU Module

Michael BrenndoerferUpdated June 14, 20253 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 109Gated Linear Units › Implementation: A Complete SwiGLU Module

Reading section 11 of 16 · Implementation: A Complete SwiGLU Module · 3 minSection 11 of 16 · 3 min

Section 11 of 16 · 3 min read

Gated Linear Units › Implementation: A Complete SwiGLU ModuleOne section · Sections in this chapter11 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Having traced through the mathematics by hand, we can now consolidate everything into a reusable implementation. This module follows the patterns used in production transformer libraries: it initializes weights using proper scaling (Xavier/Glorot initialization), supports optional biases, handles both single vectors and batched sequences, and provides a clean interface for integration into larger architectures.

The implementation below captures the complete SwiGLU FFN used in modern LLMs such as LLaMA and Mistral:

In[30]:

Code

class SwiGLUFeedForward:
    """
    Feed-forward network with SwiGLU activation.

    Implements: FFN(x) = (Swish(xW) * xV) @ W2

    This is the architecture used in LLaMA, Mistral, and other modern LLMs.
    """

    def __init__(self, d_model, d_ff, use_bias=False):
        """
        Initialize the SwiGLU feed-forward network.

        Args:
            d_model: Input and output dimension
            d_ff: Hidden dimension (note: requires 3x params vs standard FFN)
            use_bias: Whether to include bias terms (modern LLMs often omit)
        """
        self.d_model = d_model
        self.d_ff = d_ff
        self.use_bias = use_bias

        # Xavier/Glorot initialization
        scale = np.sqrt(2.0 / (d_model + d_ff))
        self.W = np.random.randn(d_model, d_ff) * scale
        self.V = np.random.randn(d_model, d_ff) * scale
        self.W2 = np.random.randn(d_ff, d_model) * scale

        if use_bias:
            self.b_w = np.zeros(d_ff)
            self.b_v = np.zeros(d_ff)
            self.b_out = np.zeros(d_model)

    def __call__(self, x):
        """
        Apply SwiGLU feed-forward transformation.

        Args:
            x: Input tensor of shape (..., d_model)

        Returns:
            Output tensor of shape (..., d_model)
        """
        # Compute both projections
        gate_input = x @ self.W
        linear_input = x @ self.V

        if self.use_bias:
            gate_input = gate_input + self.b_w
            linear_input = linear_input + self.b_v

        # SwiGLU: Swish on gate path, multiply with linear path
        hidden = swish(gate_input) * linear_input

        # Project back to model dimension
        output = hidden @ self.W2

        if self.use_bias:
            output = output + self.b_out

        return output

    def num_parameters(self):
        """Return total parameter count."""
        params = 3 * self.d_model * self.d_ff
        if self.use_bias:
            params += 2 * self.d_ff + self.d_model
        return params


# Test the module
swiglu_module = SwiGLUFeedForward(
    d_model=512, d_ff=1365
)  # Reduced d_ff for param matching

# Single vector
x_test_mod = np.random.randn(512)
y_test_mod = swiglu_module(x_test_mod)

# Batch
X_test_mod = np.random.randn(16, 512)
Y_test_mod = swiglu_module(X_test_mod)

Out[31]:

Console

SwiGLUFeedForward module test:

Configuration:
  d_model:    512
  d_ff:       1365
  use_bias:   False
  Parameters: 2,096,640

Single vector:
  Input shape:  (512,)
  Output shape: (512,)

Batch (sequence):
  Input shape:  (16, 512)
  Output shape: (16, 512)

Comparison to standard FFN (d_ff=2048):
  Standard FFN params: 2,097,152
  SwiGLU FFN params:   2,096,640

The SwiGLU module with reduced $d_{ff} = 1365$ achieves comparable parameter count to a standard FFN with $d_{ff} = 2048$ . The module correctly handles both single vectors and batched sequences, producing outputs of the expected shapes. This implementation follows the pattern used in production LLMs like LLaMA.