# Easing Optimization Paths: a Circuit Perspective

Ambroise Odonnat, Wassim Bouaziz, Vivien Cabannes

► To cite this version:

Ambroise Odonnat, Wassim Bouaziz, Vivien Cabannes. Easing Optimization Paths: a Circuit Perspective. ICASSP 2025 - IEEE International Conference on Acoustics, Speech and Signal Processing, Apr 2025, Hyderabad, India. pp.1-5, 10.1109/icassp49660.2025.10888894. hal-05069172

HAL Id: hal-05069172

https://hal.science/hal-05069172v1

Submitted on 15 May 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Easing Optimization Paths: a Circuit Perspective

Ambroise Odonnat*

Noah's Ark Lab, Inria

ambroise.odonnat@gmail.com

Wassim Bouaziz*

FAIR, Meta AI

wesbz@meta.com

Vivien Cabannes

FAIR, Meta AI

vivc@meta.com

**Abstract**—Gradient descent is the method of choice for training large artificial intelligence systems. As these systems become larger, a better understanding of the mechanisms behind gradient training would allow us to alleviate compute costs and help steer these systems away from harmful behaviors. To that end, we suggest utilizing the circuit perspective brought forward by mechanistic interpretability. After laying out our intuition, we illustrate how it enables us to design a curriculum for efficient learning in a controlled setting. The code is available at https://github.com/facebookresearch/pal.

**Index Terms**—Gradient Descent, Optimization, Pruning, Circuits, Transformers.

## I. INTRODUCTION

Deep neural networks have attracted a lot of attention for their great empirical successes in many applications such as image classification [14], natural language processing [4], protein folding prediction [12], or playing chess [24]. Recently, large language models became an emerging technology with worldwide use [20, 16]. As the scaling of these models keeps increasing, the cost of their training becomes prohibitive. This motivates studies regarding their training dynamics to minimize the cost per amount of learned intelligence. In this introductory paper, we suggest that thinking in terms of circuits could provide valuable insights into this process. This section introduces what circuits are, the setup to ground our thoughts, and a summary of our contributions.

a) **Circuits:** Artificial neural networks were inspired by the human brain, which is seen as a complex electrical circuit of interconnected neurons [23], where the flow of information is controlled by rectifiers [8]. This is notably reflected in the naming of the most used activation layer, ReLU, for Rectified Linear Unit [6]. Recently, the notion of *circuit* has gained renewed attention in the field of mechanistic interpretability which attempts to reverse-engineer deep neural networks for greater understanding and reliability. While the term circuit is polysemous, in this recent context, a circuit describes a pathway in a neural network that transforms some inputs into a given output [5]. A neural network can be thought of as a superposition of circuits and a circuit itself can be decomposed into sub-circuits. Defining circuits acting on some inputs can be somewhat subjective, aiming to capture semantically coherent units of calculation inside a network, such as algorithmic circuits, or memorization pathways [9, 18].

Thinking in terms of circuits opens new perspectives to comprehend neural networks, and information flowing through

<figure><img src="image_2.png" /><figcaption>**Fig. 1:** Analogy between neural networks and electrical circuits, different components routing the electric/information flows. A center-tapped full wave rectifier can be implemented as a 2-layer neural network with ReLU activations. Red and blue arrows represent respectfully +1 and -1 weights. The light blue represents the bias term.</figcaption></figure>

them. The crux of this work is to study the dynamics of circuits along the training of neural networks.

b) **Setup:** To carefully study the dynamics of circuits during training, and understand how the model utilizes these circuits, we focus on a simple task: the sparse modular addition problem. Inputs are sequences of $T$ tokens in $\mathbb{F}_p = \mathbb{Z}/p\mathbb{Z} \simeq [p]$, i.e., the ring of integers modulo $p$, and targets are the sum of the first $k$ terms. Formally,

$$ \begin{cases} x = (x_1, x_2, \dots, x_T) & \text{with } x_t \in \mathbb{F}_p, \\ y = \sum_{t \le k} x_t. \end{cases} $$

In other terms, input sequences $x$ live in $\mathbb{F}_p^T$, where $T$ is the sequence length and $p$ is the vocabulary size. In practice, we take $p \in \{2, 4\}$, $T \in \{8, 12\}$, and $k = 5$.

c) **Summary of contributions:** This paper utilizes the concept of circuits to provide valuable insights into how neural networks learn and optimize their performance. Our contributions can be summarized as follows:

<footnote>*Equal contribution.</footnote>
1) We explain what circuits are and their usefulness in understanding the training behavior of neural networks: gradient descent reinforces useful circuits while pruning others.

2) We discuss how gradient descent fosters sub-circuits, helping to solve complex tasks by breaking them down into intermediate reasoning steps.

3) We detail the hardness of finding circuits and how curriculum learning and data curation are useful in easing their discovery by enhancing useful sub-circuits.

Overall, our work provides a new perspective on understanding the training dynamics of neural networks and demonstrates the potential benefits of thinking in terms of circuits for optimizing their performance.

d) Architecture: While the concept of circuit dynamics is agnostic to the choice of neural networks, we focus on the Transformer architecture [25] for its great empirical successes. Specifically, we consider a one-layer transformer with cross-attention. Given an input sequence x of length T with a vocabulary size p, our model performs the following steps:

*   **Token embeddings.** Each token $x_t$ is mapped to a d-dimensional embedding via an embedding matrix $W_E \in \mathbb{R}^{d \times p}$. This results in $z_t = W_{E,x_t} \in \mathbb{R}^d$, where $W_{E,j}$ is the j-th row of $W_E$;

*   **Positional embeddings.** A learnable positional embedding $p_t \in \mathbb{R}^d$ is added to each token $z_t$ depending on its position in the sequence, which is followed by a root-mean-square (RMS) normalization layer. It results in embeddings of the form $z_t := \text{RMS}(z_t + p_t)$,

*   **Attention block.** Given a sequence $z \in \mathbb{R}^{d \times T}$, a query and a value matrices $W_Q \in \mathbb{R}^d$, $W_V \in \mathbb{R}^{d \times d}$, our cross-attention computes

$$z_A := (W_V z) \mathrm{softmax}\left(\frac{z^\top W_Q}{\sqrt{d}}\right) \in \mathbb{R}^d.$$

Since z can be set to anything thanks to $W_E$ and p, we omit the key, hence removing the need for an extra linear transformation $W_K z$. This is again followed by an RMS normalization,

$$\tilde{z}_A := \text{RMS}(z_A) \in \mathbb{R}^d.$$

*   **Feed-forward block.** Finally, a feed-forward block is applied, consisting of a two-layer MLP with GELU activation denoted by $GELU(x) = x\varphi(x)$, where $\varphi$ is the Gaussian cumulative function. It is followed by a residual connection. The output of this layer reads

$$z_O := \tilde{z}_A + W_2^\top \text{GELU}(W_1 \tilde{z}_A) \in \mathbb{R}^d,$$

with $W_1, W_2 \in \mathbb{R}^{h \times d}$ for h a hidden dimension typically set to $h = 4d$.

*   **Unembedding.** After the last transformer layer, the embedding vector $z_O \in \mathbb{R}^d$ is mapped back to the vocabulary space through an unembedding matrix $W_U \in \mathbb{R}^{p \times d}$. The network output is a probability distribution

$$p(y|x; W_U, z_O) := \mathrm{softmax}(W_U z_O).$$

<figure><img src="image_3.png" /><figcaption>(a) Typical initial update, increasing various connections.</figcaption></figure>

<caption>(b) Final profile equalizing focus on non-spurious tokens.</caption>

<caption>**Fig. 2:** Visualization of $(s_t) = \text{softmax}((x_t + p_t)^\top W_Q / \sqrt{d})$, the attention scores for a fixed sequence made up of $(x_t) = 3$, displayed at different points during training. The strength of the attention is visualized through the thickness of the arrows, while the color indicates the sign of the last updates: red for arrows that have just been thickened, blue for those that have been thinned.</caption>

The output of the model is then fed into a cross-entropy loss to predict the right target y from the input sequence x. In all our experiments, the model is trained by gradient descent with Adam optimizer [13]. Except specified otherwise, the learning rate is $lr = 10^{-3}$ and the embedding dimension is $d = 32$.

## II. IS LEARNING ABOUT PRUNING CIRCUITS?

This section discusses gradient descent as a way to foster useful circuits and prune the other ones. It equally introduces useful visualization to understand this pruning mechanism in attention modules.

a) A Gradient Step May Enforce Many Circuits: To better comprehend the effect of an update, we visualize an attention layer using a bipartite graph structure where nodes represent both the input and output sequences and the edges are weighted according to the attention scores (the thicker the edge, the higher the attention). We illustrate the attention for a sequence with repeated entries on Figure 2. To encompass the training dynamics, we store the attention maps along the iterations, and color in blue a weight that was just diminished, and in red one that was just increased. At the beginning of training, the gradient updates sometimes reinforce attention on spurious tokens. This can be seen on the first three profiles of Figure 2, where at the first epochs, among the 64 batches of size 32, some reinforce the 3 tokens among the 8 ones, i.e. $x_5, x_6$ or $x_7$, although they do not contribute to the sum y. This is because some specific memorization circuits could be taken to memorize different batches. Meanwhile, the updates
<figure><img src="image_4.png" /><figcaption>**Fig. 3:** Evolution of the attention weights through gradient descent. Each line corresponds to a training iteration and each row corresponds to an entry $x_t$ of the input sequence $x$. The darker, the higher the attention weight. Ultimately, the transformer learns to focus solely on the first $k = 5$ input tokens, which are the ones defining the output $y$, indicated by the red vertical line. More exactly, it focuses on the 0 among these tokens, before counting them and deducing the number of 1 to make its final prediction.</figcaption></figure>

tend to push the right attention scores, augmenting the focus on the first five tokens.

Our visualization helps build a better intuition of gradient descent. During training, a gradient descent step updates the network weights proportionally to how much they locally change the output. As circuits are spread across the network, several of them can be reinforced by one update. A simple thought experiment can illustrate this mechanism. Consider an input defined as the concatenation $[x_1, x_2]$ of twice the same original input $x = x_1 = x_2$. A gradient update will reinforce both circuits to go from $x_1$ to $y$ and from $x_2$ to $y$, likely creating duplicates, due to the locality of the updates.

b) Many Steps Pruned Non-Invariant Circuits: Figure 2 illustrates how, as we iterate over the data, spurious and non-spurious features are reinforced by gradient descent. The spurious part consists of the 3 tokens on the right, i.e. $x_5$, $x_6$, and $x_7$, that do not participate in the definition of $y$. The evolution of the corresponding attention weights is non-monotonic, going from blue to red and vice-versa. On the contrary, the non-spurious part tends to be redder, ultimately becoming dominant, while the spurious circuits get pruned in the last profile.

We hypothesize that the phenomenon in Figure 2 is characteristic of gradient descent in neural networks. As we iterate over data, gradient updates randomly increase or diminish connections in spurious circuits, while non-spurious connections tend to be increased more thoroughly. As we iterate over pairs $(x, y)$, the circuits that invariantly help predict the right $y$ for all $x$ become dominant, pruning the other ones.

c) Learning Sub-Circuits: When analyzing the circuit learned when $x_t \in \mathbb{F}_p$ with $p = 2$, we observe that the network first solves for $z = \sum_{t \le k} x_t[k+1]$, before solving for $y = z[p]$. This can be seen partially from Figure 3. The attention focuses on the number of zeros among the $k$ non-

<figure><img src="image_5.png" /><figcaption>(a) Pretraining with $p = 2$.</figcaption></figure>

<figure><img src="image_6.png" /><figcaption>(b) Finetuning with $p = 4$.</figcaption></figure>

<caption>**Fig. 4:** Representation on the plan of the $d = 2$ dimensional embeddings $z_A$ obtained after the attention module (see Section 1). Colors represent the sum of the $k$ first tokens. **Left:** after pretraining with $p = 2$, we observe the emergence of equivalence classes modulo 6. After finetuning with $p = 4$, equivalence classes modulo 16 appear.</caption>

spurious tokens, leading to

$$z_A = \sum_{i<k} \mathbf{1}_{x_i=0}(W_V W_{E,0} + W_V p_i) \\ = \sum_{i<k} \mathbf{1}_{x_i=0} W_V W_{E,0} \approx k - \sum_{i<k} x_i[k+1],$$

where the second equality is due to the structure of the learned position embedding and value matrix. Figure 4 shows the learned representations $z_A$ coming out of the attention block. Learning to solve our task for $p = 2$ and $k = 5$, the model learns to map the sequences into 6 clusters based on the sum modulo 6 of the first $k$ tokens of the sequence. Similarly, finetuning from that model on $p = 4$ data, the observed structure strengthens further and the embeddings are gathered in 16 different clusters. Such a structure could not emerge from $p = 4$ only on the considered model size ($d = 2$).

In other terms, the network solves our task after splitting it into sub-tasks that involve intermediate reasoning steps. This is reminiscent of how, in image processing, convolutional neural networks (CNNs) first extract low-level features, such as texture, and extract more and more fine-grained ones with increasing depth [22]. In this case, the first CNN layers resemble Gabor filters [7] which is similar to how human vision works. The discovery of sub-circuits is not surprising. During the forward pass, the information processed at a given layer is the result of sub-circuits in previous layers. During the backward pass, gradient descent will reinforce sub-circuits in these previous layers that provide useful signals to build correct predictions.

### III. EASING OPTIMIZATION PATHS

The intuition laid out in the previous section could be of great interest for several reasons. First, formalizing it into mathematical statements may help derive useful theorems. Second, it may help discover better optimization schemes. To that end, this section provides a compelling experiment in our controlled setup. It showcases how carefully chosen curriculum learning can instill the right circuits to solve a needle-in-a-haystack problem.
<figure><img src="image_7.png" /><figcaption>**Fig. 5:** Evolution of the train and test accuracy along the training iterations. $p = 4$ corresponds to the model trained from scratch with $p = 4$ and $p = 2 \to 4$ is the model first pretrained with $p = 2$ and then finetuned with $p = 4$. The red dashed line indicates the iteration at which we switch from the pretraining to the finetuning.</figcaption></figure>

When trying to solve the modular addition problem with $p = 4$, $T = 12$, and $k = 5$ for $n = 2048$ training data, we need to find a circuit specified for $2^{11}$ points that generalize to $2^{14} > 2^{14}$ points. This is akin to finding a needle in a haystack, and as shown by Figure 5, training from scratch leads to networks that memorize their training sets but do not generalize to our testing set (blueish curves). However, if we first train the network with parity data, i.e. $x_t \in \mathbb{F}_p$ with $p = 2$, we have seen that the first attention layer will implement an addition modulo $k+1$. When initializing a network with such an already implemented sub-circuit, and then training with $p = 4$, gradient descent easily finds a circuit that can generalize to the $2^{14}$ data (pinkish curves). In practice, we train the network with a dataset where $p = 2$ for 3000 epochs in full batch, before switching the training data to $p = 4$ and continuing the training for another 7000 epochs. This explains the discontinuity (horizontal dashed line) in Figure 5.

Figure 6 illustrates the potential barrier to going from the memorizing solution learned from scratch and the generalizing ones learned with curriculum for our problem with $p = 4$. Loss profile of $f_{t \cdot \theta_4+(1-t) \cdot \theta_2 \to 4}$, where $\theta$ denotes the weight of the network $f_\theta$, $\theta_4$ are the weights learned from scratch, and $\theta_2 \to 4$ are the weights learned with the curriculum technique of Figure 5. The high loss barrier separating the two models explains why the model trained from scratch fails at leaving its local minimum and learning a general solution.

Once again, the observations made on our controlled problem are meaningful more generically. For complicated problems, finding the right circuits can be extremely hard, and require a large amount of data. In these cases, networks with high capacity can simply store their training data within disjoint memorization pathways, without generalizing to unseen data. However, in many cases, there exist solutions based on sub-circuits that are easier to learn with proxy tasks. In particular, curriculum learning may help learn modular components that can later be used to learn more complicated tasks [see 1, 2, for related observations]. This may explain the

<figure><img src="image_8.png" /><figcaption>**Fig. 6:** Loss profile of $f_{(1-t) \cdot \theta_4 + t \cdot \theta_{2 \to 4}}$, where $\theta$ denotes the weights of the network $f_\theta$, $\theta_4$ are the weights learned from scratch, and $\theta_{2 \to 4}$ are the weights learned from curriculum technique of Figure 5. It shows that to go from $\theta_4$ to $\theta_{2 \to 4}$ has to cross a high potential barrier, making it hard to escape bad local minima.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>**Fig. 7:** Low-dimensional t-SNE representation of the transformer's 10000 parameters during the gradient descent. The black circle represents the initial point and the black crosses represent the end point at the end of training. For the model $p = 2 \to 4$ that is first pre-trained with $p = 2$ and then finetuned with $p = 4$, the red circle represents the switch between pretraining and finetuning.</figcaption></figure>

practical usefulness of curriculum learning [19, 3, 11, 21], as well as the usefulness of careful data curation to enhance the right circuits [16].

## IV. CONCLUSION

In this paper, we propose a new perspective on understanding the training dynamics of neural networks through the lens of circuits. It lays out intuition regarding the benefits of this perspective and showcases how it helps design curricula to easily expose useful circuits and optimize network performance. We believe that this perspective has the potential to advance the field of deep learning both in terms of practical results and theoretical understanding.

Future work could use a circuit perspective to study the interplay between gradient norms, loss spikes, and training stability [26, 10] or the impact of activation sparsity on models' performance [15, 17]. Our work could enable researchers to first elucidate them on a small scale with the benefits of an easier visualization before tackling them on a larger scale.
REFERENCES

[1] Emmanuel Abbe, Elisabetta Cornacchia, and Aryo Lotfi. Provable advantage of curriculum learning on parity targets with mixed inputs. In *Advances in Neural Information Processing Systems*, volume 36, pages 24291–24321. Curran Associates, Inc., 2023.

[2] Vivien Cabannes, Charles Arnal, Wassim Bouaziz, Alice Yang, Francois Charton, and Julia Kempe. Iteration head: A mechanistic study of chain-of-thought, 2024.

[3] Paola Cascante-Bonilla, Fuwen Tan, Yanjun Qi, and Vicente Ordonez. Curriculum labeling: Revisiting pseudolabeling for semi-supervised learning. In *Proceedings of the AAAI conference on artificial intelligence*, volume 35, pages 6912–6920, 2021.

[4] Alexis Conneau and Guillaume Lample. Cross-lingual language model pretraining. In *Advances in Neural Information Processing Systems*, volume 32. Curran Associates, Inc., 2019.

[5] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, and et al. A mathematical framework for transformer circuits. *Transformer Circuits Thread*, 2021. https://transformer-circuits.pub/2021 framework/index.html.

[6] Kunihiko Fukushima. Visual feature extraction by a multilayered network of analog threshold elements. *IEEE Transactions on Systems Science and Cybernetics*, 5(4):322–333, 1969.

[7] D. Gabor. Theory of communication. *J. Inst. Electr. Engineering*, pages 429–457, 1946.

[8] Richard H. R. Hahnloser, Rahul Sarpeshkar, Misha A. Mahowald, Rodney J. Douglas, and H. Sebastian Seung. Digital selection and analogue amplification coexist in a cortex-inspired silicon circuit. *Nature*, 405(6789):947–951, Jun 2000.

[9] Tom Henighan, Shan Carter, Tristan Hume, Nelson Elhage, Robert Lasenby, Stanislav Fort, Nicholas Schiefer, and Olah Christopher. Superposition, memorization, and double descent. *Transformer Circuits Thread*, 2023. https://transformer-circuits.pub/2023/toy-double-descent/index.html.

[10] Romain Ilbert, Ambroise Odonnat, Vasilii Feofanov, Aladin Virmaux, Giuseppe Paolo, Themis Palpanas, and Ievgen Redko. SAMformer: Unlocking the potential of transformers in time series forecasting with sharpness-aware minimization and channel-wise attention. In *ICML*, 2024.

[11] Radu Tudor Ionescu, Bogdan Alexe, Marius Leordeanu, Marius Popescu, Dim Papadopoulos, and Vittorio Ferrari. How hard can it be? estimating the difficulty of visual search in an image. In *Proceedings of CVPR*, pages 2157–2166, 06 2016.

[12] John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, and et al. Highly accurate protein structure

prediction with alphafold. *Nature*, 596:583–589, 2021.

[13] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In *International Conference on Learning Representations (ICLR)*, San Diega, CA, USA, 2015.

[14] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In *Advances in Neural Information Processing Systems*, volume 25. Curran Associates, Inc., 2012.

[15] Zonglin Li, Chong You, and Srinadh Bhojanapalli et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In *ICLR*, 2023.

[16] Meta. The llama 3 herd of models, 2024.

[17] Seyed Iman Mirzadeh, Keivan Alizadeh-Vahid, Sachin Mehta, and Carlo C del Mundo et al. ReLU strikes back: Exploiting activation sparsity in large language models. In *ICLR*, 2024.

[18] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023.

[19] Ambroise Odonnat, Vasilii Feofanov, and Ievgen Redko. Leveraging ensemble diversity for robust self-training in the presence of sample selection bias. In *Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics*, volume 238 of *Proceedings of Machine Learning Research*, pages 595–603. PMLR, 02–04 May 2024.

[20] OpenAI. Gpt-4 technical report, 2024.

[21] Emmanouil Antonios Platanios, Otilia Stretcu, Graham Neubig, Barnabás Póczos, and Tom M. Mitchell. Competence-based curriculum learning for neural machine translation. In *Jill Burstein, Christy Doran, and Thamar Solorio, editors, NAACL-HLT (1)*, pages 1162–1172. Association for Computational Linguistics, 2019.

[22] Pablo Rivas and Mehang Rai. Gabor filters as initializers for convolutional neural networks: A study on inductive bias and performance on image classification. In *LatinX in AI Workshop at ICML 2023 (Regular Deadline)*, 2023.

[23] F. Rosenblatt. The perceptron: A probabilistic model for information storage and organization in the brain. *Psychological Review*, 65(6):386–408, 1958.

[24] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. *Nature*, 529(7587):484–489, 2016.

[25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In *Advances in Neural Information Processing Systems*, volume 30. Curran Associates, Inc., 2017.

[26] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective. In *NeurIPS*, 2024.