Activation Functions¶

Activation functions are the non-linear transformations applied element-wise (or gate-wise) inside neural networks. Without them, any composition of linear layers collapses to a single linear map, and the network can represent only affine functions. This page provides a rigorous treatment of the activations used in modern transformer models, together with their ZigLlama implementations.

1. Why Non-Linearity?¶

1.1 The Linear Collapse Problem¶

This means a 100-layer "deep" network of purely linear transformations has exactly the same representational capacity as a single matrix multiply. Stacking layers buys us nothing without non-linearity.

1.2 Universal Approximation¶

The theorem guarantees existence of a good approximation but says nothing about how efficiently the network can be trained. In practice, deeper networks with well-chosen activations learn more efficiently than wide shallow ones.

2. Classical Activations¶

2.1 Rectified Linear Unit (ReLU)¶

Derivative:

The derivative is the Heaviside step function (undefined at exactly zero, where a sub-gradient of 0 is conventionally used).

Properties:

2.2 Sigmoid¶

Derivative: \(\sigma'(x) = \sigma(x)(1 - \sigma(x))\).

2.3 Hyperbolic Tangent¶

Equivalently, \(\tanh(x) = 2\sigma(2x) - 1\).

Tanh was the default activation in early neural networks. Its zero-centered output is advantageous for gradient dynamics, but it still suffers from vanishing gradients in the saturated regime.

3. Modern Activations¶

3.1 GELU -- Gaussian Error Linear Unit¶

Intuition. GELU weights the input \(x\) by the probability that a standard Gaussian random variable is less than \(x\). Large positive inputs pass through nearly unchanged (\(\Phi(x) \approx 1\)); large negative inputs are suppressed (\(\Phi(x) \approx 0\)); and intermediate values are smoothly interpolated.

Fast approximation (tanh-based):

This approximation avoids computing the error function and is used in most production implementations, including ZigLlama.

3.2 SiLU / Swish¶

SiLU is self-gated: the input acts as both the signal and the gate. It shares many properties with GELU -- smooth, non-monotonic near zero, and asymptotically linear for large positive inputs -- but uses the simpler sigmoid rather than the Gaussian CDF.

3.3 SwiGLU -- SiLU-Gated Linear Unit¶

SwiGLU is a gated linear unit variant that uses SiLU as its gating activation. The key idea is to split the computation into two parallel streams -- one processed through SiLU (the gate), the other left linear (the content) -- and then combine them via element-wise multiplication.

This architecture requires three weight matrices per feed-forward layer (gate, up, down) rather than two, but empirically yields better model quality at the same compute budget.

3.4 GeGLU -- GELU-Gated Linear Unit¶

GeGLU replaces the SiLU gate in SwiGLU with GELU. The two variants perform comparably in practice; the choice often follows from consistency with the rest of the architecture (e.g., if the model already uses GELU elsewhere).

4. Comparison Table¶

5. Activation Function Landscape¶

graph LR
    subgraph "Classical"
        ReLU
        Sigmoid
        Tanh
    end

    subgraph "Modern Smooth"
        GELU
        SiLU["SiLU / Swish"]
    end

    subgraph "Gated Variants"
        GLU["GLU (sigmoid gate)"]
        SwiGLU["SwiGLU (SiLU gate)"]
        GeGLU["GeGLU (GELU gate)"]
    end

    ReLU --> GELU
    Sigmoid --> SiLU
    Sigmoid --> GLU
    SiLU --> SwiGLU
    GELU --> GeGLU
    GLU --> SwiGLU
    GLU --> GeGLU

6. Implementation in ZigLlama¶

6.1 Type Enum and Dispatcher¶

ZigLlama defines all supported activations in a single enum and provides a compile-time dispatcher:

/// Activation function types supported in transformers
pub const ActivationType = enum {
    ReLU,
    GELU,
    SiLU,    // Also known as Swish
    GLU,     // Gated Linear Unit
    GeGLU,   // GELU-based Gated Linear Unit
    SwiGLU,  // SiLU-based Gated Linear Unit
    Tanh,
    Sigmoid,
};

/// Generic activation function dispatcher
pub fn applyActivation(
    comptime T: type,
    activation_type: ActivationType,
    input: Tensor(T),
    allocator: Allocator,
) TensorError!Tensor(T) {
    return switch (activation_type) {
        .ReLU   => relu(T, input),
        .GELU   => gelu(T, input, allocator),
        .SiLU   => silu(T, input, allocator),
        .GLU    => glu(T, input, allocator),
        .GeGLU  => geglu(T, input, allocator),
        .SwiGLU => swiglu(T, input, allocator),
        .Tanh   => tanh_activation(T, input, allocator),
        .Sigmoid => sigmoid(T, input, allocator),
    };
}

6.2 Scalar Activation Functions¶

Each scalar activation operates element-wise over a Tensor(T).

/// ReLU: max(0, x)
pub fn relu(comptime T: type, input: Tensor(T)) TensorError!Tensor(T) {
    var result = try Tensor(T).init(input.allocator, input.shape);
    for (0..input.size) |i| {
        result.data[i] = @max(0.0, input.data[i]);
    }
    return result;
}

/// GELU: tanh-based approximation
pub fn gelu(comptime T: type, input: Tensor(T), allocator: Allocator) TensorError!Tensor(T) {
    var result = try Tensor(T).init(allocator, input.shape);
    for (0..input.size) |i| {
        const x = input.data[i];
        const inner = 0.7978845608 * (x + 0.044715 * x * x * x);
        result.data[i] = 0.5 * x * (1.0 + std.math.tanh(inner));
    }
    return result;
}

/// SiLU / Swish: x * sigmoid(x)
pub fn silu(comptime T: type, input: Tensor(T), allocator: Allocator) TensorError!Tensor(T) {
    var result = try Tensor(T).init(allocator, input.shape);
    for (0..input.size) |i| {
        const x = input.data[i];
        result.data[i] = x * (1.0 / (1.0 + @exp(-x)));
    }
    return result;
}

6.3 Gated Activations¶

Gated variants split the input tensor along the last dimension. The first half is the content stream; the second half passes through the gating activation.

/// SwiGLU: content * SiLU(gate)
pub fn swiglu(comptime T: type, input: Tensor(T), allocator: Allocator) TensorError!Tensor(T) {
    const last_dim = input.shape[input.shape.len - 1];
    const half_size = input.size / 2;
    // ... shape validation omitted for brevity ...
    var result = try Tensor(T).init(allocator, output_shape);
    for (0..half_size) |i| {
        const a = input.data[i];                          // content
        const b = input.data[i + half_size];              // gate
        const sigmoid_b = 1.0 / (1.0 + @exp(-b));
        result.data[i] = a * (b * sigmoid_b);             // a * SiLU(b)
    }
    return result;
}

7. Numerical Considerations¶

7.1 GELU Approximation Accuracy¶

The tanh-based GELU approximation has a maximum absolute error of approximately \(4 \times 10^{-4}\) relative to the exact \(x\Phi(x)\) form. For inference in f32, this is well within acceptable precision.

7.2 Sigmoid Overflow¶

For very large negative inputs, \(e^{-x}\) overflows f32. ZigLlama relies on IEEE 754 semantics: \(\exp(-x) = +\infty\) yields \(1/(1+\infty) = 0\), which is the correct saturated value. No explicit clamping is required.

7.3 SiLU Near Zero¶

\(\operatorname{SiLU}(0) = 0\) exactly, and the derivative at zero is 0.5. Unlike ReLU, the gradient is non-zero at the origin, so neurons do not die.

8. Exercises¶

1. Prove that \(\operatorname{SiLU}'(x) = \sigma(x)(1 + x(1 - \sigma(x)))\).
2. Implement a "leaky" GELU variant where negative outputs are scaled by a small constant \(\alpha\) instead of suppressed.
3. Benchmark the per-element cost of GELU vs. SiLU on a 4096-element tensor using std.time.Timer. Which is faster and why?
4. Explain why SwiGLU requires three weight matrices while standard FFN requires only two, and show that the parameter counts match when \(d_{ff}^{\text{SwiGLU}} = \frac{2}{3} d_{ff}^{\text{standard}}\).

References¶