Grokking Group Multiplication with Cosets — part 1/7

Cleaned-text code points [0, 13152). Bibliography omitted here; source.md includes it.

Grokking Group Multiplication with Cosets
Dashiell Stander 1 Qinan Yu1 2 Honglu Fan 1 3 Stella Biderman 1
Abstract
The complex and unpredictable nature of deep
neural networks prevents their safe use in many
high-stakes applications. There have been many
techniques developed to interpret deep neural net-
works, but all have substantial limitations. Al-
gorithmic tasks have proven to be a fruitful test
ground for interpreting a neural network end-to-
end. Building on previous work, we completely
reverse engineer fully connected one-hidden layer
networks that have “grokked” the arithmetic of
the permutation groups S5 andS6. The models
discover the true subgroup structure of the full
group and converge on neural circuits that decom-
pose the group arithmetic using the permutation
group’s subgroups. We relate how we reverse en-
gineered the model’s mechanisms and confirmed
our theory was a faithful description of the cir-
cuit’s functionality. We also draw attention to
current challenges in conducting interpretability
research by comparing our work to Chughtai et al.
[4] which alleges to find a different algorithm for
this same problem.
1. Introduction
Many methods have been proposed to render deep neural
networks interpretable. There is both an academic interest
in understanding how neural networks do what they do
and a societal interest in ensuring that decisions made by
such models are sound, unbiased, and subject to human
review. These concerns are not new, nor are they unique to
deep neural networks. Many of the techniques developed
(such as SHAP values [ 33], saliency maps [ 52], gradient
attribution [51], dimension reduction [61], etc...) are still
widely used today, but there is an understanding that such
methods must be used as just one part of a careful analysis.
Naive applications of even the most sophisticated algorithms
1EleutherAI 2Brown University 3University of Geneva. Corre-
spondence to: Dashiell Stander <dash.stander@gmail.com>.
Proceedings of the 41 st International Conference on Machine
Learning, Vienna, Austria. PMLR 235, 2024. Copyright 2024 by
the author(s).
will give misleading results [1, 2, 9, 25].
Mechanistic interpretability seeks to find “neural circuits”
within deep neural networks, small sub-networks that act
as connected computation graphs and accomplish a task. In
“toy” (highly constrained) settings mechanistic interpretabil-
ity has been successful, with multiple examples where the
inner workings of neural networks have been successfully
reverse engineered end-to-end [20, 40, 41, 49, 63]. There
have also been encouraging early successes in finding in-
terpretable circuits within real-world models [ 18, 32, 35,
44, 57], but there is already work emerging that illustrates
how neural networks can resist common “mechanistic inter-
pretability” methods [14, 34, 60].
The toy interpretability projects that have succeeded have
done so in large part because a distinct ground truth circuit
that encodes the true nature of the task or environment
emerged in the model. We build on this tradition and study
a model that has perfectly learned to multiply permutations
of five and six elements, which in mathematics is known as
the symmetric groupsS5 andS6, which are deeply studied
and well-understood objects [ 8, 10, 15]. We succeed in
completely reverse engineering the model and enumerating
the diverse circuits that it converges on to implement the
multiplication of the symmetric group. Our work does not,
however, represent an unmitigated success for the project
of mechanistic interpretability. The prior work of Chughtai
et al. [4] studied the exact same model and setting, but came
to completely different conclusions. Understanding why
our and Chughtai et al. [4]’s interpretations of the same data
diverged required extensive effort (see Appendix 7 for a
thorough comparison). We find that even in a setting as
simple and well understood as group arithmetic, it is
incredibly difficult to do interpretability research and be
confident about one’s conclusions.
Our main contributions are as follows:
• We completely reverse engineer a one-hidden layer
fully-connected network trained on the permutation
groupsS5 andS6.
• We apply a methodology inspired by Geiger et al. [17]
to use causal experiments to thoroughly test all of the
properties of our proposed circuit.
• We survey current research in mechanistic interpretabil-
1

Grokking Group Multiplication with Cosets
ity and draw connections between the difficulty of our
work and broader challenges in the field.
2. Related Work
Mechanistic Interpretability Interpreting and reverse engi-
neering the mechanism used to complete a given task is an
active field in interpretability. Analysis of such mechanisms
and circuits are discovered mainly through a top-down ap-
proach of causal mediation analysis. In the previous work
Hanna et al. [21], Meng et al. [36], Tigges et al. [54], Wang
et al. [58], the circuits are composed at the “component
level” using the feed-forward layer and attention heads. We
analyze the mechanisms of neural networks at the circuit
level of individual and small groups of neurons, drawing
directly on the work of Nanda et al. [40; 41], Olah et al.
[43], Quirke & Barez [49], Zhong et al. [64], Zhang et al.
[63]. Our work builds directly on “A Toy Model of Univer-
sality” by Chughtai et al. [4]. We recreated precisely their
experimental setup for the groups S5 andS6, though we
came to different conclusions.
Grokking The models we study exhibit “grokking”,
wherein the model first memorizes the training set and
then much later generalizes to the held out data perfectly.
Grokking was first identified by Power et al. [47] and has
been well studied for its counter-intuitive training dynamics
[31, 37, 55, 50, 62, 59]. We conducted all the analysis on
fully grokked models with perfect test accuracy, as models
that show this behavior have often formed clean generalizing
circuits that are more easily interpreted [20, 40].
Group Theory We used many of the tools of group theory
for our analysis, in particular the well-developed representa-
tion theory of the symmetric group. Tools for analyzing data
on groups are well-laid out in Clausen & Baum [5], Cohen
& Welling[6], Diaconis [8], Kondor [29], Kondor & Trivedi
[30], Huang et al. [24], Karjol et al. [27], Plumb et al. [46].
3. Mathematical Preliminaries
This paper requires a familiarity with functions on groups,
a topic that is uncommon in machine learning research. In
this section we give an overview of the major concepts as
they are realized in the permutation groups that we study.
For a more formal introduction to group theory, please refer
to Appendix D.
3.1. Permutations and the Symmetric Group
A permutation of n elements is a map σ that sends one
ordering ofn elements to a different ordering. For example
the order-reversing permutation on four elements would be:
(1 2 3 4)
σ
7→ (4 3 2 1)
Left Perm
Left 
embedding
Right 
embedding+
Fully Connected Linear Layer
Unembed
Logits
Right Perm
Relu
Figure 1. Model Architecture: we follow the model architecture
used by Chughtai et al. [4]. The one-hot vectors of left and right
permutations pass through separate embeddings. We concatenate
the embeddings and pass them through a single fully-connected
hidden layer with ReLU activations. An unembedding matrix
transforms the activations into logits.
The identity permutation, denoted e, leaves the ordering
unchanged:
(1 2 3 4)
e
7→ (1 2 3 4)
We refer to specific permutations by identifying them
with the image of their action on the elements [n] :=
{1, 2,...,n } in increasing order. For the above example
we would simply denote the order reversing permutation on
four elements as (4 3 2 1) .
We multiply two permutations onn elementsσ,τ by com-
position, read from right to left. If σ = (4 3 2 1) and
τ = (3 2 1 4) , thenστ is the permutation we obtain by first
applyingτ and then applyingσ to the output ofτ:
(1 2 3 4)
τ
7→ (3 2 1 4)
σ
7→ (4 1 2 3)
First applying τ and then σ has the same effect as just
applying the permutation (4 1 2 3) . Additionally every per-
mutationσ has an inverseσ−1 such thatσσ−1 =e. These
properties makes all of the permutations on n elements a
group called the symmetric group, which we writeSn.
There are six permutations in S4 that do not change the
position of 4:
(1 2 3 4) (2 1 3 4) (3 2 1 4)
(1 3 2 4) (3 1 2 4) (2 3 1 4)
2

Grokking Group Multiplication with Cosets
These six permutations form a subgroup of S4 because
multiplication is closed within that subset, multiplying any
two permutations that leave 4 unchanged results in another
permutation that leaves 4 unchanged. You can see that these
six permutations are isomorphic toS3 by simply “forgetting”
about the 4 that is fixed in the fourth position. In the paper,
we will refer to the subgroup ofSn isomorphic toSn−1 that
leaves elementi fixed asHi.
One of the simplest types of permutations is a “transposi-
tion,” a permutation τ ∈ Sn that switches (“transposes”)
two elementsi,j ∈ [n] and leaves the remaining elements
fixed. Every element ofSn can be decomposed into a prod-
uct of transpositions. A given decomposition of a permuta-
tion is not unique, but the number of transpositions in the
decomposition is an invariant of the permutation. For a per-
mutationg∈Sn if a set of transpositions τ1τ2...τ k =g,
then every possible such set of transpositions will also have
k elements. The permutations that have an even number of
transpositions are referred to as “even” permutations and
those with an odd number are “odd.” The set of all even per-
mutations inSn is a subgroup referred to as the “alternating
group”An.
If we takeH4 < S4 and multiply every element of on the
left by some elementσ∈S4 then we get a left coset ofH4
denotedσH4. The transpositionτ = (4 2 3 1) switches the
elements in the first and fourth positions. The elements of
τH 4 are:
(4 2 3 1) (4 1 3 2) (4 2 1 3)
(4 3 2 1) (4 1 2 3) (4 3 1 2)
This coset is characterized by every element having 4 in
the first position. Every element of H4 has 4 in the fourth
position andτ switches the first and fourth positions. For
anyh∈H4,hτ has 4 in the first position becauseτ moves
it from the fourth. We would get a coset with all of the
elements ofS4 with 4 in the third position if we multiplied
H4 on the left by the any permutation that switches three
and four.
There are also right cosets where every element in a sub-
group is multiplied from the right. The elements of H4τ
are:
(4 2 3 1) (1 4 3 1) (3 2 4 1)
(4 3 2 1) (3 4 2 1) (2 3 4 1)
This right coset is characterized by every element having 1
in the fourth position.
There are in fact four subgroupsHi <S 4 that are isomor-
phic toS3, one where each element{1,..., 4} is fixed. In
general there are at leastn subgroups ofSn that are isomor-
phic toSn−1. Any twoHi, Hj are conjugate to each other.
Conjugation by an elementσ mapsx7→σxσ−1. So if we
haveH4 and conjugate it byσ = (1 4 3 2) , thenσH4σ−1
isH2:
(1 2 3 4) (3 2 1 4) (4 2 3 1)
(1 2 4 3) (4 2 1 3) (3 2 4 1)
If a subgroup is invariant to conjugation it is a normal sub-
group. The only normal subgroup of Sn forn >4 is the
alternating groupAn of even permutations.
We will mostly refer to groups by name, but we will denote
a general group as capital G and a general subgroup as
H≤ G. For a proper subgroup ( H̸= G), we will write
H <G. For a normal subgroup, we will useN ⊴G.
3.2. Fourier Transform over Groups
Though Group Fourier Transform is not central to our pre-
sentation of the coset circuit, it was an important tool that
we used to analyze the the activations of the trained models.
It is also a critical part of [ 4]. We introduce the concepts
here and go over the the similarities and differences between
our work and [4] in Section 7.
We begin with a presentation of the Discrete Fourier Trans-
form (DFT), and then present the Group Fourier Transform
by analogy. The DFT converts a function f defined on
{0, 1,...,n − 1} to a complex-valued function via the for-
mula:
ˆf (k) =
n−1X
t=0
f (t)e−2iπkt/n, k ∈{ 0,...,n − 1}
The DFT is commonly interpreted as a conversion from the
time domain to the frequency domain because thee−2iπkt/n
terms define a complex sinusoid with frequency 2πkt/n.
The frequency domain in this case means that these frequen-
cies provide an alternative orthonormal basis from which we
can work with functions. A function on{0, 1,...,n −1} can
be represented as a vector f =
x0 x1 ... x n−1

⊤
and its basis is given by the identity matrix In. The
DFT defines a basis transformation, much like any other.
The Fourier basis is given n vectors. The first ba-
sis vector, corresponding to k = 0 , is all ones. The
k = 1 basis vector is

1 e−2iπ/n ... e −2iπ(n−1)/n

,
and all of the rest for up to n− 1 are given by
1 e−2iπk/n ... e −2iπ(n−1)k/n

.
The DFT has a particularly nice interpretation as a function
on the cyclic group Cn, which is isomorphic to addition
modulon. Please refer to Appendix D or to references such
as [8, 15, 29] for a more detailed discussion.
The interpretation of the DFT as being over the cyclic groups
can be generalized to non-commutative groups. We go
3