[2311.07568] Feature emergence via margin maximization: case studies in algebraic tasks

Abstract:Understanding the internal representations learned by neural networks is a cornerstone challenge in the science of machine learning. While there have been significant recent strides in some cases towards understanding how neural networks implement specific target functions, this paper explores a complementary question -- why do networks arrive at particular computational strategies? Our inquiry focuses on the algebraic learning tasks of modular addition, sparse parities, and finite group operations. Our primary theoretical findings analytically characterize the features learned by stylized neural networks for these algebraic tasks. Notably, our main technique demonstrates how the principle of margin maximization alone can be used to fully specify the features learned by the network. Specifically, we prove that the trained networks utilize Fourier features to perform modular addition and employ features corresponding to irreducible group-theoretic representations to perform compositions in general groups, aligning closely with the empirical observations of Nanda et al. and Chughtai et al. More generally, we hope our techniques can help to foster a deeper understanding of why neural networks adopt specific computational strategies.

Abstract:Understanding the internal representations learned by neural networks is a cornerstone challenge in the science of machine learning. While there have been significant recent strides in some cases towards understanding how neural networks implement specific target functions, this paper explores a complementary question -- why do networks arrive at particular computational strategies? Our inquiry focuses on the algebraic learning tasks of modular addition, sparse parities, and finite group operations. Our primary theoretical findings analytically characterize the features learned by sty

related reading

• Zoom In: An Introduction to Circuitsdistill.pub
• Toy Models of Superpositiontransformer-circuits.pub
• Towards Monosemanticity: Decomposing Language Models With Dictionary Learningtransformer-circuits.pub
• [2602.16849] On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokkingarxiv.org
• On neural scaling and the quanta hypothesisericjmichaud.com
• Feature Visualizationdistill.pub
• Neural networks and deep learningneuralnetworksanddeeplearning.com
• What Would Non-Linear Features Actually Look Like? — Liv Gortonlivgorton.com
• Investigating the learning coefficient of modular addition: hackathon project — LessWronglesswrong.com
• A Geometric Calculator Inside a Neural Networkgoodfire.ai
• [2301.05217] Progress measures for grokking via mechanistic interpretabilityarxiv.org
• [2604.21691] There Will Be a Scientific Theory of Deep Learningarxiv.org