FISTA-Refined Dictionary Learning with Learned TopK Initialization for Sparse Decomposition of Transformer Residual Streams

Sparse autoencoders (SAEs) decompose transformer activations into sparse feature dictionaries, but one-shot support selection can leave many features unused and become difficult to optimize at very small active budgets. We study an SAE-like dictionary-learning procedure that combines a learned TopK encoder with a few unrolled Fast Iterative Shrinkage-Thresholding Algorithm (FISTA) steps for coefficient refinement. On GPT-2 Small layer-11 residual streams, a 12,288-feature TopK256 model with five FISTA steps achieves ΔCE = 0.0745 and a 3.09% dead-feature ratio, compared with 0.0859 and 50.76% for a matched TopK SAE. Ablations show that initialization and learned TopK feature selection explain most of the reduction in dead features, while FISTA further improves fidelity. In a 32,768-feature regime, direct exact-TopK32 training is unstable. Training becomes feasible by retaining a TopK128 candidate support while reducing the final active budget from TopK128 to TopK64 and then TopK32. Five refinement steps reach ΔCE = 0.2442 and 17.84% dead features after 264 million token presentations. A residual-gain-ranked loss over candidates at ranks 33–48 improves near-boundary feature recruitment with a modest fidelity cost. These results show that lightweight iterative correction can improve fixed-budget sparse decomposition without replacing one-shot SAE inference with a full optimization solver.