Improving Dictionary Learning with Gated Sparse Autoencoders

“Further, shrinkage (Section 4.2) is built into the Lasso (Tibshirani, 1996) and well-studied in statistical learning (Hastie et al., 2015).”

“Motivated by the superposition hypothesis, Bricken et al. (2023) and Cunningham et al. (2023) recently used sparse autoencoders (SAEs; Ng (2011)) to find sparse decompositions of model activations in terms of an overcomplete basis, or dictionary (Mallat and Zhang, 1993).”

“This algorithm is a variant of orthogonal matching pursuit (Pati et al., 1993) which solves the orgothonalisation of the residual to the span of chosen dictionary elements approximately at every step rather than exactly, but which only requires matrix multiplies rather than matrix solves and is…”

“The discontinuity in particular introduces issues with methods like integrated gradients (Sundararajan et al., 2017) that discretely approximate a path integral, as applied to SAEs by Marks et al. (2024).”

“For example, sparse coding (Elad, 2010) studies how discrete and continuous representations can involve more representations than basis vectors, like our setup in Section 1, and sparse representations are also studied in neuroscience (Olshausen and Field, 1997; Thorpe, 1989).”

“Methods explicitly motivated by learning a disentangled representation typically enforce a prior structure on the learned representation, typically that features are aligned with the basis of a latent space (Chen et al., 2018, 2016; Kim and Mnih, 2018; Mathieu et al., 2019).”

“Taking inspiration from Gated Linear Units (Dauphin et al., 2017; Shazeer, 2020), we define the gated encoder as follows: where 𝟙 [• > 0 ] is the (pointwise) Heaviside step function and ⊙ denotes elementwise multiplication.”

“Further, shrinkage (Section 4.2) is built into the Lasso (Tibshirani, 1996) and well-studied in statistical learning (Hastie et al., 2015).”

“For example, sparse coding (Elad, 2010) studies how discrete and continuous representations can involve more representations than basis vectors, like our setup in Section 1, and sparse representations are also studied in neuroscience (Olshausen and Field, 1997; Thorpe, 1989).”

“Much work makes use of the fact that many concept representations appear to be linear (Elhage et al., 2021; Gurnee et al., 2023; Olah et al., 2020; Park et al., 2023).”

“Methods explicitly motivated by learning a disentangled representation typically enforce a prior structure on the learned representation, typically that features are aligned with the basis of a latent space (Chen et al., 2018, 2016; Kim and Mnih, 2018; Mathieu et al., 2019).”

“Disentanglement (Bengio, 2013) aims to learn representations that separate out distinct, independent ‘factors of variation’ of the underlying data generating process.”

“For all results in this section, we use gradient pursuit, as described in Blumensath and Davies (2008), as our inference time optimisation (ITO) algorithm.”

“Taking inspiration from Gated Linear Units (Dauphin et al., 2017; Shazeer, 2020), we define the gated encoder as follows: where 𝟙 [• > 0 ] is the (pointwise) Heaviside step function and ⊙ denotes elementwise multiplication.”

“Indeed, a common problem when studying SAE features/other interpretable directions is that, while the tails are monosemantic, activations near zero are very noisy (see e.g. the Arabic feature in Bricken et al. (2023) or the sentiment feature in Tigges et al. (2023)).”

“8B (Biderman et al., 2023) and Gemma-7B (Gemma Team et al., 2024), and on multiple sites within models: MLP layer outputs, attention layer outputs, and residual stream activations.”

“Methods explicitly motivated by learning a disentangled representation typically enforce a prior structure on the learned representation, typically that features are aligned with the basis of a latent space (Chen et al., 2018, 2016; Kim and Mnih, 2018; Mathieu et al., 2019).”

“Motivated by the superposition hypothesis, Bricken et al. (2023) and Cunningham et al. (2023) recently used sparse autoencoders (SAEs; Ng (2011)) to find sparse decompositions of model activations in terms of an overcomplete basis, or dictionary (Mallat and Zhang, 1993).”

“Marks et al. (2024) stress-test how useful SAEs are, and find success but also rely on methods that leave many error nodes in their computational subgraphs, which represent the difference between SAE reconstructions and the ground truth.”

“…components in small and large LMs (Olsson et al., 2022), identified computational subgraphs that carry out specific tasks in small LMs (circuits; Wang et al. (2023)) and reverse-engineered how toy 11 Concretely, we do this by training baseline SAEs, freezing their weights, and then learning…”

“The superposition hypothesis (Elhage et al., 2022b, Definitions and Motivation) posits a mechanistic explanation for these observations: in an intermediate representation of dimension 𝑛 , a model will encode 𝑀 ≫ 𝑛 concepts as linear directions, where the set of concepts and their directions is…”

“Classic approaches, like interpreting neurons (i.e. directions in the standard basis) are insufficient, as many are polysemantic and tend to activate for a range of different seemingly unrelated concepts (Bolukbasi et al., 2021; Elhage et al., 2022a,b, Empirical Phenomena).”

“In parallel to our work, Taggart (2024) finds early improvements with a similar Jump ReLU (Erichson et al., 2019) architecture change to SAEs, but with a different loss function, and without addressing the problems of L1.”

“For example, sparse coding (Elad, 2010) studies how discrete and continuous representations can involve more representations than basis vectors, like our setup in Section 1, and sparse representations are also studied in neuroscience (Olshausen and Field, 1997; Thorpe, 1989).”

“…as linear directions, where the set of concepts and their directions is fixed across all inputs, but on a given input only a sparse subset of concepts are active, ensuring that there is not much simultaneous interference (Gurnee et al., 2023, Appendix A) between these (non-orthogonal) concepts.”

“Much work makes use of the fact that many concept representations appear to be linear (Elhage et al., 2021; Gurnee et al., 2023; Olah et al., 2020; Park et al., 2023).”

“…ANN InceptionV1 (Cammarata et al., 2020; Olah et al., 2018), Chris Olah and colleagues explicitly endorse applying life science research strategies to ANNs, suggesting that “neural networks are an object of empirical investigation, perhaps similar to an organism in biology” (Olah et al., 2020).”