Breaking the Block: Preserving Data Continuity to Train Superior SAEs for Instruct Models
Jiaming Li ⋅ Haoran Ye ⋅ Yukun Chen ⋅ Xinyue Li ⋅ Lei Zhang ⋅ Hamid Alinejad-Rokny ⋅ Jimmy Chih-Hsien Peng ⋅ Min Yang
Abstract
Sparse Autoencoders (SAEs) are a cornerstone of mechanistic interpretability. Existing training methods inherit the Block Training paradigm from LLM pre-training, which introduces destructive gradient noise in instruct models due to attention leakage from unrelated contexts. Using GSNR analysis, we theoretically characterize this issue and propose $\underline{\textbf{F}}$inetuning-$\underline{\textbf{a}}$ligned $\underline{\textbf{S}}$equential $\underline{\textbf{T}}$raining ($\textit{FAST}$), a sequential training paradigm specifically designed for instruct models. $\textit{FAST}$ aligns SAE training with the data distribution and activation patterns of instruct models, substantially improving both reconstruction fidelity and feature interpretability. Experimental results show that $\textit{FAST}$ achieves higher GSNR, a significantly lower log-scaled MSE of 0.6468 compared to the baseline’s 5.1985, and a near-zero Delta Loss (-0.51\% to 0.37\%). Moreover, on Llama-3.2-3B-it, $\textit{FAST}$ produces 21.1\% high-quality features, substantially outperforming baseline methods that achieve 7.0\% and 10.2\%. We further find that intervening on special token activations through SAEs can improve generation quality, revealing new opportunities for fine-grained control. Our codes are available as open source at https://github.com/Geaming2002/FAST.
Lay Summary
Understanding large language models requires tools like Sparse Autoencoders to translate internal neural activity into understandable concepts. However, standard training methods create destructive errors when applied to modern instruct models. To solve this, we introduce $\textit{FAST}$, a sequential training method that precisely matches how conversational models process data. Our experiments show $\textit{FAST}$ drastically reduces reconstruction errors and discovers up to three times more high-quality, interpretable features than previous approaches. Furthermore, we demonstrate that $\textit{FAST}$ can adjust specific internal model states to directly improve text generation.
Successful Page Load