Grokking Group Multiplication with Cosets — part 3/7

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

Grokking Group Multiplication with Cosets
0 50 100 150 200 250
# of Epochs (thousands)
0
0.5
1
1.5
2
Coset Concentration
0 50 100 150 200 250
# of Epochs (thousands)
0
10
20
30
40Validation Loss
Figure 4. The paired evolution of the the validation loss and minH∈Sub(H) CH, which encodes the formation of coset circuits. Displayed
is the S5 model with random seed 1. Different runs will form coset circuits at different times in training, but the effect is representative.
0%
25%
50%
75%
100%Post-Ablation Accuracy
Remove minCH < x
1e-081e-071e-061e-050.00010.001
0.1 0.5 1.0
Coset Concentration Threshold
0%
25%
50%
75%
100%
Remove minCH > x
Figure 5. We perform ablations by re-calculating the accuracy after
removing any neuronsNi that have minH∈Sub(G) CH (Ni) greater
than (top figure) or less than (bottom figure) the thresholds on the
x-axis.
accuracy is still 100%. It is not until we set the threshold
to 10−6 that the 25th percentile moves at all. When we set
the threshold at 10−7 the performance for many models
has collapsed, but the median model has had 42 neurons
removed and the median accuracy is still 100%. Recall also
that the neuron shown in the far right of Figure 3 has a coset
concentration of 10−5.
The overwhelming majority of neurons are identifiable as
coset neurons. Of those neurons, those with the very highest
concentration on cosets account for the largest portion of
each model’s performance.
6.3. Causal Interventions
To rigorously test the properties of the coset circuit, we care-
fully designed causal experiments to test specific properties
of in the circuits. We observe a circuit’s behavior over the
entire data distribution (the full groupSn) and we see that
our model of the circuit is consistent with the behavior of
the true circuit. To confirm that our model of the circuit is
correct, however, we need to “break” the circuit in targeted
ways and test that it behaves in the way we predict. Neural
circuits are complex enough that observational evidence is
not enough. We aggregated runs over 128 S5 models of
different and recorded their average loss and accuracy. Ini-
tially, over the initial models without intervention, we have
accuracy extremely close to 1.
Embedding Exchange The left and right embeddings en-
code different information—membership in right and left
cosets, respectively—and cannot be interchanged. To test
this we intervene to switch the left and right embeddings.
After the intervention, we observed a significant drop in
accuracy to 0 and a rise in loss. This aligns with our ex-
pectation that the membership is an important property that
can’t be switched.
Switch Permutation Sign The pre-activations are sym-
metric about the origin and the sign of the pre-activations
does not matter, only whether or not the pre-activations is
equal to zero. The relative sign of the left and right pre-
activations should matter a lot. To test this, we have three
tests: changing the sign of just the left embeddings, just the
right embeddings, and both embeddings. In the case where
we change both the sign and with commutative property, we
can still expect the left and right activation to cancel out.
Therefore, we should see a near-perfect accuracy and near-0
loss. The result is as expected. When we change the sign
of only the left or right embedding, such cancellation law
doesn’t hold anymore. Therefore, we observe a 0 accuracy
in both cases.
Absolute Value Non-linearity The circuit can create a
perfect 0-1 coset membership switch with multiple neurons
on constructive interference, but every single neuron is noisy
and fundamentally limited by the ReLu non-linearity. To
7

Grokking Group Multiplication with Cosets
Table 1. Causal interventions aggregated over 128 runs on S5 with different sizes
Intervention Mean Accuracy Mean Loss
Base Model 99.99% 1.97e-6
Embedding Swap 1% 4.76
Switch Left and Right Sign 100% 1.97e-6
Switch Left Permutation Sign 0% 22.39
Switch Right Permutation Sign 0% 22.36
PerturbN (0, 0.1) 99.99% 2.96e-6
PerturbN (0, 1) 97.8% 0.0017
Absolute Value Non-Linearity 100% 3.69e-13
PerturbN (1, 1) 88% 0.029
PerturbN (−1, 1) 98% 0.0021
test this, we replace the ReLU activation function with the
absolute value function x7→| x|. We observe perfect ac-
curacy and an even lower loss that a half of the original
loss.
Distribution Change It is essential to the functioning of
each neuron that a large proportion of the pre-activations
are close to zero. To test this we compare how adding noise
from aN (−1, 1) andN (1, 1) affect the performance of the
model. We can see that changing the distribution of the
activation in PerturbN (−1, 1) changes the performance
less significantly thanN (1, 1). This indicates that the coset
requires 0 as a threshold value to decide the membership.
The results of these interventions can be viewed in Table 1
7. The Group Composition via
Representations Algorithm
Our experimental setup is identical to that of Chughtai
et al. [4], but our analysis led us to a different conclu-
sion.Chughtai et al. [4] proposed the “Group Composition
via Representations” (GCR) algorithm. They show that,
given an irrepρ ofSn, argmaxc∈Sn tr[ρ(a)ρ(b)ρ(−1c)] =
ab and propose that this is the algorithm the model is imple-
menting. This requires that not only store the matrix irreps,
but that the model perform the matrix multiplication within
its mechanism. We find that most of the evidence [ 4] put
forward is also consistent with coset circuits. The other evi-
dence we were not able to independently replicate. We also
find evidence that, to our understanding, is not consistent
with the GCR algorithm but is explained by coset circuits.
7.1. Our Interpretation of the Evidence for GCR
Chughtai et al. [4] put forward four main pieces of evidence,
which we restate here for clarity: (1) Correlation between
the model’s logits and characters of a learned representation
ρ. (2) The embedding and unembedding layers function
as a “lookup table” for the representations of the input ele-
mentsρ(a), ρ(b) and the inverse of the targetρ(c−1). (3)
The neurons in the linear layer calculate the matrix product
ρ(a)ρ(b) = ρ(ab). (4) Ablations showing that the circuit
they identify is responsible for the majority of the model’s
performance. Many of these points are equally consistent
with the coset circuit and the other we could not find evi-
dence for.
Ablations Though we do not perform all of the exact ab-
lations that Chughtai et al. [4] perform, we also find that
the weights that show high Fourier concentration and per-
form the coset multiplication are integral to the model’s
performance, see Section 6.2.
Irrep Look Up Table We were not able to find any evi-
dence that the embedding or unembedding layers function
as a look-up table for any representation except for the
one-dimensional sign representation. We did find that the
model’s weights and activations concentrate on specific
irreps in the group Fourier basis. This is due, however,
to concentration on cosets of specific subgroups, not be-
cause the matrix representations are realized anywhere in
the weights. The relationship between functions that are
constant on cosets and specific irreps is shown in Appendix
G.2.
Logit Attribution The trace of a group representation is re-
ferred to as the “character” and often denotedχ. We find that
the model’s logits correlate with the character χρ(abc−1)
when the irrep ρ appears in the Fourier transform of the
model’s weights. This is not, however, because the model
has implemented the matrix product ρ(ab)ρ(c−1), but be-
cause the model is “counting” the number of cosets that
ab and c are both in. We prove in G.2, if the cosets are
of conjugate subgroups that have their Fourier transform
concentrated on the irrepρ (as we observe for the models
in question), then the number of shared cosets will also
correlate with the characters ofρ.
Matrix Multiplication of Irreps We were not able to find
8

Grokking Group Multiplication with Cosets
any evidence that the linear layer implements matrix multi-
plication, again excluding scalar multiplication of the sign
irrep.
7.2. Evidence GCR Does Not Explain
Concentration on Cosets In the standard basis the pre-
activations of the overwhelming majority of neurons con-
centrate heavily on the cosets of subgroups. This is behavior
is not predicted by the GCR algorithm.
The Difference Between Subgroups and Irreps The GCR
algorithm and coset circuit cannot be equivalent because
there is not, in fact, a one-to-one relationship between cosets
and irreps. Most subgroups ofSn have their Fourier trans-
forms concentrate on more than a single group (see Table
5 for the spectral properties of all of the subgroups ofS5),
indeed this needs to be the case as there are many more
subgroups than irreps. Please refer to Table 4 for a concrete
comparison and Appendix 7.2 for an asymptotic analysis.
We also observe coset circuits for some subgroups such as
D104 will have coset circuits concentrated on both (3, 2) or
(2, 2, 1), depending on the run. The GCR algorithm would
treat these as different circuits, though their behavior is in
fact identical.
Unembedding Correlations of Neurons We observe that
the correlation between in the unembedding of neurons
that concentrate on the same coset is on average 81.4% (see
Table 3). The correlation between neurons concentrated only
on the same conjugacy class of subgroup (e.g.H1 andH2)
is on average−0.2%. The neurons that represent subgroups
in the same conjugacy class will oftentimes, though not
always, be concentrated on the same irrep. The model is
treating cosets together but the irreps and conjugacy classes
separately.
Coset Circuit Specific Causal Interventions The property
that the loss goes down when we replace the ReLU activa-
tion function with absolute value is a very strange property
that GCR does not predict.
The concentration of the model’s activations on irreps ofSn
is striking evidence and the GCR algorithm that [4] detail
could indeed solve the problem of group multiplication. The
coset circuit is also consistent with all of the evidence that
[4] provide and is additionally consistent with evidence that
the GCR algorithm does not explain.
8. Discussion and Conclusion
We performed a circuit level analysis to discover the con-
crete mechanism a one layer fully connected network uses
to solve group multiplication inS5 andS6. We showed that
4The dihedral group of order 10, the symmetry group of a
pentagon.
the model decomposes S5 andS6 into its cosets and uses
this structural information to perfectly implement the task.
Though our work concerns a toy problem, we highlight
a core takeaway that applies broadly to the field of inter-
pretability: we must treat proposed neural mechanisms as
theories until they have been thoroughly tested.
When we identify what we believe to be a circuit within
a larger network found via techniques such as [ 7, 19], we
have taken the first step towards mechanistically understand-
ing how a model performs a task. The evidence we have
for the circuit’s role in that task is, however, fundamentally
observational and correlational. The nodes in the circuit’s
computation graph are causally connected, but the relation-
ship between the action of those nodes is only observed to
be correlated to a certain task with respect to a distribution.
This is valuable information to have, but the understanding
that it imparts is limited and must be recognized as such.
When beginning this project we quickly noticed that the
activations of sub-circuits of our model were concentrated
on specific irreps ofSn. It was only with additional investi-
gation that we were able to attach semantic meaning to this
phenomenon. We observed that the neurons concentrated on
a single irrep were activating for specific subgroups. The hy-
pothesis of the coset circuits had formed, but it was still only
a theory. The facts we had observed were incontrovertible,
but their reason was unclear. It was only after performing
the causal experiments detailed in Section 6.3 that we
became confident we understood the mechanism. The
simple reality is that more than one theory can be consis-
tent with observational data, especially when that data only
comes from a small subset of the full distribution. There is
a long history of scholarship showing that interpretability
techniques, including state-of-the-art, can give be mislead-
ing and contradictory results [1, 2, 3, 9, 14, 23, 25, 34, 35].
In doing this work we had many advantages not available
when interpreting real-world models: access to the entire
distribution, an orthonormal basis for the function space of
the network, and a relatively small model. The task of mul-
tiplication inSn is deterministic and very well studied, we
had many mathematical tools to bring to bear in analyzing
the model. Even still, this project was quite challenging and
the circuits we found surprised us. Interpreting real models
will be even difficult. We encourage future work to apply
interpretability tools cautiously and validate observational
results with rigorous experimental tests.
Impact Statement
This paper presents work whose goal is to make the function
and mechanisms of deep neural networks interpretable to
humans. We present methods for reasoning about counter-
factual and out of distribution behavior in the models that
9