🤖 machine learning

Sign-Aware Gated Sparse Autoencoders: Modeling Anticorrelated Features with Bi-Jump-ReLU Activations

This paper introduces the Sign-Aware Gated Sparse Autoencoder (SA-GSAE), which employs a novel Bi-Jump-ReLU activation and two-sided gating to efficiently model anticorrelated features within shared latents, thereby significantly reducing dead neurons and improving reconstruction performance compared to standard non-negative SAEs across various LLM architectures.

Original authors: Bartosz Wieciech, Zmnako Awrahman, Marcin Czelej, Victor Hugo Jaramillo Velasquez, Wioletta Stobieniecka

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Bartosz Wieciech, Zmnako Awrahman, Marcin Czelej, Victor Hugo Jaramillo Velasquez, Wioletta Stobieniecka

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to understand human language. To do this, the computer builds a massive "dictionary" of concepts. In the old way of doing this (using standard tools called Sparse Autoencoders), the computer had a weird rule: it could only use positive numbers.

The Problem: The "One-Way Street" Dictionary

Because the computer could only use positive numbers, it couldn't represent a concept and its opposite with a single entry.

The Old Way: If the computer wanted to learn about "Pressure Too High," it had to create a specific dictionary entry for it. If it also wanted to learn about "Pressure Too Low," it had to create a completely separate entry for that, even though they are just two sides of the same coin.
The Result: The dictionary became bloated and inefficient. It was like having two separate keys for a lock—one for "Open" and one for "Close"—when you really just needed one key that could turn both ways. This wasted space and forced the computer to learn two separate things for what is actually one relationship.

The Solution: The "Two-Way Street" Dictionary

The authors of this paper invented a new tool called SA-GSAE (Sign-Aware Gated Sparse Autoencoder). Think of this as upgrading the dictionary to allow two-way streets.

Instead of forcing the computer to pick a "High" key or a "Low" key, this new tool lets a single dictionary entry represent both directions at once.

If the value is positive, it means "Pressure Too High."
If the value is negative, it means "Pressure Too Low."
If the value is zero, it means "Pressure is normal."

They achieved this using a special activation function they named Bi-Jump-ReLU. You can think of this as a smart gatekeeper that sits at the dictionary entry. It has a "dead zone" in the middle (where nothing happens), but if the signal is strong enough to the left, it opens the gate for "Negative," and if it's strong enough to the right, it opens the gate for "Positive."

Why This Matters: The "Half-Size" Advantage

Because this new tool can pack two opposite ideas into one slot, the authors found they could build a dictionary that is half the size but works just as well (or even better) than a full-sized, old-fashioned dictionary.

The Analogy: Imagine you have a suitcase full of clothes. The old method required you to pack a "Left Shoe" bag and a "Right Shoe" bag separately. The new method lets you pack them into a single "Shoe" bag that holds both. You can now fit the exact same amount of clothing into a suitcase that is half the size.
The Proof: The authors tested this on real AI models (Pythia-1B and SmolLM3-3B). They found that their "half-size" dictionary:
1. Reduced "Dead" Space: It stopped wasting slots on unused entries (reducing "dead fraction" by huge factors, sometimes 100x to 500x).
2. Kept Quality: It reconstructed the language just as accurately as the big, bloated dictionaries.
3. Avoided Crashes: Interestingly, when they tried to use the full-size version of their new tool, it sometimes crashed or performed poorly on specific parts of the AI. The "half-size" version avoided these crashes entirely, suggesting that less is sometimes more.

The "Safety Net" (Auxiliary Loss)

The authors also discovered a crucial training trick. To teach the gatekeeper (the Bi-Jump-ReLU) when to open the gate for "Positive" or "Negative," they used a "safety net" called auxiliary supervision.

Without the safety net: The gatekeeper got confused, stopped opening, and the whole system failed (98% of the dictionary became useless).
With the safety net: The gatekeeper learned exactly when to switch directions, making the system robust and efficient.

Summary

In short, this paper introduces a smarter way to teach AI models about opposites. By allowing a single dictionary entry to hold both positive and negative meanings, the authors created a system that is smaller, cleaner, and more efficient, using half the space to do the same job. They proved that you don't need a massive dictionary to understand complex ideas; you just need a dictionary that knows how to turn both ways.

Technical Summary: Sign-Aware Gated Sparse Autoencoders

Problem Statement

Sparse Autoencoders (SAEs) are a standard tool for extracting interpretable features from Large Language Models (LLMs) by learning overcomplete dictionaries of latent features. However, standard SAE variants (e.g., ReLU-based or standard Gated SAEs) enforce non-negativity on latent activations ( $f_i(x) \geq 0$ ). This constraint forces the model to learn separate latent features for diametrically opposed concepts (e.g., "pressure too high" vs. "pressure too low") along the same semantic axis.

This "sign splitting" wastes dictionary capacity, as two latents are required to represent a single bidirectional feature direction. While classical sparse coding allows for signed coefficients, standard SAE implementations typically lack mechanisms to represent both positive and negative evidence along a shared decoder direction without incurring L1 shrinkage artifacts or requiring doubled model width.

Methodology: Sign-Aware Gated SAE (SA-GSAE)

The authors propose the Sign-Aware Gated SAE (SA-GSAE), an architecture designed to model anticorrelated features using a single latent encoding. The method generalizes the Gated SAE by introducing a two-sided gated sparsity mechanism with signed magnitude and auxiliary supervision.

1. Bi-Jump-ReLU Activation

The core innovation is the Bi-Jump-ReLU activation function. For each latent $i$ , the model computes a decoder-aligned projection $t_i(x)$ . A learnable "dead zone" $[-\delta^-_i, \delta^+_i]$ is defined via thresholds.

Polarity Selection: A gate pre-activation $\pi_i(x)$ determines the polarity $s_i(x) \in \{+1, -1, 0\}$ . If $\pi_i(x)$ exceeds the positive threshold, the latent fires positively; if it falls below the negative threshold, it fires negatively.
Signed Magnitude: Once the gate selects a polarity, a separate magnitude path computes the activation value. Crucially, this path uses a gain parameter ( $g^+_i$ or $g^-_i$ ) and a shared magnitude offset, avoiding the subtraction of thresholds from the active magnitude (which causes L1 shrinkage).
Shared Direction: Both positive and negative activations map to the same decoder column $D_{:,i}$ , allowing a single latent to represent a bidirectional semantic axis.

2. Training Objective

The training loss combines three components:

Main Reconstruction: $\|x - D a(x) - b_{dec}\|^2_2$ , where $a(x)$ contains the signed Bi-Jump-ReLU activations.
Gate Sparsity Penalty: A two-sided hinge loss $\Omega_{gate}(\pi_i(x))$ applied to the gate pre-activations to encourage sparsity without shrinking the active magnitudes.
Auxiliary Reconstruction: An auxiliary loss $L_{aux}$ using a stop-gradient decoder path. This loss supervises the gate directly, teaching it which side of the dead zone to cross, while the main reconstruction learns the signed magnitude. This separation prevents the gate from collapsing (a common failure mode in signed SAEs).

3. Parameter Efficiency Argument

The authors provide a theoretical argument for parameter efficiency. If a fraction $p$ of ground-truth features appear as anticorrelated pairs, a standard SAE requires width $H$ , while the Sign-Aware model requires $H_{\pm} = H(1 - p/2)$ .

The Sign-Aware model adds 4 scalar parameters per latent (two thresholds, two magnitude scales) compared to a standard Gated SAE.
Net parameter savings occur when $p > \frac{8}{d_{in} + 4}$ .
For typical transformer hidden sizes ( $d_{in} = 2048$ ), this threshold is approximately $0.39\%$ . The authors argue that even a small fraction of anticorrelated pairs makes the Sign-Aware approach strictly more parameter-efficient.

Key Contributions

Architecture: Introduction of the SA-GSAE with Bi-Jump-ReLU, enabling signed magnitude estimation and two-sided gated sparsity. Ablation studies confirm that the two-sided gate and auxiliary loss are the "load-bearing" components; removing the auxiliary loss causes catastrophic failure (98% dead latents).
Symmetric-Magnitude Default: The authors demonstrate that tying the positive and negative magnitude scales ( $r^+_i = r^-_i$ ) is practically indistinguishable from independent scaling ( $|\Delta R^2| \approx 0.0015$ ) on real LLM activations. They recommend this symmetric variant as the default for parameter efficiency.
Empirical Validation: Extensive evaluation on Pythia-1B and SmolLM3-3B across three hookpoints (MLP output, attention output, residual stream).

Results

The paper compares a half-width SA-GSAE ( $H = 16,384$ ) against a full-width Gated SAE ( $2H = 32,768$ ) and a matched-width Gated SAE.

Dead Fraction Reduction: The half-width SA-GSAE drastically reduces the fraction of dead (inactive) latents compared to the full-width Gated SAE.
- On MLP-output and Pythia-1B residual stream cells, the reduction is by factors of 100× to 500× (geometric mean).
- On attention cells and SmolLM3-3B residual stream, the reduction is 2× to 4×.
Reconstruction Fidelity:
- On 3 of 6 cells (both MLP-output hookpoints and Pythia-1B residual stream), the half-width SA-GSAE strictly Pareto-dominates the full-width Gated SAE across the entire swept $L_0$ (sparsity) overlap, achieving equal or better $R^2$ with significantly fewer dead latents.
- On the remaining 3 cells, the $R^2$ gap is negligible (within $0.025$), while the dead fraction is still significantly reduced.
Stability: At full width, the SA-GSAE exhibits a reproducible reconstruction collapse on the SmolLM3-3B residual stream site (MSE up to 6.0, $R^2$ to -4.3). The half-width variant entirely avoids this pathology, providing a practical motivation for operating at half width.
Bipolar Usage: Calibration analysis shows that on MLP-output cells, most latents carry signal on both polarities (high "bipolar fraction"), driving the capacity wins. On attention layers, bipolar structure is concentrated in a smaller set of top-activation latents.

Significance and Claims

The paper claims to be the first two-sided, gated, no-shrinkage SAE for signed feature sharing. Its primary significance lies in:

Efficiency: Demonstrating that sign-awareness allows for a reduction in dictionary width (half-width) without sacrificing reconstruction quality, effectively resolving the "sign splitting" inefficiency.
Structural Insight: Showing that the ability to share decoder directions for anticorrelated pairs is not just a theoretical possibility but a practical necessity for efficient feature learning in LLMs, particularly in MLP outputs.
Robustness: Identifying that the auxiliary loss is critical for training stability and that a symmetric-magnitude variant is sufficient for high performance.

The authors treat SAE dictionaries as useful decompositions rather than unique ground-truth bases. They conclude that the half-width SA-GSAE with symmetric magnitudes is the recommended operating point, offering superior capacity utilization and stability compared to standard non-negative Gated SAEs.