Grokking Group Multiplication with Cosets — part 2/7
Cleaned-text code points [13152, 25642). Bibliography omitted here; source.md includes it.
Grokking Group Multiplication with Cosets over the construction in Appendix E and F. The high level interpretation, however, is the same. For functions from Sn→ C there is an orthonormal basis that is equivariant to translations and convolutions. The frequencies for the Fourier transform overSn are given by the partitions ofn. The “highest” frequencies can be interpreted as representing functions that are constant on permutations that all agree on a small number of elements of [n] [12]. 4. Model Architecture As shown in Figure 1, the model we study contains separate left and right embeddings, followed by a fully connected linear layer with ReLU activations, and an unembedding layer. We use the same architecture as in [ 4] to enable consistent comparisons. 1 • One hot vectors xg with length|G|. • Two embedding matrices, El, Er with dimensions (d, |G|), where d is embedding dimension. Sn is non-abelian, i.e. not commutative, and the separate embeddings are to give the model extra capacity. • A linear layer W with dimension (w, 2d),w denoting the width of the linear layer. After the linear layer we apply the ReLU pointwise nonlinearity. • An unembedding layer U with dimension (|G|, w), which transforms the outputs of the ReLU and linear layer to into logit space for the group. We also note that the firstd columns of the linear layer will only act on the left embeddings and the secondd columns will only act on the right embeddings, so we can analyze W as the concatenation of two (w, d) matrices: W = [L R]. W Elxg Erxh = LElxg + RErxh Throughout the paper will refer to the values LElxg, RErxh, and their sum as “pre-activations” to denote that the ReLU activation function has not been applied. Post- ReLU values we refer to as “activations.” 5. Coset Circuits 5.1. Sign Neurons Implement the Sign Circuit The even permutations form a subgroup called the alternat- ing group An. The two cosets of An are the group itself and all of the odd permutations, τAn. The multiplication 1All code necessary for reproducing results and analysis is available at https://www.github.com/dashstander/ sn-grok Lel -1 Odd Left Pre-activation 1 Even Rer 1 Odd Right Pre-activation -1 Even + +++ -2 0 0 ReLU ReLU 0 2 Activation 67% Even, 33% Odd Activation 100% Odd 2 Figure 2. A diagram showing the four possible paths through a single neuron (i.e. one row of REr) that implements part of a “sign circuit.” The model stores whether a permutation is “even” or “odd” in the embeddings, represented in the left or right pre- activation values. The pre-activations are added together and then the ReLU activation is applied. The neuron only fires when the left permutation is even and the right is odd. If the neuron does not fire, then in 1/3 cases the product is odd and 2/3 it is even. of even and odd permutations has similar features to the addition of even and odd integers (hence the name). The sign map on a permutation inSn, sgn, is given by: sgn(σ) = ( 1 σ∈An −1 σ∈τAn An “even” permutation that is in An is mapped to 1 and an “odd” permutation not inAn is mapped to−1. For any σ,ρ∈Sn, the sign of their product is the product of their signs: sgn(σρ) = sgn(σ) sgn(ρ). The one-layer model that we train uses this relationship to help solve the general group multiplication. Every single model we trained had at least two neurons dedicated to encoding the sign of the permutation product. Though the model cannot use the alternating group to completely solve multiplication inSn, this sign circuit is emblematic of the general coset circuits the model forms. Consider the neuron shown in Fig. 2. The left pre- activations are given byL(σ) = sgn(σ) and the right pre- activations are R(σ) = − sgn(σ). The full action of the neuron is given by ReLU(L(σl) + R(σr)) and there are three cases: 1. sgn(σl) = sgn( σr)⇒ sgn(σlσr) = 1 . In this case 4 Grokking Group Multiplication with Cosets L(σl) andR(σr) destructively interfere, cancelling out to 0. Both the pre-activation and activation are 0. 2. sgn(σl) = −1, sgn(σr) = 1 ⇒ sgn(σlσr) = −1. In this caseL(σl) andR(σr) reinforce each other and sum to a positive value. Since 2 > 0, the activation value is 2. 3. sgn(σl) = 1 , sgn(σr) = −1⇒ sgn(σlσr) = −1. Like in (2) the productσlσr is an odd permutation and L(σl) andR(σr) constructively interfere, though this timeL(σl) +R(σr) =−2, which is less than 0. Thus ReLU clips the pre-activation and sends it to 0. 5.2. Conjugate Subgroup Circuit All four ways to multiply two cosets ofAn are well-defined. For each of the four options (even-even, odd-even, etc...) we know which coset of An the product will be in, but no other subgroup of Sn has this property. The model instead learns to use sets of conjugate subgroups. Recall thatHi <S n is the subgroup isomorphic toSn−1 that fixes the elementi∈ [n] in theith place andτij is the permutation that swaps i and j. Any two Hi and Hj are conjugate to each other, τijHiτij = Hj and τijHjτij = Hi. This means that there are two shared cosets betweenHi andHj, becauseHiτij = τijHj andHjτij = τijHi. The model implements the full group multiplication by picking out the shared cosets of conjugate subgroups. As an example, consider a neuron that corresponds to H1 for the left permutation and H5 for the right permutation. The shared coset is H1τ15 =τ15H5, the set of all σ∈S5 withσ(1) = 5 . The pre-activations for the left and right permutations will be: L(σ) = 4 σ∈H1 2 σ∈H1τ12 0 σ∈H1τ13 −2 σ∈H1τ14 −4 σ∈H1τ15 R(σ) = −4 σ∈τ15H5 −2 σ∈τ25H5 0 σ∈τ35H5 2 σ∈τ45H5 4 σ∈H5 (1) The final activation is still ReLU(L(σl) +R(σr)), but now there are twenty-five possible pairs of cosets. All twenty- five combinations can be boiled down to two meaningful cases: 1. IfL(σl) +R(σr) = 0 , thenσlσr is in the shared coset H1τ15. 2. IfL(σl) +R(σr)̸= 0, thenσlσr is not in the shared cosetH1τ15. Each left cosetyH5 has a paired right cosetH1x such that H1xyH5 = H1τ15 = τ15H5. The discrete values that L andR can take are precisely tuned so that those pairs of left and right cosets cancel out. Just like with the sign neuron, information about the pre-activation being negative is lost with the ReLU. This lost information has to be made up with extra neurons that correspond to (H1,H 5) and assign different values to the cosets. For example, a different neuron that uses −L(σl)− R(σr) will fail to fire for a different set of permutations. The combination ReLU(L(σl) +R(σr)) + ReLU(−L(σl)−R(σr)) will be much closer to a perfect on/off switch for coset membership. 5.3. Decoding Permutations with Coset Membership There are n2 combinations of (Hi,Hj) subgroups. Each pair can be interpreted directly as encoding the set of per- mutations withi in thejth position. Because of the way the coset neurons function, each neuron is better understood as firing when the value in the jth position is certainly noti. Then2 combinations of (Hi,Hj) uniquely identify each ele- ment ofSn. We can use the outputs of twenty-five (Hi,Hj) neurons as a code that uniquely encodes each element ofS5. By analyzing the unembedding layer to see how the model makes use of (Hi,Hj) neurons, we see that this is almost exactly what the model does. This same construction works for every subgroup ofSn except forAn. 6. The Process of Reverse Engineering 6.1. Identifying Coset Circuits The first step in attempting to reverse engineer the mecha- nisms of a neural network is to spend some time staring at the weights and activations. Even a small one-layer model such as ours is too large to visualize all at once. It was not until we looked closely at the pre-ReLU activations that we produced a histogram similar to Figure 3. The left and right pre-activations of one neuron were nearly constant on the distinct cosets of the Frobenius group of order 20 (F20), one of the subgroups ofS5. 2 Further investigation revealed that almost every neuron had this property of only producing a discrete number of values that corresponded directly to the cosets of one of the subgroups ofS5 orS6. For a function f :G→ R, we defineCH (f ) to be the degree to whichf concentrates on the cosetsH≤G: CH (f ) := P gH Var[f|gH ] Var[f ] Where Var[f|gH ] is the variance off when the domain is 2F20 is equivalent to the group of affine transformations x7→ ax + b, where a, b, x are in the field with five elements and a̸= 0. 5 Grokking Group Multiplication with Cosets 3 0 3 6 0 5 10 15 20# of Permutations 100,000 Steps Coset F20 F20(12354) F20(12435) F20(12453) F20(12534) F20(12543) 3 0 3 6 110,000 Steps 3 0 3 6 120,000 Steps 3 0 3 6 Pre-Activation Value 130,000 Steps Figure 3. An illustration of the phenomenon of “concentration on cosets,” depicting the 115th neuron from seed 11. We show the evolution of the left pre-activations (the pre-ReLU outputs of a layer) of training on anF20 neuron from 100k to 130k steps. The seed of the neuron’s functionality is already present at 100k steps, where it fires very strongly and negatively for permutations in the coset F20(1 2 3 5 4) , but it takes time for the action of the neuron to “clean up” on the other cosets of F20. The distribution found at 130k steps does not change very much afterwards. Noticing this common pattern of neurons taking on these discrete values was a striking piece of evidence that required further investigation. restricted to the cosetgH. IntuitivelyCH (f ) calculates the degree to which restricting to the cosets of H reduces the variance off. IfCH (f )< 1 it implies that the activations f can meaningfully be understood better by looking at the values that it takes on the cosets of some subgroup. Recall that a single neuron is a function Ni : Sn×Sn→ R is the sum of two functionsG→ R, one for the left and right permutations, respectively. We can calculate minCH for each. Take as an exampleNl 115, the neuron shown in Figure 3. At 100,000 steps (on the far left) Var[Nl 115] = 5 .23. Its activations are not concentrated on the specific cosets of F20, however, and CF20 (Nl 115) = 2 .96. At 130,000 steps (on the far right) Var[Nl 115] has increased to 9.06, but CF20 (Nl 115)< 10−5. The distribution within each coset of F20 has close to zero variance. We see a typical example of what this looks like for the entire model in Fig. 4. As the validation loss approaches a small value, there is a rapid transition from the median coset concentration being approximately 1, to a minuscule value. Even if it is apparent that a neuron is taking on discrete values and is a good candidate for being a coset neuron, it is difficult to tell by sight which subgroup the neuron is activating for.S5 andS6 only have 156 and 1,455 subgroups, respectively, 3 so it is tractable to do an exhaustive search and calculate argminH∈Sub(G)CH (f ) the subgroup that minimizes the variance of f for ev- ery neuron in the model. Running these calculations shows that for the 128 S5 models and 100 S6 models we 3Sequence A005432 OEIS [42] trained over 99.2% of the neurons in the linear layer had minH∈Sub(G)CH (f )< 1.0, and the vast majority of those were less than 10−6. With the ability to calculate directly which neurons corre- sponded to which subgroup, our theories for exactly what the neurons were representing fell into place. The next step was to confirm that these neurons were actually responsible for the models’ performance. 6.2. Ablations We have described how coset neurons function and how they can be identified. We will now show via ablations that coset neurons are not solely sufficient but also neces- sary to implement multiplication inSn. We conduct abla- tions by removing neurons which have a coset concentration minH∈Sub(G)CH (Ni) above a threshold. If coset circuits are in fact responsible for the performance of our models, then we expect to see no change in the ac- curacy when the neurons that have not converged onto the cosets of a subgroup are removed from the model. This is precisely what we see on the far right of Figure 5. Of the 128 S5 models that we trained, 126 models saw no change in the accuracy when we removed the neurons with minH∈Sub(G)CH (Ni)≥ 1 (the far right of Figure 5). Re- call that if CH (Ni)≥ 1, restricting to the cosets of H at best does not change the variance off. Of the two models that did show a decrease in accuracy, they decreased to 99% and 98%. We see more between-run variation when we remove more neurons. The median model has 24 out of 128 neurons with minCH (Ni)≥ 10−5, but the 50th and 25th percentile 6