SwiGLU activation

Background

SwiGLU is the gated activation in the feed-forward block of modern LLMs (PaLM, LLaMA). A GLU (gated linear unit) splits its input in half and lets one half gate the other; SwiGLU uses the Swish gate $\text{swish}(z) = z\,\sigma(z)$. It consistently outperforms plain ReLU/GELU feed-forward layers at the same parameter budget.

Problem statement

Implement SwiGLU(x) for input of shape (batch, 2d). Split the last dimension into halves $x_1, x_2$ and return:

$$\text{SwiGLU}(x) = x_1 \odot \text{swish}(x_2) = x_1 \odot \big(x_2 \cdot \sigma(x_2)\big)$$

where $\sigma$ is the sigmoid.

Input

• x — np.ndarray of shape (batch_size, 2d).

Output

Returns an np.ndarray of shape (batch_size, d).

Examples

Example 1

Input:  [[1, -1, 1000, -1000]]
Output: [[1000.0, 0.0]]

Explanation: split into $x_1=[1,-1]$ and $x_2=[1000,-1000]$. $\sigma(1000)\approx1$ so $\text{swish}(1000)\approx1000$; $\sigma(-1000)\approx0$ so $\text{swish}(-1000)\approx0$. Then $x_1\odot\text{swish}(x_2)=[1\cdot1000,\,-1\cdot0]=[1000, 0]$.

Constraints

• The last dimension is even; split it into equal halves $x_1$ (first) and $x_2$ (second).
• Gate with $\text{swish}(x_2) = x_2\,\sigma(x_2)$; output $= x_1 \odot \text{swish}(x_2)$.
• The output has half the last-dimension size: (batch, d).

Notes

• The GLU family $x_1 \odot g(x_2)$ lets the network learn which features to pass through — Swish is the smooth, non-monotonic gate that works best empirically.
• In a real FFN, $x_1$ and $x_2$ come from two separate linear projections of the input; here they are handed to you pre-projected as the two halves.