Group-SAE: Efficient Training of Sparse Autoencoders for Large Language Models via Layer Groups

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Ghilardi, Davide, Belotti, Federico, Molinari, Marco, Ma, Tao, Palmonari, Matteo
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911163973894144
author Ghilardi, Davide
Belotti, Federico
Molinari, Marco
Ma, Tao
Palmonari, Matteo
author_facet Ghilardi, Davide
Belotti, Federico
Molinari, Marco
Ma, Tao
Palmonari, Matteo
contents SAEs have recently been employed as a promising unsupervised approach for understanding the representations of layers of Large Language Models (LLMs). However, with the growth in model size and complexity, training SAEs is computationally intensive, as typically one SAE is trained for each model layer. To address such limitation, we propose \textit{Group-SAE}, a novel strategy to train SAEs. Our method considers the similarity of the residual stream representations between contiguous layers to group similar layers and train a single SAE per group. To balance the trade-off between efficiency and performance, we further introduce \textit{AMAD} (Average Maximum Angular Distance), an empirical metric that guides the selection of an optimal number of groups based on representational similarity across layers. Experiments on models from the Pythia family show that our approach significantly accelerates training with minimal impact on reconstruction quality and comparable downstream task performance and interpretability over baseline SAEs trained layer by layer. This method provides an efficient and scalable strategy for training SAEs in modern LLMs.
format Preprint
id arxiv_https___arxiv_org_abs_2410_21508
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Group-SAE: Efficient Training of Sparse Autoencoders for Large Language Models via Layer Groups
Ghilardi, Davide
Belotti, Federico
Molinari, Marco
Ma, Tao
Palmonari, Matteo
Computation and Language
Artificial Intelligence
SAEs have recently been employed as a promising unsupervised approach for understanding the representations of layers of Large Language Models (LLMs). However, with the growth in model size and complexity, training SAEs is computationally intensive, as typically one SAE is trained for each model layer. To address such limitation, we propose \textit{Group-SAE}, a novel strategy to train SAEs. Our method considers the similarity of the residual stream representations between contiguous layers to group similar layers and train a single SAE per group. To balance the trade-off between efficiency and performance, we further introduce \textit{AMAD} (Average Maximum Angular Distance), an empirical metric that guides the selection of an optimal number of groups based on representational similarity across layers. Experiments on models from the Pythia family show that our approach significantly accelerates training with minimal impact on reconstruction quality and comparable downstream task performance and interpretability over baseline SAEs trained layer by layer. This method provides an efficient and scalable strategy for training SAEs in modern LLMs.
title Group-SAE: Efficient Training of Sparse Autoencoders for Large Language Models via Layer Groups
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2410.21508