TLDR: The research paper introduces Signed Quadratic Shrink (SQS), a novel activation function designed for Gated Linear Units (GLUs). SQS enables neural networks to learn interpretable features directly from their weights, addressing the performance and data efficiency drawbacks of existing weight-based interpretability methods like Bilinear MLPs. Experimental results on image classification (MNIST, Fashion MNIST) and language modeling (Tiny Stories) demonstrate that SQS achieves performance competitive with, and often superior to, state-of-the-art activation functions (e.g., GELU, SwiGLU) while successfully preserving the interpretability of learned features.
Understanding how artificial intelligence models make decisions is a critical step towards building more reliable and trustworthy systems. While many techniques focus on analyzing how activations in a neural network contribute to its output, a more direct approach involves deriving meaningful features directly from the network’s weights. This method, known as weight-based interpretability, offers the promise of stronger guarantees and greater computational efficiency.
However, existing methods for achieving this often come with trade-offs, such as reduced performance or inefficiency in data usage. A new research paper introduces an innovative solution to this challenge: the Signed Quadratic Shrink (SQS) activation function.
Introducing Signed Quadratic Shrink (SQS)
The SQS activation function is specifically designed for Gated Linear Units (GLUs), a type of neural network architecture. Its primary goal is to enable GLUs to learn features that are inherently interpretable from their weights, without sacrificing performance or data efficiency. This is a significant advancement, as it aims to combine the best of both worlds: the clarity of interpretable models with the high performance of state-of-the-art activation functions like SwiGLU and GEGLU.
Traditional Bilinear Multi-Layer Perceptrons (MLPs) have been explored for their ability to yield interpretable features due to their simpler structure, which essentially omits a standard activation function. This allows their output to be analyzed through weight spectra, revealing underlying patterns. However, these Bilinear MLPs often fall short in terms of overall performance and how efficiently they learn from data compared to more advanced GLUs.
SQS addresses these shortcomings by modifying the activation process. Standard quadratic activation functions can suffer from issues like vanishing gradients (where changes become too small to learn effectively) or exploding gradients (where changes become too large and unstable). SQS introduces a clever modification that includes a ‘shift’ and a ‘shrinking factor’ to the quadratic function. This makes the function behave in a quasi-linear way for very small or very large inputs, while maintaining a signed quadratic behavior for intermediate values. Crucially, this design helps to preserve the valuable weight structure properties that make Bilinear MLPs interpretable, while enhancing their performance.
Experimental Validation: Interpretability and Performance
The researchers conducted extensive experiments to evaluate SQS against other popular activation functions like GeLU and ReLU, as well as traditional Bilinear GLUs. They tested these models on various datasets, including MNIST (handwritten digits), Fashion MNIST (clothing items), and Tiny Stories (language modeling).
For interpretability, the study focused on how SQS-GLUs learn meaningful features. By analyzing the eigenvectors derived from the SQS-GLU weights, they found that these eigenvectors visually corresponded to distinct classes in datasets like MNIST and Fashion MNIST. This means that the model was learning features that humans could easily understand and relate to the input data. Furthermore, the eigenvectors from SQS-GLUs showed a high degree of similarity to those from Bilinear MLPs, confirming that SQS successfully retains the interpretability benefits of its predecessors.
In terms of performance, SQS demonstrated impressive results. On image classification tasks (MNIST and Fashion MNIST), SQS-GLUs converged faster and achieved lower loss compared to ReLU-GLUs and Bilinear MLPs. Its convergence rate was comparable to, and in some cases, even surpassed, that of leading activation functions like SwiGLU and GELU, while often achieving lower final loss. SQS-GLUs were also quicker to reach high accuracy thresholds (80%, 85%, and 90%) and maintained competitive final accuracy.
For language modeling tasks using the Tiny Stories dataset, SQS-GLUs also showed superior performance, yielding the best results in terms of both loss and perplexity at various training stages.
Also Read:
- Dynamic Noise Injection for Robust Deep Learning
- Unveiling the Silent Thought Processes of Large Language Models
A Step Forward for Explainable AI
The introduction of the Signed Quadratic Shrink (SQS) activation function marks a significant step forward in the field of mechanistic interpretability. By enabling neural networks to learn interpretable features directly from their weights while achieving competitive or even superior performance, SQS offers a powerful tool for developing more transparent, reliable, and robust AI models. This research paves the way for a future where the inner workings of complex machine learning systems are not just efficient, but also understandable. You can read the full research paper here: Preserving Bilinear Weight Spectra with a Signed and Shrunk Quadratic Activation Function.
