Grokking Group Multiplication with Cosets — part 6/7

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

Grokking Group Multiplication with Cosets
The purpose the group Fourier transform serves is largely analogous to the one served by the classical Fourier transform: it
provides an alternate orthogonal basis with which to analyze functions from a groupG to either R or C.
To motivate the transition from the classical Fourier theory to the Fourier theory over groups, we start with a brief recall of
the definitions.
The classical Fourier transform over real numbers converts a complex-valued Lebesgue-integrable functionf : R→ C into
a function from the complex unit circleS1 to C with following formula:
ˆf (ξ) =
Z ∞
−∞
f (x)e−2πiξxdx. (5)
Taking one step further in abstraction, we note thate−2πiξx as a function ofx has the defining properties of turning additions
into multiplications (being a group homomorphism) and always having complex norm 1:
e−2πiξ(x1+x2) =e−2πiξx 1·e−2πiξx 2,
|e−2πiξx| = 1.
We call such functions thecharacters of R, though they are often thought of as frequencies. One can prove that all characters
of R can be written ase−2πiξx for a suitableξ∈ R.
Looking back at (5), the properties we need in order to define the Fourier transform over R are:
• R has the Lebesgue measure (allowing for integration to happen).
• R is a group (so that the characters make sense as group homomorphisms from R to the unit circle groupS1⊂ C).
Now, if we are given a finite groupG, the Fourier transform of a finite group is an operator converting a mapf :G→ C
into a function between Irr(G) and the set of linear operatorsM (V ).
Definition F.1. Given a groupG, the Fourier transformof a mapf :G→ C is a function ˆf from Irr(G) to the union of
M (Cn) for alln such that
ˆf (ρ) =
X
a∈G
f (a)ρ(a)
for an irreducible representationρ.
The analogy comes from the following similar facts:
• G, as a finite set, has the invariant discrete measure (where the “integration” becomes the sum).
• G is a group, and the irrepsρ are in a sense the ”smallest” group homomorphisms fromG toGL(n, C) (note that the
images ofρ similarly have complex-norm-1 determinants due toG being a finite group).
For more details and applications, one can refer to, for example, Elias M. Stein [11]. We would like to note that there is also
an inverse transform that restores the original functionf from ˆf:
f (g) = 1
|G|
X
ρ∈Irr(G)
dρ tr[ ˆf (ρ)ρ(g−1)] (6)
G. The Coset Circuit (with more math)
We did not introduce it in the main body of our paper because it would distract from the core of our results, but for the
first half of our investigation the Fourier transform over the symmetric group was integral to our investigation. We were
building directly on [4] who had shown striking results around the weights of single-layer models showing high degrees of
correlation with the irreps of the symmetric group. We wished to cast those results in the language of the group Fourier
transform. Even when we realized that the mechanism of the model was based around cosets it became extremely important
to understand why our coset circuit was so concentrated in Fourier space.
19

Grokking Group Multiplication with Cosets
G.1. Harmonic Analysis on the Symmetric Group
The presentation in the Appendix E was given in terms of functions on C because it is required for arbitrary groups. For Sn
all of the irreps are rational [15] and the Fourier transform of functions onSn can safely be defined over R.
In this section we describe how we use the Fourier transform to analyze the weights and activations of an MLP. The inputs to
the model are two one-hot vectors, xl, xr, which multiply the embedding matrices Elxl and Erxr. El and Er ared×|G|
matrices, whered is the embedding dimension and|G| is the size of the group. The columns are the embedding vectors for a
single elementg∈G. The normal approach would be to try and look at the column spaces of El and Er, as these columns
are the inputs to the model. However, since each row of El and Er and each value of that row is associated with a single
element ofG, we instead treat each row of the embedding as a functionf :G→ R.
In fact, anywhere in the model where a matrix or set of activations has|G| in the shape we can expand into the Fourier basis.
For non-abelian groups, each Fourier frequency is an irrep, and the Fourier transform for each irrep is matrix-valued. This is,
on its face, less interpretable than what we started with. Following the techniques outlined in Diaconis [8], however, we can
expand the function at each elementg∈G into a new Fourier basis. Concretely, if our functionf :G→ R is represented
as a vector, we know from 6 that each element of the vector is a sum of the Fourier components:


f (g1)
f (g2)
...
f (g|G|)

 = 1
|G|


P
ρdρ tr[ ˆf (ρ)ρ(g−1
1 )]P
ρdρ tr[ ˆf (ρ)ρ(g−1
2 )]
...P
ρdρ tr[ ˆf (ρ)ρ(g−1
|G|)]


We can keep track of all of the Fourier components at once by purposefully not completing the sum from 6), but instead
keep each term into a new dimension:
1
|G|


dρ1 tr[ ˆf (ρ1)ρ1(g−1
1 )] ... d ρk tr[ ˆf (ρk)ρk(g−1
1 )]
dρ1 tr[ ˆf (ρ1)ρ1(g−1
2 )] ... d ρk tr[ ˆf (ρk)ρk(g−1
2 )]
... ...
dρ1 tr[ ˆf (ρ1)ρ1(g−1
|G|)] ... d ρk tr[ ˆf (ρk)ρk(g−1
|G|)]


Though this may seem like it is only making the data more complicated, it gives us many tools for analyzing the data. In
particular, it turns out that the weights and activations are sparse in this new basis, which gives us a small path forward in
analyzing the mechanisms.
Corollary G.1. IfHi andHj are conjugate subgroups ofG such that the only two double cosets areHigHj andHiHj,
then each right coset Hix has a paired left coset yHj wherey = x−1g such that for all hx ∈ Hix and hy ∈ yHj,
hxhy∈HigHj
Lemma G.2. Letf :G→ C be constant on the cosets ofH≤G and non-zero on at least one coset. Then ˆf (ρ) = 0 if the
restriction ofρ toH,ρ|H does not contain the trivial representation as a subrepresentation.
Proof. The functionf can be decomposed as the sum of functions
fxH (σ) =
(
αx σ∈xH
0 otherwise
for each cosetxH. Because Fourier transform ˆf is invariant under translation we may, without loss of generality, analyze
only the functionfH. For a givenαx, ˆfxH (ρ) = ˆfx
H (ρ) = ρ(x) ˆfH (ρ) for allx∈G. Recall the definition of ˆfH (ρ) from
20

Grokking Group Multiplication with Cosets
F.1:
ˆfH (ρ) =
X
g∈G
fH (g)ρ(g) (7)
=αH
X
h∈H
ρ|H (h) (8)
=αH
X
h∈H
T−1[
M
τi∈T
τi(h)]T (9)
=αHT−1[
M
τi∈T
X
h∈H
τi(h)]T (10)
where in 9 we decomposeρ|H into a direct sum of irreps ofH. But because each τi is irreducible,P
h∈Hτi(h) = 0 unless
τi is the trivial irrep. Thus, unless the decomposition ofρH into irreps ofH includes the trivial representation, ˆf|H = 0
G.2. Logits and Counting Cosets
In Chughtai et al. [4], one way of justifying the GCR algorithm is to study the correlation between the character functions
and the neuron activations. We would like to argue that the correlation between the GCR and the coset membership counting
function may already exist, and in some simple cases it can be made explicit.
More precisely, we are measuring the correlation between the character function χ(ρ) of an irrep ρ with a set function
f :G→ C. In this section, we provide an explicit characterization of f in terms of trace and irreps, when f counts the
membership of cosets.
We are specifically interested in the following situation:
Lemma G.3. Suppose f : G→ C is a function such that its Fourier transform ˆf is nonzero only on an irreducible
representationρ and the trivial representation. Let ˆf (ρ) = A∈M (Cn). We have the following explicit formula:
f (σ) = dρ
|G| tr(A·ρ(σ−1)) +|H|
|G|. (11)
Proof. This is immediate by the Fourier inversion formula.
In this case, althoughf is not directly written in terms of tr(ρ(σ−1)),f is correlated with tr(ρ(σ−1)) depending on how
muchA is concentrated to the diagonal and how even are the diagonal entries. For the rest of the section, we show that
under certain conditions, Equation (11) applies verbatim to the functions that count membership of cosets for a collection of
conjugate subgroups.
Given a subgroupH≤G, let 1H be the function that takes value 1 on the subgroupH, and takes 0 otherwise. The action of
G on cosetsG/H induces a representation ofG on C|G/H| by permuting the basis accordingly. We call it the permutation
representation ofG onG/H.
Lemma G.4. The Fourier transform of 1H is nonzero only at the irreducible components of the permutation representation
ofG onG/H.
Proof. By definition, the Fourier transform of 1H on an irrepρ is
c1H (ρ) =
X
a∈H
ρ(a).
Notice that the image of P
a∈H
ρ(a) are invariant underH due to the symmetry of this expression.
LetV be the vector space whereρ acts on. Under the action of the subgroupH throughρ, one can decomposeV as irreps
ofH. We group them into two parts:
V =VH⊕V′,
21

Grokking Group Multiplication with Cosets
whereVH is a direct sum of copies of trivial representation ofH (or in other words, the invariant subspace ofV underH),
andV′ is the direct sum of nontrivial irreducible components ofV .
We immediately see the following by definition:
X
a∈H
ρ(a)|V H =|H|· IdV H.
Also by definition, nontrivial irreps ofH do not have invariant subspaces since they do not admit proper sub-representations.
Therefore, X
a∈H
ρ(a)|V ′ = 0.
As a result, c1H (ρ) is simply a scaled projection to the invariant subspace of V . Whether it is zero depends on whether
ResHρ has any trivial components.
By Frobenius reciprocity,
⟨IndG
H (1H ),χ (ρ)⟩ =⟨1H,χ (ResH (ρ))⟩H,
whereχ(ρ) is the character of the irrepχ∈ Irr(G) given by its traces, and⟨·⟩ is the inner product between class functions.
The left-hand side⟨IndG
H (1H ),χ (ρ)⟩ is nonzero if and only ifρ is an irreducible component of the permutation representation
ofG onG/H. The right-hand side⟨1H,χ (ResH (ρ))⟩H is nonzero if and only if dim(VH )̸= 0
Note that this lemma also works for 1gH for a cosetgH, since Fourier transforms turns the translation action byg into group
multiplication byρ(g).
In the double coset circuit, we are specifically interested in the membership counting functions. More specifically, let
H1,··· ,Hn be a collection of conjugate subgroups of G. Given an element σ∈ G, define the membership counting
function as
F (σ) =
nX
i=1
1σHi.
Combining all previous results, we have the following corollary describing the membership counting functionF .
Corollary G.5. If the permutation representation ofG onG/H1 has only 2 irreducible components, the Fourier transform
ˆF of the membership counting functionF is nonzero only at these 2 irreducible components. In particular, the equation (11)
applies toF .
One may wonder how restrictive it is for the permutation representation onG/H to only have 2 irreducible components.
The follow lemma shows that it applies to our case whenG =Sn andH =Sn−1.
Lemma G.6. ForSn and the subgroupSn−1 fixing one element, the permutation representation has only two irreducible
components.
Proof. The natural representation ofSn on Cn (by permuting the basis) decomposes as a direct sum of trivial representation
and the standard representation of dimensionn− 1.
Indeed, we see that when looking at the action of an individual neuron on the prediction space (i.e. “if this neuron fires,
which predictions become more likely and which less?”), we see that it is only neurons that are predicting the same coset
that are correlated. The average pairwise correlation of neuron actions is uncorrelated, as is the correlation of neurons
associated with the same irrep. Refer to Table 3 for the full results.
22

Grokking Group Multiplication with Cosets
Table 3. The correlation of unembedding neurons. Neurons that correspond to the same coset are averaged together in the unembedding,
leading to the unembedding vectors being highly correlated.
Mean Correlation Std Dev Correlation
Within Coset 0.814 0.445
Within Subgroup Conjugacy Class -0.002 0.222
Baseline -0.003 0.163
.
G.3. An Asymptotic Analysis
Our theory of coset circuits and the GCR algorithm of [4] cannot be equivalent because there is no one-to-one relationship
between irreps and subgroups. Even forS5, there are more subgroups than irreps. Quantitatively speaking, the irreps already
fail to catch up with the number of subgroups. For the direct comparison ofSn refer to
Asymptotically, the number of subgroups ofSn is bounded below as follows (see Pyber [48, Corollary 3.3]):
2( 1
16 +o(1))n2
≤| Sub(Sn)|,
whereas the number of irreps ofSn is asymptotically the following (see Erdos [13]):
| Irr(Sn)|∼ 1
4n· 3
1
2
eπ( 2
3 )
1
2n
1
2
.
We see that the former has a much higher asymptotic growth than the latter.
In practice, as can be seen in Table 5, many subgroups concentrate on more than one irrep. We do not have an explanation
for why the coset circuits always do concentrate one irrep. In practice, the different values for the cosets are arranged so
that the contributions of all but one irrep cancel out. We hypothesize that it may have something to do with the margin
maximization effect discussed in [38]. As we mention in the main body, we observe that subgroups which concentrate on
more than one irrep will form coset circuits that concentrate entirely on any of the irreps, while still behaving equivalently.
We do not think that there is in fact a connection between what the circuit is doing the irrep.
Table 4. The number of subgroups and the number of irreps from S5 to S12. The numbers of subgroups use the A005432 sequence of the
OEIS [42]. The numbers of irreps corresponds to the number of integer partitions of n and use the A000041 sequence of the OEIS [42].
S5 S6 S7 S8 S9 S10 S11 S12
Number of subgroups 156 1455 11300 151221 1694723 29594446 404126228 10594925360
Number of irreps 7 11 15 22 30 42 56 77
23

Grokking Group Multiplication with Cosets
H. Extra Graphs
H.1. Distribution over Subgroups and Cosets
0
30
60
90
120Number of Runs
A4 A5 F20
0% 40%80%
0
30
60
90
120
S3
0% 40%80%
S4
0% 40%80%
% of Neurons in Linear Layer
S2 × S3
(a) 128 Models trained on S5
0
25
50
75
100Number of Runs
A5 A6 A*
5
0%25%50%75%
0
25
50
75
100
S3 × S3
0%25%50%75%
S5
0%25%50%75%
% of Neurons in Linear Layer
S*
5 (b) 100 Models trained on S6
Figure 6. Distribution of coset circuits for models trained on S5 and S6 with different initial seeds. Every model has a few sign circuit
neurons that correspond to An < S n, but the model cannot completely solve the task with only the sign circuit, so there are never more
than a few. Every other subgroup could, with enough neurons, be used to completely solve the the multiplication, but in general if a model
primarily uses a single subgroup it is Sn−1 (in the main body of the paper we refer to these subgroups as Hi, for the element i∈ [n]
that is fixed). Every model has at least a few Sn−1 neurons. Many models use a mix of subgroups and there is often a “long tail” of a
subgroup being represented by only one or two neurons. The subgroups marked with asterisks,A∗
5 and S∗
5 , correspond to the “exceptional”
subgroups of S6, which come from an outer automorphism that only S6 has [26]. These subgroups are isomorphic to S5 and A5, but not
conjugate to the subgroups that come from fixing an element in{1, . . . ,6}.
24