GLU Activation Functions

Gated linear unit (GLU) feed-forward blocks.

SwiGLU, GEGLU, and ReGLU are gated FFN units that replace the standard Linear -> activation -> Linear block with two projections whose elementwise product forms the gate:

FFN(x) = (act(x W_gate + b_gate)) * (x W_up + b_up),  then down-project.

They are not pure elementwise activations (they own linear weights), so they are implemented here as drop-in FFN modules consumed by HybridLayer.

Reference: Shazeer (2020), GLU Variants Improve Transformer, arXiv:2002.05202.

class src.model.activation_function.glu.GatedFFN(*args: Any, **kwargs: Any)[source]

Bases: Module

Gated feed-forward network (SwiGLU / GEGLU / ReGLU).

Computes:

y = dropout( down( act(gate_proj(x)) * up_proj(x) ) )

where gate_proj and up_proj map hidden -> intermediate and down_proj maps intermediate -> hidden.

Parameters:

  • hidden_size – Input/output feature dimension.

  • intermediate_size – Gated intermediate dimension.

  • gate_fn – Elementwise gating function applied to the gate projection. One of F.silu (SwiGLU), F.gelu (GEGLU), or F.relu (ReGLU). Default: F.silu (SwiGLU).

  • bias – Whether to include bias in the projections. Default: False.

  • dropout – Dropout probability after gating. Default: 0.0.

  • proj_factory – Optional callable (in, out, bias) -> nn.Module used to build the projections (e.g. BitLinear for BitNet compatibility). When None, plain nn.Linear is used.

gate_proj

Linear hidden -> intermediate feeding the gate function.

up_proj

Linear hidden -> intermediate value pathway.

down_proj

Linear intermediate -> hidden output projection.

Reference: Shazeer (2020), arXiv:2002.05202.

forward(x: torch.Tensor) → torch.Tensor[source]

Apply the gated FFN.

Parameters:

x – Input tensor of shape (..., hidden_size).

Returns:

Tensor of shape (..., hidden_size).

src.model.activation_function.glu.make_gated_ffn(kind: str, hidden_size: int, intermediate_size: int, bias: bool = False, dropout: float = 0.0, proj_factory: Callable[[int, int, bool], torch.nn.Module] | None = None) → GatedFFN[source]

Build a GatedFFN for the requested GLU variant.

Parameters:

  • kind – One of "swiglu", "geglu", "reglu".

  • hidden_size – Input/output feature dimension.

  • intermediate_size – Gated intermediate dimension.

  • bias – Whether projections include bias. Default: False.

  • dropout – Dropout probability after gating. Default: 0.0.

  • proj_factory – Optional projection factory (e.g. BitLinear).

Returns:

A GatedFFN configured for the requested variant.

Raises:

ValueError – If kind is not a recognized GLU variant.