Towards Interpretable and Inference-Optimal COT Reasoning with Sparse Autoencoder-Guided Generation
This work addresses interpretability and optimization in reasoning for AI researchers, though it appears incremental as it builds on existing methods like sparse autoencoders and clustering.
The authors tackled the problem of improving mathematical reasoning in large language models by balancing exploitation and exploration, using sparse autoencoders and clustering to guide generations, resulting in enhanced reasoning quality.
We propose a novel method that leverages sparse autoencoders (SAEs) and clustering techniques to analyze the internal token representations of large language models (LLMs) and guide generations in mathematical reasoning tasks. Our approach first trains an SAE to generate sparse vector representations for training tokens, then applies k-means clustering to construct a graph where vertices represent token clusters and weighted edges capture sequential token transitions. Using this graph, we define an edge-weight based reward function to quantify adherence to established reasoning traces, thereby identifying exploitative reasoning trajectories. Additionally, we measure generation diversity from clustering to assess the extent of exploration. Our findings indicate that balancing both exploitation and exploration is crucial for achieving high accuracy in mathematical reasoning tasks. During generation, the SAE can serve as a scalable reward model to guide generations, ensuring a balanced trade-off between exploitation and exploration. This prevents extreme behaviors in either direction, ultimately fostering a higher-quality reasoning process in LLMs.