NNX SwiGLU FFN

Multiply two learned branches

A feed-forward network transforms each token independently. SwiGLU uses two projections of the same input, multiplies their outputs after applying SiLU to one branch, then projects back to the model width:

output = down(silu(gate(x)) * up(x))
silu(z) = z * sigmoid(z)

Let D be the model width and F the hidden width. Gate and up each map D features to F features. Their elementwise product also has F features; down maps those back to D. The same weights are used at every token position, without mixing information between positions.

What the gate does

The SiLU branch supplies one multiplier per hidden feature and token. A multiplier near zero suppresses that feature's contribution to the product. But this is not a probability or an on/off switch: SiLU can be negative or greater than one. It can reverse a feature's sign or amplify it as well as shrink it.

Gate and up have separate learned parameters. Reusing the same projection would still produce a multiplicative expression, but it would tie the two branches and define a different model. Likewise, applying SiLU to up instead of gate changes the result for the seeded parameters here.

Mathematically, SiLU is smooth and nonzero at finite negative inputs. ReLU instead returns zero for negative inputs. Neither observation alone tells you which activation gives better results on a particular training task.

A complete example

Here D=1 and F=2. The gate projection returns [0, log(3)], while up returns [x, 2*x]. The down projection adds the two products and a bias.

import jax
import jax.numpy as jnp
from flax import nnx

class SwiGLU(nnx.Module):
    def __init__(self, d_model, d_ff, rngs):
        self.gate = nnx.Linear(d_model, d_ff, rngs=rngs)
        self.up = nnx.Linear(d_model, d_ff, rngs=rngs)
        self.down = nnx.Linear(d_ff, d_model, rngs=rngs)

    def __call__(self, x):
        return self.down(jax.nn.silu(self.gate(x)) * self.up(x))

model = SwiGLU(1, 2, nnx.Rngs(0))
model.gate.kernel[...] = jnp.zeros((1, 2))
model.gate.bias[...] = jnp.array([0.0, jnp.log(3.0)])
model.up.kernel[...] = jnp.array([[1.0, 2.0]])
model.up.bias[...] = jnp.zeros(2)
model.down.kernel[...] = jnp.ones((2, 1))
model.down.bias[...] = jnp.array([1.0])
x = jnp.array([[1.0], [2.0]])
expected = 1.0 + 1.5 * jnp.log(3.0) * x
assert jnp.allclose(model(x), expected)

Since sigmoid(log(3)) is 3/4, SiLU(log(3)) is 3*log(3)/4. The first gate is zero, so only the second hidden feature contributes. If both gate outputs were zero, the final output would still include the down projection's bias; it would not necessarily be zero.

Compare parameter counts at the same width

A two-layer D-to-F-to-D network has 2*D*F matrix entries. SwiGLU has 3*D*F, or 50% more at the same hidden width F. With the default biases used here, the exact totals are:

• Two-layer network: 2*D*F + F + D.
• SwiGLU: 3*D*F + 2*F + D.

Matching matrix counts requires a SwiGLU hidden width two-thirds of the original hidden width. For example, an original width of 4*D corresponds to approximately (8/3)*D for SwiGLU, rounded to a suitable integer. This comparison ignores biases and activation costs. Runtime also depends on array sizes, kernels and hardware; the parameter ratio is not a measured speed ratio. Use the supplied d_ff in this exercise.

The chart shows scalar SiLU compared with ReLU. It does not show the other projection, the elementwise product or the complete network output.