SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziqian Zhong*, Ziming Liu*, Max Tegmark, Jacob Andreas
Massachusetts Institute of Technology {ziqianz, zmliu, tegmark, jda}@mit.edu
Do neural networks, trained on well-understood algorithmic tasks, reliably rediscover known algorithms for solving those tasks? Several recent studies, on tasks ranging from group arithmetic to in-context linear regression, have suggested that the answer is yes. Using modular addition as a prototypical problem, we show that algorithm discovery in neural networks is sometimes more complex. Small changes to model hyperparameters and initializations can induce discovery of qualitatively different algorithms from a fixed training set, and even parallel implementations of multiple such algorithms. Some networks trained to perform modular addition implement a familiar Clock algorithm (previously described by Nanda et al. [1]); others implement a previously undescribed, less intuitive, but comprehensible procedure we term the Pizza algorithm, or a variety of even more complex procedures. Our results show that even simple learning problems can admit a surprising diversity of solutions, motivating the development of new tools for characterizing the behavior of neural networks across their algorithmic phase space.
Mechanistically understanding deep network models—reverse-engineering their learned algorithms and representation schemes—remains a major challenge across problem domains. Several recent studies [2, 3, 4, 5, 1] have exhibited specific examples of models apparently re-discovering interpretable (and in some cases familiar) solutions to tasks like curve detection, sequence copying and modular arithmetic. Are these models the exception or the rule? Under what conditions do neural network models discover familiar algorithmic solutions to algorithmic tasks?
In this paper, we focus specifically on the problem of learning modular addition, training networks to compute sums like  $8 + 6 = 2$  (mod 12). Modular arithmetic can be implemented with a simple geometric solution, familiar to anyone who has learned to read a clock: every integer is represented as an angle, input angles are added together, and the resulting angle evaluated to obtain a modular sum (Figure 1, left). Nanda et al. [1] show that specific neural network architectures, when trained to perform modular addition, implement this Clock algorithm. In this work, we show that the Clock algorithm is only one part of a more complicated picture of algorithm learning in deep networks. In particular, networks structurally similar to the ones trained by Nanda et al. preferentially implement a qualitatively different approach to modular arithmetic, which we term the Pizza algorithm (Figure 1, right), and sometimes even more complex solutions. Models exhibit sharp algorithmic phase transitions [6] between the Clock and Pizza algorithms as their width and attention strength vary, and often implement multiple, imperfect copies of the Pizza algorithm in parallel.
$^{1}$ Code is available at https://github.com/fjzzq2002/pizza.
Figure 1: Illustration of the Clock and the Pizza Algorithm.
Our results highlight the complexity of mechanistic description in even models trained to perform simple tasks. They point to characterization of algorithmic phase spaces, not just single algorithmic solutions, as an important goal in algorithm-level interpretability.
Organization In Section 2, we review the Clock algorithm [1] and show empirical evidence of deviation from it in models trained to perform modular addition. In Section 3, we show that these deviations can be explained by an alternative Pizza algorithm. In Section 4, we define additional metrics to distinguish between these algorithms, and detect phase transitions between these algorithms (and others Non-circular algorithms) when architectures and hyperparameters are varied. We discuss the relationship between these findings and other work on model interpretation in Section 5, and conclude in Section 6.
2 Modular Arithmetic and the Clock Algorithm
Setup We train neural networks to perform modular addition  $a + b = c$  (mod  $p$ ), where  $a, b, c = 0, 1, \dots, p - 1$ . We use  $p = 59$  throughout the paper. In these networks, every integer  $t$  has an associated embedding vector  $\mathbf{E}_t \in \mathbb{R}^d$ . Networks take as input embeddings  $[\mathbf{E}_a, \mathbf{E}_b] \in \mathbb{R}^{2d}$  and predict a categorical output  $c$ . Both embeddings and network parameters are learned. In preliminary experiments, we train two different network architectures on the modular arithmetic task, which we refer to as: Model A and Model B. Model A is a one-layer ReLU transformer [7] with constant attention, while Model B is a standard one-layer ReLU transformer (see Appendix F.1 for details). As attention is not involved in Model A, it can also be understood as a ReLU MLP (Appendix G).
2.1 Review of the Clock Algorithm
As in past work, we find that after training both Model A and Model B, embeddings  $(\mathbf{E}_a,\mathbf{E}_b$  in Figure 1) usually describe a circle [8] in the plane spanned by the first two principal components of the embedding matrix. Formally,  $\mathbf{E}_a\approx [\cos (w_ka),\sin (w_ka)]$  where  $w_{k} = 2\pi k / p,k$  is an integer in  $[1,p - 1]$ . Nanda et al. [1] discovered a circuit that uses these circular embeddings to implement an interpretable algorithm for modular arithmetic, which we call the Clock algorithm.
Algorithm	Learned Embeddings	Gradient Symmetry	Required Non-linearity
Clock	Circle	No	Multiplication
Pizza	Circle	Yes	Absolute value
Non-circular	Line, Lissajous-like curves, etc.	N/A	N/A
Table 1: Different neural algorithms for modular addition
"If a meeting starts at 10, and lasts for 3 hours, then it will end at 1." This familiar fact is a description of a modular sum,  $10 + 3 = 1$  (mod 12), and the movement of a clock describes a simple algorithm for modular arithmetic: the numbers 1 through 12 are arranged on a circle in  $360^{\circ} / 12 = 30^{\circ}$  increments, angles of  $10 \times 30^{\circ}$  and  $3 \times 30^{\circ}$  are added together, then this angle is evaluated to determine that it corresponds to  $1 \times 30^{\circ}$ .
Remarkably, Nanda et al. [1] find that neural networks like our Model B implement this Clock algorithm, visualized in Figure 1 (left): they represent tokens  $a$  and  $b$  as 2D vectors, and adding their polar angles using trigonometric identities. Concretely, the Clock algorithm consists of three steps: In step 1, tokens  $a$  and  $b$  are embedded as  $\mathbf{E}_a = [\cos(w_k a), \sin(w_k a)]$  and  $\mathbf{E}_b = [\cos(w_k b), \sin(w_k b)]$ , respectively, where  $w_k = 2\pi k / p$  (an everyday clock has  $p = 12$  and  $k = 1$ ). Then the polar angles of  $\mathbf{E}_a$  and  $\mathbf{E}_b$  are added (in step 2) and extracted (in step 3) via trigonometric identities. For each candidate output  $c$ , we denote the logit  $Q_{abc}$ ; the predicted output is  $c^* = \operatorname{argmax}_c Q_{abc}$ .
Crucial to this algorithm is the fact that the attention mechanism can be leveraged to perform multiplication. What happens in model variants when the attention mechanism is absent, as in Model A? We find two pieces of evidence of deviation from the Clock algorithm in Model A.
2.2 First Evidence for Clock Violation: Gradient Symmetricity
Since the Clock algorithm has logits:
$$
Q _ {a b c} ^ {\text {C l o c k}} = \left(\mathbf {E} _ {a, x} \mathbf {E} _ {b, x} - \mathbf {E} _ {a, y} \mathbf {E} _ {b, y}\right) \mathbf {E} _ {c, x} + \left(\mathbf {E} _ {a, x} \mathbf {E} _ {b, y} + \mathbf {E} _ {a, y} \mathbf {E} _ {b, x}\right) \mathbf {E} _ {c, y}, \tag {1}
$$
(see Figure 1) the gradients of  $Q_{abc}$  generically lack permutation symmetry in argument order:  $\nabla_{\mathbf{E}_a}Q_{abc}\neq \nabla_{\mathbf{E}_b}Q_{abc}$ . Thus, if learned models exhibit permutation symmetry  $(\nabla_{\mathbf{E}_a}Q_{abc} = \nabla_{\mathbf{E}_b}Q_{abc})$ , they must be implementing some other algorithm.
We compute the 6 largest principal components of the input embedding vectors. We then compute the gradients of output logits (unnormized log-probabilities from the model) with respect to the input embeddings. We then project them onto these 6 principal components (since the angles relevant to the Clock and Pizza algorithms are encoded in the first few principal components). These projections are shown in Figure 2. While Model B demonstrates asymmetry in general, Model A exhibits gradient symmetry.
Figure 2: Gradients on first six principal components of input embeddings.  $(a,b,c)$  in the title stands for taking gradients on the output logit  $c$  for input  $(a,b)$ . x and y axes represent the gradients for embeddings of the first and the second token. The dashed line  $y = x$  signals a symmetric gradient.
2.3 Second Evidence for Clock Violation: Logit Patterns
Inspecting models' outputs, in addition to inputs, reveals further differences. For each input pair  $(a,b)$ , we compute the output logit assigned to the correct label  $a + b$ . We visualize these correct logits from Models A and B in Figure 3. Notice that the rows are indexed by  $a - b$  and the columns by  $a + b$ . From Figure 3, we can see that the correct logits of Model A have a clear dependency on  $a - b$  in that within each row, the correct logits are roughly the same, while this pattern is not observed in Model B. This suggests that Models A and B are implementing different algorithms.
Figure 3: Correct Logits of Model A & Model B. The correct logits of Model A (left) have a clear dependence on  $a - b$ , while those of Model B (right) do not.
3 An Alternative Solution: the Pizza Algorithm
How does Model A perform modular arithmetic? Whatever solution it implements must exhibit gradient symmetry in Figure 2 and the output patterns in Figure 3. In this section, we describe a new algorithm for modular arithmetic, which we call the Pizza algorithm, and then provide evidence that this is the procedure implemented by Model A.
3.1 The Pizza Algorithm
Unlike the Clock algorithm, the Pizza algorithm operates inside the circle formed by embeddings (just as pepperoni are spread all over a pizza), instead of operating on the circumference of the circle. The basic idea is illustrated in Figure 1: given a fixed label  $c$ , for all  $(a,b)$  with  $a + b = c$  (mod  $p$ ), the points  $\mathbf{E}_{ab} = (\mathbf{E}_a + \mathbf{E}_b) / 2$  lie on a line though the origin of a 2D plane, and the points closer to this line than to the lines corresponding to any other  $c$  form two out of  $2p$  mirrored "pizza slices", as shown at the right of the figure. Thus, to perform modular arithmetic, a network can determine which slice pair the average of the two embedding vectors lies in. Concretely, the Pizza algorithm also consists of three steps. Step 1 is the same as in the Clock algorithm: the tokens  $a$  and  $b$  are embedded at  $\mathbf{E}_a = (\cos(w_k a), \sin(w_k a))$  and  $\mathbf{E}_b = (\cos(w_k b), \sin(w_k b))$ , respectively. Step 2 and Step 3 are different from the Clock algorithm. In Step 2.1,  $\mathbf{E}_a$  and  $\mathbf{E}_b$  are averaged to produce an embedding  $\mathbf{E}_{ab}$ . In Step 2.2 and Step 3, the polar angle of  $\mathbf{E}_{ab}$  is (implicitly) computed by computing the logit  $Q_{abc}$  for any possible outputs  $c$ . While one possibility of doing so is to take the absolute value of the dot product of  $\mathbf{E}_{ab}$  with  $(\cos(w_k c / 2), \sin(w_k c / 2))$ , it is not commonly observed in neural networks (and will result in a different logit pattern). Instead, Step 2.2 transforms  $\mathbf{E}_{ab}$  into a vector encoding  $|\cos(w_k(a - b) / 2)|(\cos(w_k(a + b)), \sin(w_k(a + b)))$ , which is then dotted with the output embedding  $U_c = (\cos(w_k c), \sin(w_k c))$ . Finally, the prediction is  $c^* = \operatorname{argmax}_c Q_{abc}$ . See Appendix A and Appendix L for a more detailed analysis of a neural circuit that computes  $\mathbf{H}_{ab}$  in a real network.
The key difference between the two algorithms lies in what non-linear operations are required: Clock requires multiplication of inputs in Step 2, while Pizza requires only absolute value computation, which is easily implemented by the ReLU layers. If neural networks lack inductive biases toward implementing multiplication, they may be more likely to implement Pizza rather than Clock, as we will verify in Section 4.
3.2 First Evidence for Pizza: Logit Patterns
Both the Clock and Pizza algorithms compute logits  $Q_{abc}$  in Step 3, but they have different forms, shown in Figure 1. Specifically,  $Q_{abc}(Pizza)$  has an extra multiplicative factor  $|\cos(w_k(a - b)/2)|$  compared to  $Q_{abc}(Clock)$ . As a result, given  $c = a + b$ ,  $Q_{abc}(Pizza)$  is dependent on  $a - b$ , but  $Q_{abc}(Clock)$  is not. The intuition for the dependence is that a sample is more likely to be classified correctly if  $\mathbf{E}_{ab}$  is longer. The norm of this vector depends on  $a - b$ . As we observe in Figure 3, the logits in Model A indeed exhibit a strong dependence on  $a - b$ .
3.3 Second Evidence for Pizza: Clearer Logit Patterns via Circle Isolation
To better understand the behavior of this algorithm, we replace the embedding matrix  $\mathbf{E}$  with a series of rank-2 approximations: using only the first and second principal components, or only the third and fourth, etc. For each such matrix, embeddings lie in a two-dimensional subspace. For both Model A and Model B, we find that embeddings form a circle in this subspace (Figure 4 and Figure 5, bottom). We call this procedure circle isolation. Even after this drastic modification to the trained models' parameters, both Model A and Model B continue to behave in interpretable ways: a subset of predictions remain highly accurate, with this subset determined by the periodicity of the  $k$  of the isolated circle. As predicted by the Pizza and Clock algorithms described in Figure 1, Model A's accuracy drops to zero at specific values of  $a - b$ , while Model B's accuracy is invariant in  $a - b$ . Applying circle isolation to Model A on the two principal components (one circle) yields a model with  $32.8\%$  overall accuracy, while retaining the first six principal components (three circles) yields an overall accuracy of  $91.4\%$ . See Appendix D for more discussion. By contrast, Model B achieves  $100\%$  when embeddings are truncated to the first six principal components. Circle isolation thus reveals an error correction mechanism achieved via ensembling: when an algorithm (clock or pizza) exhibits systematic errors on subset of inputs, models can implement multiple algorithm variants in parallel to obtain more robust predictions.
Correct Logits of Model A after Circle Isolation
Figure 4: Correct logits of Model A (Pizza) after circle isolation. The rightmost pizza is accompanying the third pizza (discussed in Section 3.4 and Appendix D). Top: The logit pattern depends on  $a - b$ . Bottom: Embeddings for each circle.
Using these isolated embeddings, we may additionally calculate the isolated logits directly with formulas in Figure 1 and compare with the actual logits from Model A. Results are displayed in Table 2. We find that  $Q_{abc}(Pizza)$  explains substantially more variance than  $Q_{abc}(Clock)$ .
Why do we only analyze correct logits? The logits from the Pizza algorithm are given by  $Q_{abc}(Pizza) = |\cos(w_k(a - b)/2)|\cos(w_k(a + b - c))$ . By contrast, the Clock algorithm has logits  $Q_{abc}(Clock) = \cos(w_k(a + b - c))$ . In a word,  $Q_{abc}(Pizza)$  has an extra multiplicative factor  $|\cos(w_k(a - b)/2)|$  compared to  $Q_{abc}(Clock)$ . By constraining  $c = a + b$  (thus  $\cos(w_k(a + b - c)) = 1$ ), the factor  $|\cos(w_k(a - b)/2)|$  can be identified.
(Unexpected) dependence of logits  $Q_{abc}(Clock)$  on  $a + b$ : Although our analysis above expects logits  $Q_{abc}(Clock)$  not to depend on  $a - b$ , they do not predict its dependence on  $a + b$ . In Figure 5, we surprisingly find that  $Q_{abc}(Clock)$  is sensitive to this sum. Our conjecture is that Step 1 and Step 2 of the  $Clock$  are implemented (almost) noiselessly, such that same-label samples collapse to the same point after Step 2. However, Step 3 (classification) is imperfect after circle isolation, resulting in fluctuations of logits. Inputs with common sums  $a + b$  produce the same logits.
Correct Logits of Model B after Circle Isolation
Figure 5: Correct logits of Model B (Clock) after circle isolation. Top: The logit pattern depends on  $a + b$ . Bottom: Embeddings for each circle.
Circle	wk	Qabc(clock) FVE	Qabc(pizza) FVE
#1	2π/59·17	75.41%	99.18%
#2	2π/59·3	75.62%	99.18%
#3	2π/59·44	75.38%	99.28%
Table 2: After isolating circles in the input embedding, fraction of variance explained (FVE) of all Model A's output logits  $\left( {{59} \times  {59} \times  {59}\text{of them}}\right)$  by various formulas. Both model output logits and formula results' are normalized to mean 0 variance 1 before taking FVE.  ${w}_{k}$  ’s are calculated according to the visualization. For example,distance between 0 and 1 in Circle #1 is 17,sow  ${w}_{k} = {2\pi }/{59} \cdot  {17}$  .
3.4 Third Evidence for Pizza: Accompanied & Accompanying Pizza
The Achilles' heel of the Pizza algorithm is antipodal pairs. If two inputs  $(a, b)$  happen to lie antipodally, then their middle point will lie at the origin, where the correct "pizza slice" is difficult to identify. For example in Figure 1 right, antipodal pairs are (1,7), (2,8), (3,9) etc., whose middle points all collapse to the origin, but their class labels are different. Models cannot distinguish between, and thus correctly classify, these pairs. Even for odd  $p$ 's where there are no strict antipodal pairs, approximately antipodal pairs are also more likely to be classified incorrectly than non-antipodal pairs.
Intriguingly, neural networks find a clever way to compensate for this failure mode. we find that pizzas usually come with "accompanying pizzas". An accompanied pizza and its accompanying pizza complement each other in the sense that near-antipodal pairs in the accompanied pizza become adjacent or close (i.e., very non-antipodal) in the accompanying pizza. If we denote the difference between adjacent numbers on the circle as  $\delta$  and  $\delta_{1}$ ,  $\delta_{2}$  for accompanied and accompanying pizzas, respectively, then  $\delta_{1} = 2\delta_{2} (\mathrm{mod} p)$ . In the experiment, we found that pizzas #1/#2/#3 in Figure 4 all have accompanying pizzas, which we call pizzas #4/#5/#6 (see Appendix D for details). However, these accompanying pizzas do not play a significant role in final model predictions<sup>2</sup>. We conjecture that training dynamics are as follows: (1) At initialization, pizzas #1/#2/#3 correspond to three different "lottery tickets" [9]. (2) In early stages of training, to compensate the weaknesses (antipodal pairs) of pizzas #1/#2/#3, pizzas #4/#5/#6 are formed. (3) As training goes on (in the presence of weight decay), the neural network gets pruned. As a result, pizzas #4/#5/#6 are not significantly involved in prediction, although they continue to be visible in the embedding space.
${}^{2}$  Accompanied pizzas #1/#2/#3 can achieve 99.7% accuracy, but accompanying pizzas #4/#5/#6 can only achieve 16.7% accuracy.
4 The Algorithmic Phase Space
In Section 3, we have demonstrated a typical Clock (Model A) and a typical Pizza (Model B). In this section, we study how architectures and hyperparameters govern the selection of these two algorithmic "phases". In Section 4.1, we propose quantitative metrics that can distinguish between Pizza and Clock. In Section 4.2, we observe how these metrics behave with different architectures and hyperparameters, demonstrating sharp phase transitions. The results in this section focus on Pizza models, but other algorithmic solutions to modular addition are also discovered, and explored in more detail in Appendix B.
We wish to study the distribution of Pizza and Clock algorithms statistically, which will require us to distinguish between two algorithms automatically. In order to do so, we formalize our observations in Section 2.2 and 2.3, arriving at two metrics: gradient symmetry and distance irrelevance.
4.1.1 Gradient Symmetricity
To measure the symmetry of the gradients, we select some input-output group  $(a,b,c)$ , compute the gradient vectors for the output logit at position  $c$  with respect to the input embeddings, and then compute the cosine similarity. Taking the average over many pairs yields the gradient symmetry. Definition 4.1 (Gradient Symmetry). For a fixed set  $S\subseteq \mathbb{Z}_p^3$  of input-output pairs<sup>3</sup>, define gradient-symmetry of a network  $M$  with embedding layer  $E$  as
$$
s _ {g} \equiv \frac {1}{| S |} \sum_ {(a, b, c) \in S} s i m \left(\frac {\partial Q _ {a b c}}{\partial \mathbf {E} _ {a}}, \frac {\partial Q _ {a b c}}{\partial \mathbf {E} _ {b}}\right),
$$
where  $\operatorname{sim}(a, b) = \frac{a \cdot b}{|a||b|}$  is the cosine-similarity,  $Q_{abc}$  is the logit for class  $c$  given input  $a$  and  $b$ . It is clear that  $s_g \in [-1, 1]$ .
As we discussed in Section 2.2, the Pizza algorithm has symmetric gradients while the Clock algorithm has asymmetric ones. Model A and Model B in Section 3 have gradient symmetricity  $99.37\%$  and  $33.36\%$ , respectively (Figure 2).
4.1.2 Distance Irrelevance
To measure the dependence of correct logits on differences between two inputs, which reflect the distances of the inputs on the circles, we measure how much of the variance in the correct logit matrix depends on it. We do so by comparing the average standard deviation of correct logits from inputs with the same differences and the standard deviation from all inputs.
Definition 4.2 (Distance Irrelevance). For some network  $M$  with correct logit matrix  $L$  ( $L_{i,j} = Q_{ij,i+j}$ ), define its distance irrelevance as
$$
q \equiv \frac {\frac {1}{p} \sum_ {d \in \mathbb {Z} _ {p}} \mathrm {s t d} \left(L _ {i , i + d} \mid i \in \mathbb {Z} _ {p}\right)}{\mathrm {s t d} \left(L _ {i , j} \mid i , j \in \mathbb {Z} _ {p} ^ {2}\right)},
$$
where std computes the standard deviation of a set. It is clear that  $q \in [0,1]$ .
Model A and Model B in Section 3 give distance irrelevance 0.17 and 0.85, respectively (Figure 3). A typical distance irrelevance from the Pizza algorithm ranges from 0 to 0.4 while a typical distance irrelevance from Clock algorithm ranges from 0.4 to 1.
4.1.3 Which Metric is More Decisive?
When the two metrics have conflicting results, which one is more decisive? We consider distance irrelevance as the decisive factor of the Pizza algorithm, as the output logits being dependent on the distance is highly suggestive of Pizza. On the other hand, gradient symmetricity can be used to rule out the Clock algorithm, as it requires multiplying (transformed) inputs which will result in asymmetric gradients. Figure 6 confirmed that at low distance irrelevance (suggesting pizza) the gradient symmetricity is almost always close to 1 (suggesting non-clock).
3To speed-up the calculations, in our experiments  $S$  is taken as a random subset of  $\mathbb{Z}_p^3$  of size 100.
Figure 6: Distance irrelevance vs gradient symmetricity over all the main experiments.
4.2 Identifying algorithmic phase transitions
How do models "choose" whether to implement the Clock or Pizza algorithm? We investigate this question by interpolating between Model A (transformer without attention) and Model B (transformer with attention). To do so, we introduce a new hyperparameter  $\alpha$  we call the attention rate.
For a model with attention rate  $\alpha$ , we modify the attention matrix  $M$  for each attention head to be  $M' = M\alpha + J(1 - \alpha)$ . In other words, we modify this matrix to consist of a linear interpolation between the all-one matrix and the original attention (post-softmax), with the rate  $\alpha$  controlling how much of the attention is kept. The transformer with and without attention corresponds to the case where  $\alpha = 1$  (attention kept) and  $\alpha = 0$  (constant attention matrix). With this parameter, we can control the balance of attention versus linear layers in transformers.