2023

A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations

Chughtai, Bilal, Chan, Lawrence, Nanda, Neel

Understand

Universality is a key hypothesis in mechanistic interpretability -- that different models learn similar features and circuits when trained on similar tasks.

In this work, we study the universality hypothesis by examining how small neural networks learn to implement group composition.
We present a novel algorithm by which neural networks may implement composition for any finite group via mathematical representation theory.
We then show that networks consistently learn this algorithm by reverse engineering model logits and weights, and confirm our understanding using ablations.

Built on

Groups and Representations , volume 162 of Graduate Texts in Mathematics
Alperin, J. L. and Bell, R. B · 1995
Earlier work this paper cites.
Data Structures for Statistical Computing in Python
McKinney, W · 2010
Earlier work this paper cites.
Collaborative data science
Inc., P. T · 2015
Earlier work this paper cites.
In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning, April 2015
Neyshabur, B., Tomioka, R., and Srebro, N · 2015
Earlier work this paper cites.
Convergent Learning: Do different neural networks learn the same representations?, February 2016
Li, Y., Yosinski, J., Clune, J., Lipson, H., and Hopcroft, J · 2016
Earlier work this paper cites.
SymPy: Symbolic computing in Python
Meurer, A., Smith, C. P., Paprocki, M., Čertík, O., Kirpichev, S. B., Rocklin, M., Kumar, Am., Ivanov, S., Moore, J. K., Singh, S., Rathnayake, T., Vig, S., Granger, B. E., Muller, R. P., Bonazzi, F., Gupta, H., Vats, S., Johansson, F., Pedregosa, F., Curry, M. J., Terrel, A. R., Roučka, Š., Saboo, A., Fernando, I., Kulal, S., Cimrman, R., and Scopatz, A · 2017
Earlier work this paper cites.
Feature Visualization
Olah, C., Mordvintsev, A., and Schubert, L · 2017
Earlier work this paper cites.
Attention Is All You Need, December 2017
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I · 2017
Earlier work this paper cites.
Insights on representational similarity in neural networks with canonical correlation, October 2018
Morcos, A. S., Raghu, M., and Bengio, S · 2018
Earlier work this paper cites.
The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks, March 2019
Frankle, J. and Carbin, M · 2019
Earlier work this paper cites.
Similarity of Neural Network Representations Revisited, July 2019
Kornblith, S., Norouzi, M., Lee, H., and Hinton, G · 2019
Earlier work this paper cites.
PyTorch: An Imperative Style, High-Performance Deep Learning Library, December 2019
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S · 2019
Earlier work this paper cites.

Similar

Language Models are Few-Shot Learners, July 2020
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D · 2020
Cited alongside, same era.
Thread: Circuits
Cammarata, N., Carter, S., Goh, G., Olah, C., Petrov, M., Schubert, L., Voss, C., Egan, B., and Lim, S. K · 2020
Cited alongside, same era.
Array programming with NumPy
Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E · 2020
Cited alongside, same era.
Zoom In: An Introduction to Circuits
Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S · 2020
Cited alongside, same era.
Revisiting Model Stitching to Compare Neural Representations, June 2021
Bansal, Y., Nakkiran, P., and Barak, B · 2021
Cited alongside, same era.
A mathematical framework for transformer circuits, 2021
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C · 2021
Cited alongside, same era.
Multimodal neurons in artificial neural networks
Goh, G., †, N. C., †, C. V., Carter, S., Petrov, M., Schubert, L., Radford, A., and Olah, C · 2021
Cited alongside, same era.
Thinking Like Transformers, July 2021
Weiss, G., Goldberg, Y., and Yahav, E · 2021
Cited alongside, same era.
Unifying Grokking and Double Descent
Davies, X., Langosco, L., and Krueger, D · 2022
Cited alongside, same era.
Predictability and Surprise in Large Generative Models
Ganguli, D., Hernandez, D., Lovitt, L., DasSarma, N., Henighan, T., Jones, A., Joseph, N., Kernion, J., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Elhage, N., Showk, S. E., Fort, S., Hatfield-Dodds, Z., Johnston, S., Kravec, S., Nanda, N., Ndousse, K., Olsson, C., Amodei, D., Amodei, D., Brown, T., Kaplan, J., McCandlish, S., Olah, C., and Clark, J · 2022
Cited alongside, same era.
Acquisition of Chess Knowledge in AlphaZero
McGrath, T., Kapishnikov, A., Tomašev, N., Pearce, A., Hassabis, D., Kim, B., Paquet, U., and Kramnik, V · 2022
Cited alongside, same era.
A Comprehensive Mechanistic Interpretability Explainer & Glossary
Nanda, N · 2022
Cited alongside, same era.

Then

In-context Learning and Induction Heads, September 2022
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C · 2022
Later among the works it cites.
Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, January 2022
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V · 2022
Later among the works it cites.
The MultiBERTs: BERT Reproductions for Robustness Analysis, March 2022
Sellam, T., Yadlowsky, S., Wei, J., Saphra, N., D’Amour, A., Linzen, T., Bastings, J., Turc, I., Eisenstein, J., Das, D., Tenney, I., and Pavlick, E · 2022
Later among the works it cites.
Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 small, November 2022
Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J · 2022
Later among the works it cites.
Emergent Abilities of Large Language Models, October 2022
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., Chi, E. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W · 2022
Later among the works it cites.
Unveiling Transformers with LEGO: A synthetic reasoning task, July 2022
Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., and Wagner, T · 2022
Later among the works it cites.
Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit, January 2023
Barak, B., Edelman, B. L., Goel, S., Kakade, S., Malach, E., and Zhang, C · 2023
Closest in time.
Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task, February 2023
Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., and Wattenberg, M · 2023
Closest in time.
Tracr: Compiled Transformers as a Laboratory for Interpretability, January 2023
Lindner, D., Kramár, J., Rahtz, M., McGrath, T., and Mikulik, V · 2023
Closest in time.
TransformerLens, January 2023
Nanda, N · 2023
Closest in time.
Progress measures for grokking via mechanistic interpretability, January 2023
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J · 2023
Closest in time.
Individual differences among deep neural network models
Mehrer, J., Spoerer, C. J., Kriegeskorte, N., and Kietzmann, T. C · 2041
Closest in time.

Beyond the bibliography

alphaXiv searches the wider corpus for related work and actual follow-ups.

Open on alphaXiv

alphaXiv is searching for related work…

Reading the bibliography…