Skip to content

← Flax step 48 of 100

Medium Primitives

SwiGLU Feed-Forward Network

Jump to the task

Multiply two learned feature paths

A conventional feed-forward network projects an input into a hidden space, applies an activation, then projects back. SwiGLU uses two separate input projections: one supplies values, and the other scales them element by element after a SiLU activation.

For an input x of shape (T, D) and hidden width F = d_ff:

gate = x @ W_gate + b_gate       # (T, F)
up   = x @ W_up   + b_up         # (T, F)
h    = silu(gate) * up           # (T, F), elementwise product
out  = h @ W_down + b_down       # (T, D)

The three projections have separate parameters. Each row is processed independently, using the same weights. This block does not mix tokens or include attention, normalization or a residual addition.

A gate is not necessarily a mask

SiLU is silu(z) = z * sigmoid(z). Unlike sigmoid alone, it is not restricted to [0, 1]. For example:

silu(-3) ≈ -0.142278
silu(0)  = 0
silu(3)  ≈ 2.857722

A gate can therefore suppress, amplify or reverse the sign of an up feature. It is not a boolean mask or an attention probability. Apply SiLU only to gate, then multiply by the unactivated up values. Swapping the two paths changes the result for fixed parameters.

Work through a small example

These hand-set weights make the two paths easy to see. Biases are zero in this example; the exercise uses Linen's seeded initialization.

import jax
import jax.numpy as jnp

x = jnp.array([[1., -1.]])
W_gate = jnp.array([[1., 0., 2.], [4., 0., -1.]])
W_up = jnp.array([[2., 5., -1.], [0., 0., 0.]])
W_down = jnp.array([[1., 0.], [0., 1.], [1., -1.]])

gate = x @ W_gate                 # [[-3, 0, 3]]
up = x @ W_up                     # [[2, 5, -1]]
h = jax.nn.silu(gate) * up
out = h @ W_down

assert jnp.allclose(h, jnp.array([[-.2845552, 0., -2.8577223]]), atol=1e-6)
assert jnp.allclose(out, jnp.array([[-3.1422775, 2.8577223]]), atol=1e-6)

The hidden width is three, but the output width is two. The down projection both combines hidden features and restores the input width.

Count parameters and operations separately

With biases enabled, as in this exercise, SwiGLU has:

gate: D*F weights + F biases
up:   D*F weights + F biases
down: F*D weights + D biases
total: 3*D*F + 2*F + D

A two-layer Dense(F) -> ReLU -> Dense(D) block has 2*D*F + F + D parameters. For D=4, F=8, these totals are 116 and 76 respectively. Ignoring biases, reducing SwiGLU's hidden width to two-thirds of the conventional block's width matches their weight counts. Integer rounding and biases can prevent an exact total-parameter match.

SwiGLU has three dense matrix products; the conventional block has two. Activations and the gate product are elementwise operations, not extra matrix products. Matching parameter counts does not establish equal runtime or better model quality. Those require measurements for the model and training setup being compared.

Problem

Implement swiglu_ffn_forward(seed, x, d_ff):

  1. Define a Linen module with a d_ff field and a compact call method. Let D = x.shape[-1] and F = int(d_ff).
  2. Apply nn.Dense(F, name="gate") and nn.Dense(F, name="up") to x. Keep the default initializers and enabled biases.
  3. Compute h = jax.nn.silu(gate) * up and apply nn.Dense(D, name="down") to h.
  4. Initialize the module directly with jax.random.PRNGKey(int(seed)) and x, apply the resulting variables to x, and return out.reshape(-1).

x is a float32 array of shape (T, D), with positive dimensions. d_ff is positive. Return a one-dimensional array of length T * D, with features varying fastest. Keep seed and d_ff static if JIT-compiling the whole function.

Preserve the three explicit layer names: they affect seeded initialization. A fused gate/up projection can represent the same function if its weights and biases are copied appropriately, but initializing a new fused layer does not reproduce these separately initialized parameters.

The plot below shows SiLU itself and its derivative, not the complete SwiGLU network.

Loading visualization…