Nested Learning: The Illusion of Deep Learning Architectures — part 6/14
Cleaned-text code points [77003, 91982). Bibliography omitted here; source.md includes it.
Momentum-based optimizers are the major components of modern machine learning models’ training ( Kingma and Ba, 2014 ; Jordan et al., 2024 ; Duchi et al., 2011 ) . To explain momentum-based optimizers as associative memories, let us start from a simple gradient descent algorithm: W t + 1 = W t − η t ∇ W t ℒ ( W t , 𝒙 t + 1 ) , \displaystyle W_{t+1}=W_{t}-\eta_{t}\nabla_{W_{t}}\mathcal{L}(W_{t};\bm{x}_{t+1}), (32) which updates the current state of the weights based on the momentary gradient (surprise). This update rule does not incorporate the previous tokens and also the loss landscape that have been traversed so far, resulting in slower (or less robust) convergence in many scenarios. To fix this, momentum-based gradient descent methods incorporate an Exponential Moving Averages (EMAs) of past gradients: W ℓ t + 1 = W ℓ t + 𝒎 ℓ t + 1 \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}+{\bm{m}_{\ell}}_{{}_{t+1}} 𝒎 ℓ t + 1 = α ℓ , t + 1 𝒎 ℓ t − η ℓ , t + 1 ∇ W ℓ t ℒ ( W ℓ t ; 𝒙 t + 1 ) = α ℓ , t + 1 𝒎 ℓ t − η ℓ , t + 1 𝜹 ℓ 𝒙 ^ ℓ − 1 ⊤ , \displaystyle{\bm{m}_{\ell}}_{{}_{t+1}}=\alpha_{\ell,t+1}{\bm{m}_{\ell}}_{{}_{t}}-\eta_{\ell,t+1}\nabla_{{W_{\ell}}_{{}_{t}}}\mathcal{L}\left({W_{\ell}}_{{}_{t}};\bm{x}_{t+1}\right)=\alpha_{\ell,t+1}{\bm{m}_{\ell}}_{{}_{t}}-\eta_{\ell,t+1}\>\bm{\delta}_{\ell}\>\hat{\bm{x}}_{\ell-1}^{\top}, (33) where matrix (or vector) 𝒎 t \bm{m}_{t} is the momentum at state t t and α t \alpha_{t} and η t \eta_{t} are (adaptive) learning and momentum rates, respectively, and 𝜹 ℓ \bm{\delta}_{\ell} and 𝒙 ^ ℓ − 1 \hat{\bm{x}}_{\ell-1} are defined the same as in Equation 29 . Similar to Equation 31 and one of the examples in Section 3.1 , assuming α t + 1 = 1 \alpha_{t+1}=1 , the momentum term can be viewed as the result of optimizing the following objective with gradient descent: min 𝒎 ⟨ 𝒎 𝒙 ^ ℓ − 1 , 𝜹 ℓ ⟩ . \displaystyle\min_{\bm{m}}\>\>\langle\bm{m}\>\hat{\bm{x}}_{\ell-1},\bm{\delta}_{\ell}\rangle. (34) The case of α t + 1 ≠ 1 \alpha_{t+1}\neq 1 is equivalent to GD on the above minimization plus an ℓ 2 \ell_{2} -regularization on the momentum term. Thus, momentum can indeed be viewed as an associative memory module that learns how to compress the past gradients of the objective into its parameters. Contrary to Equation 31 , which was a simple 1-level associative memory and the update was directly applied to the memory, here the state of the momentum determines the update for the weights. In other words, it is a 2-level optimization procedure, in which the inner-loop learns the momentum and the outer-loop uses the state of the momentum to update the weights.
From this perspective, we can generalize the definition of momentum from EMAs to any arbitrary associative memory module that aims to compress the past gradients or maps the input of each token to its corresponding local error. This generalized momentum can be expressed as: W ℓ t + 1 = W ℓ t + 𝒎 ℓ t + 1 , \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}+{\bm{m}_{\ell}}_{{}_{t+1}}, (35) where 𝒎 ℓ \bm{m}_{\ell} is the solution of the following associative memory, optimized by gradient descent: min 𝒎 ℒ ~ ( 𝒎 , 𝒙 ^ ℓ − 1 , − 𝜹 ℓ ) . \displaystyle\min_{\bm{m}}\>\>\tilde{\mathcal{L}}\left(\bm{m};\>\hat{\bm{x}}_{\ell-1},-\bm{\delta}_{\ell}\right). (37) Here, the objective ℒ ~ ( ⋅ ) \tilde{\mathcal{L}}(\cdot) is different from the original objective of the problem at hand, and ℒ ~ ( ⋅ ) \tilde{\mathcal{L}}(\cdot) is the objective that defines the momentum and measures the quality of its mappings. In fact, the momentum term in this formulation aims to adapt in-context (recall that the context of the momentum is the gradients) to the local error rates based on the input of the layer. Most popular optimizers are formulated as element-wise update rule (for computational efficiency reasons) and so in Appendix B , we first explore the element-wise associative memory formulation of momentum and connect it to popular optimizers such as Adam ( Kingma and Ba, 2014 ) . Showing that Adam can be viewed as the optimal associative memory to the L 2 L_{2} -regression objective that aims to predict the variance of gradients, we discuss other similar algorithms such as RMSProp ( Hinton et al., 2012 ) , SignSGD and its momentum-based variants ( Bernstein et al., 2018 ) , NAdam ( Dozat, 2016 ) , AMSGrad ( Reddi et al., 2016 ) , RAdam ( Liu et al., 2020 ) , and Lion ( Chen et al., 2023 ) are also instances of an associative memory that aims to compress the gradients. We then go beyond element-wise formulation and show that AdaGrad ( Duchi et al., 2011 ) is also an associative memory module. Due to the connection of AdaGrad with optimizers such as Shampoo ( Gupta et al., 2018 ) and Soap ( Vyas et al., 2025 ) –i.e., as the approximation of the preconditioning term–we then conclude that all these optimizers can be re-formulated as associative memory. Next, we discuss another class of optimizers based on preconditioning and reformulate them from NL’s perspective in more details:
Preconditioning and Approximation of Hessian. Another class of algorithms is preconditioning algorithms where the idea is to approximate Hessian inverse to mimic the behavior of Newton’s algorithm. Formally, gradient descent with preconditioning is defined as: W ℓ t + 1 = W ℓ t − η t + 1 𝑷 − 1 t + 1 𝒈 ℓ t + 1 , \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}-\eta_{t+1}\>\bm{P}^{\>-1}_{t+1}\>{\bm{g}_{\ell_{t+1}}}, (38) where preconditioner 𝑷 t + 1 \bm{P}_{t+1} is often a positive-definite matrix. A critical interpretation of preconditioner is their role in performing gradient descent in a transformed coordinate system, which can be viewed as a mapping from gradients to that system of interest. Accordingly, we reformulate and interpret the preconditioner in Equation 38 as an associative memory that maps the set of gradients (or a function of gradients denoted as 𝒈 \bm{g} ) to the system of our choice, denoted as 𝒈 ^ \hat{\bm{g}} : W ℓ t + 1 = W ℓ t − η t + 1 𝑷 t + 1 − 1 ( 𝒈 ℓ t + 1 ) , \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}-\eta_{t+1}\>\bm{P}_{t+1}^{\>-1}\!\left({\bm{g}_{\ell_{t+1}}}\right), (39) where internally (in a nested level), 𝑷 t + 1 \bm{P}_{t+1} learns how to perform this mapping using an objective: min 𝑷 ℒ ~ ( 𝑷 ( 𝒈 ^ ) , 𝒈 ) . \displaystyle\min_{\bm{P}}\quad\tilde{\mathcal{L}}\left(\bm{P}\left(\hat{\bm{g}}\right);\bm{g}\right). (40) Given this viewpoint, the main question is about finding the best coordinate system that can empower the compression process. The most simple variant is an identity mapping, where we preserve the metric system and use 𝑷 \bm{P} to map 𝒈 \bm{g} (i.e., gradients in this case) to itself, resulting in preconditioning terms in Adam ( Kingma and Ba, 2014 ) and AdaGrad ( Duchi et al., 2011 ) , as discussed in Appendix B . These results, along with the representation of Adam and its variants as associative memories, show that not only momentum-based optimizers are associative memories, but they also can be decomposed into a set of nested learning problems, each of which optimized with gradient descent. In a more general form, however, one can use more nested levels and optimize the inner problems in Equation 40 with gradient descent, resulting in: 𝑷 t + 1 = 𝑷 t + 1 − ζ t + 1 ∇ 𝑷 t ℒ ~ ( 𝑷 t , 𝒈 t + 1 , 𝒈 ^ t + 1 ) . \displaystyle\bm{P}_{t+1}=\bm{P}_{t+1}-\zeta_{t+1}\nabla_{\bm{P}_{t}}\tilde{\mathcal{L}}\left(\bm{P}_{t};{\bm{g}}_{t+1},\hat{\bm{g}}_{t+1}\right). (41) In the NL framework, to design an effective preconditioning, one needs to find the right choice of 𝒈 ^ \hat{\bm{g}} and ℒ ~ . \tilde{\mathcal{L}}. This viewpoint can also lead to other classes of algorithms with gradient/momentum orthogonalization: e.g., Muon and its variants ( Jordan et al., 2024 ; Cesista, 2025 ; Keigwin et al., 2025 ) . Recalling Muon optimizer ( Jordan et al., 2024 ) : W ℓ t + 1 = W ℓ t + NewtonSchulz k ( 𝒎 ℓ t + 1 ) \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}+\texttt{NewtonSchulz}_{k}\left({\bm{m}_{\ell}}_{{}_{t+1}}\right) 𝒎 ℓ t + 1 = α ℓ , t + 1 𝒎 ℓ t − η ℓ , t + 1 ∇ W ℓ t ℒ ( W ℓ t ; 𝒙 t + 1 ) , \displaystyle{\bm{m}_{\ell}}_{{}_{t+1}}=\alpha_{\ell,t+1}{\bm{m}_{\ell}}_{{}_{t}}-\eta_{\ell,t+1}\nabla_{{W_{\ell}}_{{}_{t}}}\mathcal{L}\left({W_{\ell}}_{{}_{t}};\bm{x}_{t+1}\right), (42) where NewtonSchulz k ( ⋅ ) \texttt{NewtonSchulz}_{k}(\cdot) performs k k steps of Newton-Schulz orthogonalization process. From the above discussion about the general formulation of preconditioning, one can see NewtonSchulz k ( ⋅ ) \texttt{NewtonSchulz}_{k}(\cdot) operator as a mapping from gradients of momentum term to a proper metric system. The choice of proper coordinate system in Muon is to orthogonalize the gradients and so we aim to find a mapping 𝑷 ( ⋅ ) \bm{P}\left(\cdot\right) by minimizing a loss function min 𝑷 ℒ ~ ( 𝑷 , 𝑶 , 𝒎 ) \min_{\bm{P}}\tilde{\mathcal{L}}(\bm{P};\bm{O},\bm{m}) where objective ℒ ~ ( ⋅ , ⋅ , ⋅ ) \tilde{\mathcal{L}}(\cdot;\cdot,\cdot) measures the quality of mapping from 𝑶 \bm{O} to either 𝒎 \bm{m} or 𝒈 \bm{g} by 𝑷 ( ⋅ ) \bm{P}(\cdot) . A critical challenge in this process is that the parameter 𝑶 \bm{O} itself is not given and so the mapping requires learning both the mapping and the proper orthogonal space. A simple formulation measuring orthogonalization, can be achieved by defining the objective as: ℒ ~ ( 𝑷 ( 𝒈 ) , 𝒈 ) = ‖ 𝑷 ( 𝒈 ) ⊤ 𝑷 ( 𝒈 ) − 𝑰 ‖ 𝑭 2 , \displaystyle\tilde{\mathcal{L}}(\bm{P}(\bm{g});\bm{g})=\>\|\bm{P}(\bm{g})^{\top}\bm{P}(\bm{g})-\bm{I}\|^{2}_{\bm{F}}, (43) where 𝑷 ( 𝒈 ) \bm{P}(\bm{g}) is the orthogonal space that we aim to directly learn from gradients. This objective ensures that the gradients (or momentum) and their mapping are relatively close while the mapping is to an orthogonal space. Optimizing the above objective to find 𝑶 = 𝑷 ( 𝒈 ) \bm{O}=\bm{P}(\bm{g}) with one step of gradient descent results in: 𝑶 i + 1 = 𝑶 i − ζ i + 1 ∇ 𝑶 i ℒ ~ ( 𝑶 i , 𝒈 t ) = 𝑶 i − ζ i + 1 ( 𝑶 i − 𝒈 t + 2 𝑶 i ( 𝑶 i ⊤ 𝑶 i − 𝑰 ) ) , \displaystyle\bm{O}_{i+1}=\bm{O}_{i}-\zeta_{i+1}\nabla_{\bm{O}_{i}}\>\tilde{\mathcal{L}}\left(\bm{O}i;\bm{g}_{t}\right)=\bm{O}_{i}-\zeta_{i+1}\left(\bm{O}_{i}-\bm{g}_{t}+2\bm{O}_{i}\left(\bm{O}_{i}^{\top}\bm{O}_{i}-\bm{I}\right)\right), (44) which recovers the 3-degree polynomial (initial value 𝑶 0 = 𝒈 t \bm{O}_{0}=\bm{g}_{t} ). In a summary, the higher-frequency level learns the orthogonal mapping and then the lower-frequency process use the learned mapping to optimize the weights. Later, in Section 4.4 , we discuss a more general viewpoint that considers NewtonSchulz k ( ⋅ ) \texttt{NewtonSchulz}_{k}\left(\cdot\right) as a polynomial mapping to enhance the capacity of the memory.
### 4.3 Long Context in Optimizers: An Example of Continual Learning with Orthogonal Tasks
When removing the boundary between train and test time, and moving towards models that can continually learn for a long period of time, the role of (online) optimizers becomes more prominent: mainly due to the need for finding effective “solutions" rather than converging faster. Finding an effective solution, however, requires global understanding of the objective to avoid “local minima" as well as moving toward directions that might cause (catastrophic) forgetting of long past learned tasks. From associative memory perspective and as discussed earlier, the momentum term is expected to be a memory of past gradients, helping the optimization process to have a more global view of the loss landscape. The current design of momentum, however, acts as a simple low-pass filter that smoothifies the gradient updates and thus has limited capacity with only incorporating information from recent past.
To better illustrate this limitation, let β > 0 \beta>0 be the momentum’s decay term, and so the contribution of i i -th gradient before to the current state of the momentum can be calculated as β i ( 1 − β ) \beta^{i}\left(1-\beta\right) . Considering the cumulative sum of gradients’ contributions to the current state of the momentum term (i.e., 𝑺 t = ∑ i = 0 t β i ( 1 − β ) \bm{S}_{t}=\sum_{i=0}^{t}\beta^{i}(1-\beta) ) and the commonly used value of β = 0.9 \beta=0.9 in optimization setups, the last 6 gradients (resp. 43 gradients) are responsible for at least 50% (resp. 99%) of the cumulative contribution, i.e., 𝑺 t \bm{S}_{t} . This indicates that gradients and generally the global information beyond only past 43 steps contribute less than 1%, limiting the understanding of the objective landscape and the ability to find effective solutions. This simple example shows that the current design is limited even in incorporating the long past information, let alone its ability to properly retrieve information needed for the current state of the momentum.
Coming back to the setup of continual learning and considering one of its simple variants with orthogonal tasks. We let { ( 𝒯 i , 𝒟 i , ℒ i ) } i = 1 n \{(\mathcal{T}_{i},\mathcal{D}_{i},\mathcal{L}_{i})\}_{i=1}^{n} be the set of tasks, their corresponding data, and their objectives such that for task 𝒯 i \mathcal{T}_{i} : ℒ i ( W ) = 𝔼 ( 𝒙 , 𝒚 ) ∼ 𝒟 i [ ( W ⊤ 𝒙 − 𝒚 ) 2 ] , \displaystyle\mathcal{L}_{i}(W)=\mathbb{E}_{(\bm{x},\bm{y})\thicksim\mathcal{D}_{i}}\left[(W^{\top}\bm{x}-\bm{y})^{2}\right], (45) and gradients live in orthogonal directions of { 𝒖 i } i = 1 n \{\bm{u}_{i}\}_{i=1}^{n} . Given the task 𝒯 t \mathcal{T}_{t} (for large enough t > 1 t>1 ), when optimizing the problem with gradient descent with momentum, and after many steps on task t t , the gradients now point along 𝒖 t \bm{u}_{t} . Accordingly, the momentum term is gradually shifted and now is approximately in 𝒖 t \bm{u}_{t} direction. This can lead to gradual forgetting about past gradients, which in turn may cause catastrophic forgetting in the model. That is, the optimization process can move the weights in a direction that damages the performance on previous tasks, mainly because the optimizer has no memory of the old gradient subspace that it should avoid. This failure is not about the capacity of the model but the memory management of the optimization process, failing at finding an effective solution.
Motivated by this observation, we next discuss more expressive variants of momentum that are capable of better memory management and higher memory capacity:
### 4.4 More Expressive Designs for Momentum as an Associative Memory
So far we discussed that (1) momentum term can be viewed as an associative memory that aims to compress the (past) gradients into its parameters; and (2) for developing models that can continually learn for a long period of time and on diverse sets of tasks, the optimization process need proper information about long past and the global properties of loss landscape. Next, we discuss how Nested Learning and associative memory viewpoint could result in designing optimizers with diverse memory management/structure: