gated_linear_unit

gated_linear_unit

Gated Linear Unit activations.

Ref: Shazeer (2020), https://arxiv.org/abs/2002.05202

GatedLinearUnit

GatedLinearUnit(input_dimension, hidden_dimension, bias=False, gate_activation=None)

Bases: Module

Gated Linear Unit with configurable activation.

GLU(x) = act(xW_gate) ⊗ (xW_value)

SwiGLU uses SiLU, GeGLU uses GELU. The gating doubles the parameter count compared to a simple linear + activation.

Parameters:

Name	Type	Description	Default
input_dimension	int	Input feature dimension.	required
hidden_dimension	int	Hidden dimension (output will be this size).	required
bias	bool	Whether to use bias in linear layers.	False
gate_activation	Module | None	Activation applied to the gate projection.	None

Source code in src/versatil/models/layers/gated_linear_unit.py

def __init__(
    self,
    input_dimension: int,
    hidden_dimension: int,
    bias: bool = False,
    gate_activation: nn.Module | None = None,
):
    super().__init__()
    self.gate_proj = nn.Linear(input_dimension, hidden_dimension, bias=bias)
    self.value_proj = nn.Linear(input_dimension, hidden_dimension, bias=bias)
    self.gate_activation = gate_activation or nn.SiLU()
    self._init_weights()

forward

forward(x)

Apply the gated projection: activation(gate(x)) * value(x).

Source code in src/versatil/models/layers/gated_linear_unit.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply the gated projection: activation(gate(x)) * value(x)."""
    gate = self.gate_activation(self.gate_proj(x))
    value = self.value_proj(x)
    return gate * value

SwiGLU

SwiGLU(input_dimension, hidden_dimension, bias=False)

Bases: GatedLinearUnit

Swish-Gated Linear Unit: GLU with SiLU (Swish) activation.

Source code in src/versatil/models/layers/gated_linear_unit.py

def __init__(self, input_dimension: int, hidden_dimension: int, bias: bool = False):
    super().__init__(
        input_dimension=input_dimension,
        hidden_dimension=hidden_dimension,
        bias=bias,
        gate_activation=nn.SiLU(),
    )

GeGLU

GeGLU(input_dimension, hidden_dimension, bias=False)

Bases: GatedLinearUnit

GELU-Gated Linear Unit: GLU with GELU(approximate='tanh') activation.

Source code in src/versatil/models/layers/gated_linear_unit.py

def __init__(self, input_dimension: int, hidden_dimension: int, bias: bool = False):
    super().__init__(
        input_dimension=input_dimension,
        hidden_dimension=hidden_dimension,
        bias=bias,
        gate_activation=nn.GELU(approximate="tanh"),
    )