Towards a unified and verified understanding of group-operation networks

TL;DR

This paper uncovers new internal structures of neural networks trained on finite group operations, providing a unified explanation that improves understanding and verification of model performance, especially for symmetric groups.

Contribution

It introduces a comprehensive explanation of neural networks trained on group operations, unifying previous approaches and verifying model internals with improved efficiency and accuracy guarantees.

Findings

Models approximate equivariance in each input argument.

Explanation yields 3x faster accuracy guarantees than brute force.

Achieves >=95% accuracy bound for 45% of trained models.

Abstract

A recent line of work in mechanistic interpretability has focused on reverse-engineering the computation performed by neural networks trained on the binary operation of finite groups. We investigate the internals of one-hidden-layer neural networks trained on this task, revealing previously unidentified structure and producing a more complete description of such models in a step towards unifying the explanations of previous works (Chughtai et al., 2023; Stander et al., 2024). Notably, these models approximate equivariance in each input argument. We verify that our explanation applies to a large fraction of networks trained on this task by translating it into a compact proof of model performance, a quantitative evaluation of the extent to which we faithfully and concisely explain model internals. In the main text, we focus on the symmetric group S5. For models trained on this group, our explanation yields a guarantee of model accuracy that runs 3x faster than brute force and gives a >=95% accuracy bound for 45% of the models we trained. We were unable to obtain nontrivial non-vacuous accuracy bounds using only explanations from previous works.