Grokking Group Multiplication with Cosets — part 5/7

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

Grokking Group Multiplication with Cosets
A. Author Contributions
Dashiell Wrote the code for training and for calculating the Group Fourier Transform over Sn. Performed the initial
analyses of models trained on S5 and initially found what we came to call the coset circuit. Designed and ran causal
experiments to confirm our understanding of the coset circuit. Derived formal properties of the coset circuit. Participated in
discussions throughout the project and in writing the paper.
Qinan Ran training jobs and performed the bulk of circuit analysis on S6, designed and ran ablation experiments and
causal interchange interventions, participated in discussions throughout the project and the writing of the paper.
Honglu Derived formal properties of the coset circuit, participated in the discussions throughout the project, and the
writing of the paper.
Stella Helped scope the problem and identify and plan the core experiments. Advised on the interpretation of the analysis
and the writing of the paper.
B. Structure of the Appendix
In the Appendix we provide more of the mathematical background needed to fully describe some of our results and
techniques. In particular, we explain the Group Fourier Transform and how we used to to analyze our models. We do this
because we believe it is of independent interest and also because it is necessary to fully explain where our results and those
of Chughtai et al. [4] diverge.
In Appendix C we go over the precise experimental set up of the models that we trained.
In Appendix D we introduce the necessary concepts from group theory needed to rigorously talk about the more mathematical
aspects of our results.
In Appendix E we introduce representation theory, representations of the symmetric group, and the group Fourier transform.
In Appendix G we return to the coset circuit and coset neurons, with the presentation grounded in the mathematical concepts
introduced in Appendices D and E.
Finally, in Appendix H we present extra graphs that did not fit in the main paper and in Appendix I we present a table of all
of the conjugacy classes of subgroups ofS5.
C. Experiment Details
We conducted experiments focusing on the permutation group ofS5 andS6. All models were trained on NVIDIA GeForce
RTX 2080 GPUs. All models were implemented in PyTorch Paszke et al.[45] and trained with the Adam optimizer [28]
with a fixed learning rate of 0.001, weight decay set to 1.0,β1 = 0.9 andβ2 = 0.98. At the beginning of each training run,
the training set is sampled uniformly from all|Sn|2 combinations of permutations. Each optimization step was made on
the entire training set. Using our setup a singleS5 model trained in approximately 8 hours and a singleS6 model trained
in approximately 100 hours, though multiple training jobs could be scheduled on a single GPU. Analysis and reverse
engineering was performed with Vink et al. [56], Nanda & Bloom [39], Harris et al. [22], GAP [16], Stein et al. [53].
Table 2. Experiment hyperparameters.
Group % Train Set Num. Runs Num. Epochs Linear Layer Size Embedding Size
S5 40% 128 250,000 128 256
S6 40% 100 50,000 256 512
D. Group Theory
In this section, let us recall some basic definitions and propositions in group theory that are relevant to this paper.
14

Grokking Group Multiplication with Cosets
D.1. Groups
A groupG is a nonempty set equipped with a special element e∈ G called the identity and a multiplication operator·
satisfying the following:
• (inverse) For each element a∈G, there exists an elementb∈G such thata·b =b·a =e.
• (identity) For each element a∈G,a·e =e·a =a.
• (associativity) For elements a,b,c ∈G, we have (a·b)·c =a· (b·c).
The inverse ofa∈G is denoted bya−1.
Example D.1. The set of integers Z along with the addition + form a group. The identity element is 0. Also with the
addition, the same is true for the set of rational numbers Q, the set of real numbers R and the set of complex numbers C.
Example D.2. The symmetric group introduced in Section 3.1 along with the composition of permutations satisfies the
group axioms. The identity element is the identity permutation leaving each element unchanged.
Example D.3. The set of natural numbers N and addition do not form a group. The reason being that the inverse elements
do not exist except for 0.
Definition D.4. Given a groupG, a subgroupH is a subset ofG such that
• a·b∈H for anya,b∈H.
• e∈H.
• a−1∈H.
One can check thatH along with the multiplication satisfies the group axiom as well. H being a subgroup ofG is denoted
byH≤G.
D.2. Cosets and double cosets
Definition D.5. Given a proper subgroup H < Gand an element g∈ G, the set gH :={gh|h∈ H} is called a left
H-coset. Similarly,Hg :={hg|h∈H} is called a rightH-coset.
gH is sometimes called a coset if the subgroupH is clear from the context. When we do not mention whether it is a left
coset or a right coset, left coset is the default.
Lemma D.6. Two cosetsg1H andg2H are either the same subset ofG or disjoint (i.e.,g1HTg2H =∅).
Lemma D.7. IfG is a finite group, any twoH-cosets have the same number of elements.
As a result, one can pick suitable representative elements (but not unique)g1,··· ,gn∈G, so thatg1H,··· ,gnH form a
partition ofG. Because the cosets have equal sizes, we can also conclude that|G| is always divisible by|H|.
Definition D.8. Given two subgroupsH,L<G and an elementg∈G, the setHgL :={hgl|h∈H,l∈L} is called the
(H,L )-double coset, or the double coset if the pair (H,L ) is clear from the context.
Double cosets enjoy the similar property as cosets:
Lemma D.9. Two double cosetsHg 1L andHg 2L are either the same or disjoint.
As a result, G can be similarly decomposed as a disjoint union of (H,L )-double cosets. However, when G is finite,
(H,L )-double cosets do not always come with equal sizes. So the decomposition is not equal-sized.
For simplicity, we call the (H,H )-double coset theH-double coset.
15

Grokking Group Multiplication with Cosets
D.3. Normal Subgroups
Definition D.10. A subgroupN is normal inG, denotedN ⊴G, if for anyg∈G and anyn∈N, we havegng−1∈N.
A subgroup is normal if and only if the left and right cosets are the same, i.e., for anyg∈G, gN =Ng . Normal subgroups
are important because they are precisely the groups for which the set ofN-cosetsG/N has a natural group structure.
Definition D.11. Given a groupG and a normal subgroup N ⊴ G, the quotient groupG/N is defined to be the set of
N-cosets endowed with the multiplication given bygN·hN =ghN for anyg,h∈G.
The well-definedness of the multiplication is a consequence ofN being normal and its group axioms are straightforward to
check.
Example D.12. IfG is commutative (for everyg,h∈G, we havegh =hg), every subgroupH≤G is normal.
Example D.13. IfG =Sn, the subgroupSn−1 fixing the first element is not a normal subgroup. On the other hand, the
alternating subgroupAn (consisting of even permutations) is a normal subgroup ofSn.
The double cosets of a normal subgroup are simply the usual cosets.
Lemma D.14. Given a normal subgroupH ⊴G, the leftH-coset and the rightH-coset are in one-to-one correspondence.
Furthermore, the set ofH-double cosets is also in one-to-one correspondence toH-cosets.
Proof. By definition,gHg−1 =H. Therefore,gH =Hg.HgH =gHH =gH.
D.4. Conjugate Subgroups
The cosets of a normal subgroupN ⊴G themselves form a group. Ifx, y∈G andx∈gN but y∈hN, thenxy∈ghN .
IfG is not abelian, however, many or even all subgroups are not normal and do not have this property. For a non-normal
subgroupH, ag /∈H gives rise to a different conjugate subgroupgHg−1.
In general, the relationship between the cosets of H andgHg−1 is complex, but they will have at least one left and one
right coset in common:Hg−1 =g−1(gHg−1). Every right cosetHx will have a left coset pairy(gHg−1) such that when
multiplied, right coset on the left and left coset on the right,Hxy (gHg−1) = Hg−1, specifically whenxy =g−1.
This relationship between the cosets of pairs of conjugate subgroups is not as powerful as that of the cosets of normal
subgroups, but conjugate subgroups are guaranteed to exist in non-abelian groups, whereas there are many simple groups
without normal subgroups at all.
This relationship between pairs of conjugate subgroups is also useful enough that it is used by every model we trained. In
general, we have the following:
Lemma D.15. For anyH≤G and an elementg∈G, the set of conjugate elementsgHg−1 forms a subgroup ofG.
If the conjugate subgroupgHg−1 is different thanH, the left and right cosetsgH,Hg are different.
The double coset circuits operate by first identifying a pair of different conjugate subgroups H andgHg−1. It exploits
the fact that the left cosetgH and the right coset (gHg−1)g are the same subset ofG, which will be fully generalized and
elaborated in the later sections.
D.5. An important case
When a groupG decomposes as only two disjointH-double cosets, any pair of subgroups conjugate toH shares a left coset
with another’s right coset.
Lemma D.16. LetH1,...,H n be conjugate subgroups of G, such that for each Hi the double coset HigHi is equal to
eitherHi orG\Hi. Then for each pair of subgroupsHi andHj there exists ag∈G such thatHig =gHj. Moreover, the
only double cosets ofHi andHj areHigHj =gHj andHixHj =G\gHj.
Proof. Ifi =j, for anyh∈Hi the shared coset is the subgroup itself. If i̸=j, becauseHi andHj are conjugate, there
exists ag∈G such thatHj =g−1Hig. The left coset is equal to the right coset:
gHj =g(g−1Hig) = Hig
16

Grokking Group Multiplication with Cosets
Notice that the double cosetHigHj =Hi(Hig) = Hig. But forx̸=g:
HixHj =Hixg−1Hig (2)
= (G\Hi)g (3)
=G\Hig (4)
E. Representation Theory
E.1. Preliminaries
Definition E.1. Given a groupG, a representation ofG is a group homomorphismρV :G→GL(V ) for some finite (but
nonzero) dimensional vector spaceV over a fieldk. When we do not specifically mentionk, we use C as the default.
In other words, a representation maps a group elementg to a linear operatorf (g) : V →V whereV is a vector space of
dimensiond, so that the group multiplication becomes compositions of linear operators (f (g·h) = f (g)◦f (h)). Without
explicit specifications, all representations in this paper are assumed to be over complex numbers. Recall also that finite
dimensional linear operators can be represented as matrices, and composition of linear operators is then given as matrix
multiplication.
When the context is clear, sometimes we omit the subscriptV in the notationρV .
The representations of finite groups have a rich and beautiful theory (see Diaconis [8], Fulton & Harris, Joe [15]). Here, we
recall a few basic definitions and facts without going into details.
Definition E.2. A representationρV :G→GL(V ) is a sub-representation ofρW :G→GL(W ) ifV can be identified as
a linear subspace ofW so thatρW (g) restricts toρV (g) for allg∈G.
Example E.3. For any groupG, the mapG→GL(V ) sending all elements to the identity matrix is a representation. When
dim(V ) = 1 , we call it the trivial representation ofG.
Definition E.4. Given two representations ρV,ρW of G, the direct sum of vector spaces V ⊕ W admits a natural
representation ofG by lettingρV,ρW act on each component separately. We call this the direct sum of representations
ρV,ρW , and denote it byρV⊕ρW .
Definition E.5. Similarly, given two representationsρV,ρW , the tensor productV⊗W admits a natural representation
ofG by acting onV,W separately and extend by linearity. We call this the tensor product of representationsρV,ρW , and
denote it byρV⊗ρW .
Definition E.6. A representationρ of a groupG is irreducible, if it does not have sub-representations other thanρ.
We denote the set of all irreducible representations ofG by Irr(G)
Lemma E.7. A representationρ of a finite groupG is a direct sum of irreducible representations.
Example E.8. The trivial representation ofG is irreducible.
Example E.9. The permutation representation mapsSn→GL(C3), i.e. 3×3 matrices with a single 1 in each row and
column and zeros everywhere else.
(2 1 3)7→


0 1 0
1 0 0
0 0 1

 (3 2 1)7→


0 0 1
0 1 0
1 0 0


You can see that the matrices of the permutation representation act on the basis vectors of C3:


0 0 1
0 1 0
1 0 0




x
y
z

 =


z
y
x


17

Grokking Group Multiplication with Cosets
What it means to be a representation is that the group multiplication becomes matrix multiplication, so just as
(2 1 3)(3 2 1) = (2 3 1) , 

0 1 0
1 0 0
0 0 1




0 0 1
0 1 0
1 0 0

 =


0 0 1
1 0 0
0 1 0


Example E.10. The permutation representation is reducible, because there is a subspace of C3 that is invariant to it’s action.


0 0 1
0 1 0
1 0 0




x
x
x

 =


x
x
x


Note that there is no permutation matrix acting on the vector
x x x T
that will change it, because all of the components
are equal.
As it turns out, there are no irreducible representations ofS3 that are three-dimensional. The largest irrep ofS3 isρ(2,1),
which is made of 2×2 matrices. The matrices of the (2, 1) irrep ofS3 are as follows:
(1 2 3)7→
1 0
0 1

(2 1 3)7→
−1 0
0 1

(3 2 1)7→
 1/2 −
√
3/2√
3/2 −1/2

(1 3 2)7→
−1/2
√
3/2√
3/2 1 /2

(3 1 2)7→
−1/2
√
3/2
−
√
3/2 −1/2

(2 3 1)7→
−1/2
√
3/2
−
√
3/2 −1/2

We leave it as an exercise to the reader to verify thatρ(2,1)(2 1 3)ρ(2,1)(3 2 1) = ρ(2,1)(2 3 1) .
Trace is an important notion in linear algebra. Taking trace of a representation induces an important map fromG to C.
Definition E.11. LetρV be a representation ofG. The character ofρV is a mapχ(ρV ) : G→ C given byχ(ρV )(g) =
tr(ρV (g)).
Lemma E.12. The character χ(ρV ) takes the same value on a conjugacy class of G. In other words, χ(ρV )(h) =
χ(ρV )(ghg−1).
To distill this property for a wider range of functions, we have the following definition:
Definition E.13. Letf :G→ C be a map. Iff (h) = f (ghg−1) for anyg,h∈G,f is called a class function.
For a finite groupG, the set of class functions form a finite-dimensional vector space. There is an important inner product
between class functions.
Definition E.14. The inner product of two class functionsϕ,ψ are defined as:
⟨ϕ,ψ⟩ = 1
|G|
X
g∈G
ϕ(g)ψ(g).
As we require the class functions to take the same values on conjugacy classes, the dimension of the vector space of class
functions is equal to the number of conjugacy classes inG. On the other hand, we have the following important theorem:
Theorem E.15. The characters of Irr(G) forms an orthonormal basis in the vector space of class functions.
Lemma E.16. For a finite groupG, Irr(G) is a finite set. Furthermore, the order of Irr(G) is equal to the number of
conjugacy classes inG.
F. Fourier transform over finite groups
Despite being mostly perceived as a powerful tool in physics and engineering, the Fourier transform has also been successfully
applied in group theory thanks to its generalization to locally compact abelian groups as well as an analog over finite groups.
18