SwiGLU Feed-Forward Block

Write the SwiGLU feed-forward block that replaced the ReLU MLP inside the transformer layers of models such as LLaMA and PaLM.

The classic transformer MLP is relu(x @ W1) @ W2. SwiGLU uses two input projections. One goes through the SiLU activation and acts as a gate, and the other is multiplied by that gate element by element:

silu(z) = z * sigmoid(z) = z / (1 + exp(-z))
hidden  = silu(x @ W_gate) * (x @ W_up)
out     = hidden @ W_down

The gate lets the block decide, per hidden unit and per token, how much of the up projection to pass on.

Function signature

def swiglu_ffn(x, W_gate, W_up, W_down):

Example

Input

x      = np.array([[1.0, -1.0]])
W_gate = np.eye(2)
W_up   = np.array([[2.0, 0.0], [0.0, 1.0]])
W_down = np.array([[1.0], [1.0]])
swiglu_ffn(x, W_gate, W_up, W_down)

Output

array([[1.73105858]])

Reasoning

The gate input is [1, -1], and silu gives [0.7311, -0.2689]. The up projection is [2, -1]. Their product is [1.4621, 0.2689], and the down projection sums them to 1.7311. Note that the negative gate still lets a little through. ReLU would have blocked it.

Constraints and notes

• x has shape (T, d), W_gate and W_up have shape (d, h), and W_down has shape (h, d_out).
• There are no bias terms.