Activating neurons with Gated Linear Units (GLU) and Friends

Activation functions are a crucial component of every neural network. They introduce non-linearity to the model, allowing it to learn complex patterns in the data.

where \(g\) is the activation function, \(W_1\) and \(W_2\) are the weight matrices, \(b_1\) and \(b_2\) are the bias vectors, and \(x\) is the input vector.

We would end up with a linear model, which is not very useful for learning complex mappings.

In this post, we will explore activation functions derived from the Gated Linear Unit (GLU) as introduced in Shazeer et al. (2020), and their effectiveness in learning a task when applied to the feed-forward module in the Transformer architecture.

Introduction

The Transformer model processes the input data by alternating between multi-head attention and “point-wise feed-forward networks” (FFN). In its basic form, the FFN is typically applied on each input token independently, and consists of two linear layers with a non-linear activation function in between.

Ramachandran et al. (2017) introduced by the Swish activation function:

where we assume \(\beta = 1\) for simplicity.

Above we can see the plots of the activation functions discussed in this section.

Gated Linear Units (GLU) and Variants

where the bias term is ommited for simplicity.

Notice that all of these layers have 3 weight matrices as opposed to the original 2 in the FFN module. Appropriately, the dimensionality of the projections are adjusted to match the total number of parameters in the original FFN module so that a fair comparison can be made.

Experiments

Using the above FFN proposals, Shazeer et al. (2020) conducted experiments in order to evaluate the effectiveness of the different activation functions.

A T5 model is trained on a denoising objective of predicting missing text segments, and subsequently fine-tuned on various language understanding tasks.

First, they evaluated the log-perplexity by training the model on the C4 dataset, and evaluating it on a held-out validation set. The GEGLU and SwiGLU activation functions outperformed the rest.

Next, they fine-tune each individual pre-trained model on a set of language understanding tasks from the GLUE and SuperGLUE benchmarks. While the results are noisy, the GLU variants perform best on all of the tasks.

Conclusion

In this post, we have explored several variations of the Gated Linear Unit (GLU) activation function, and their effectiveness in learning a task when applied to the feed-forward module in the Transformer architecture.

The results suggest that the GEGLU and SwiGLU activation functions outperform the rest in terms of log-perplexity on the C4 dataset, and fine-tuning on various language understanding tasks.

So what variation do State-of-the-Art models use today?

Google’s Gemma uses the GEGLU activation function:

class GemmaMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        quant: bool,
    ):
        super().__init__()
        self.gate_proj = Linear(hidden_size, intermediate_size, quant)
        self.up_proj = Linear(hidden_size, intermediate_size, quant)
        self.down_proj = Linear(intermediate_size, hidden_size, quant)

    def forward(self, x):
        gate = self.gate_proj(x)
        gate = F.gelu(gate, approximate="tanh")
        up = self.up_proj(x)
        fuse = gate * up
        outputs = self.down_proj(fuse)
        return outputs

Whereas, Meta’s LLaMA uses the SwiGLU activation function:

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

Which one is better? Noam puts it best:

“We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.”