Nested Learning: The Illusion of Deep Learning Architectures โ part 7/14
Cleaned-text code points [91982, 107836). Bibliography omitted here; source.md includes it.
Extension: More Expressive Association. As discussed earlier, vanilla momentum term can be viewed as a value-less associative memory. To allow more expressive associative memory and following the original definition of associative memory (i.e., mapping keys to values), we let value parameter ๐ i = ๐ท i \bm{v}_{i}=\bm{P}_{i} and so the momentum aims to minimize: min ๐ โจ ๐ โ โ ( W i ; ๐ i ) โค , ๐ท i โฉ , \displaystyle\min_{\bm{m}}\;\;\langle\bm{m}\>\nabla\mathcal{L}(W_{i};\bm{x}_{i})^{\top},\bm{P}_{i}\rangle, (46) or equivalently, minimizes โจ ๐ , ๐ท i โ โ ( W i ; ๐ i ) โฉ \langle\bm{m},\bm{P}_{i}\nabla\mathcal{L}(W_{i};\bm{x}_{i})\rangle . Using gradient descent for updating the momentum results in the update rule of: W i + 1 = W i + ๐ฆ i + 1 \displaystyle W_{i+1}=W_{i}+\mathbf{m}_{i+1} ๐ฆ i + 1 = ๐ฆ i โ ฮท i ๐ท i โ โ ( W i ; ๐ i ) , \displaystyle\mathbf{m}_{i+1}=\mathbf{m}_{i}-\eta_{i}\bm{P}_{i}\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right), (47) where we also included a forget gate ฮฑ i + 1 \alpha_{i+1} for fogetting the past. This update rule can be viewed as preconditioning the momentum GD. Here, the momentum term with preconditioning is interpreted as an associative memory that learns how to compress the mappings between ๐ท i \bm{P}_{i} and the gradient term โ โ โ ( W i , ๐ i ) \nabla\mathcal{L}(W_{i};\bm{x}_{i}) . While any reasonable choice (e.g., random features) of preconditioning may also be used to improve the expressivity of the initial version of GD with momentum. Note that our discussion on preconditioning in Equation 39 is different from here and indicates that one can learn the preconditioner as a proper mapping, while the above formulation shows when using preconditioning, momentum term (as a memory for gradients) aims to map them to their mapping function.
Extension: More Expressive Objectives. Revisiting the formulation of momentum: for a given gradients โ โ โ ( W i , ๐ i ) \nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right) , the momentum mapping is based on dot-product similarity as the internal objective (as discussed above) and so its update rule is a Hebbian-rule ( Hebb, 2005 ) . From an associative memory perspective, this update rule has a limited capacity ( Storkey, 1997 ) and makes the update rule of momentum independent of its current state. Thus it limits its ability to track/compress the loss landscape information. A natural extension is to replace the internal objective with L 2 L_{2} -regression loss (for measuring the corresponding key-value mapping fitness) and minimize the loss function โฅ ๐ โ โ ( W i ; ๐ i ) โค โ ๐ท i โฅ 2 2 \|\bm{m}\nabla\mathcal{L}(W_{i};\bm{x}_{i})^{\top}-\bm{P}_{i}\|^{2}_{2} , resulting in the update rule of: W i + 1 = W i + ๐ฆ i + 1 , \displaystyle W_{i+1}=W_{i}+\mathbf{m}_{i+1}, (48) ๐ฆ i + 1 = ๐ฆ i ( ฮฑ i + 1 โ โ โ ( W i ; ๐ i ) โค โ โ ( W i ; ๐ i ) ) โ ฮท t ๐ท i โ โ ( W i ; ๐ i ) , \displaystyle\mathbf{m}_{i+1}=\mathbf{m}_{i}\;\left(\alpha_{i+1}-\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right)^{\top}\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right)\right)-\eta_{t}\bm{P}_{i}\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right), (49) This update is based on delta-rule ( Prados and Kak, 1989 ) and so it allows the memory (momentum) to better manage its limited capacity (i.e., ๐ช โก ( N ) \mathcal{O}(N) ) and better memorize the series of past gradients. For example, we may learn to forget some of the past gradients during the optimization process (similar to what happens when we move from linear attention to delta rule in associate memory). We refer to the variants of this type of momentum term as Delta Momentum variants.
Extension: More Expressive Memory. Viewing momentum as a compressor or a memory that store past gradients into its elements (parameters), its capacity not only depends on its update rule (similar to the above), but it also requires more expressive structure that allows for larger capacity. The current formulation is based on a linear layer (i.e., matrix-valued) to compress the past gradient values, but this linear nature may limit the capability to only learn linear mappings of past gradients. To increase the learning capacity of this module, one can use more complex mappings such as replacing a linear matrix-valued memory for momentum with an MLP. This design allows momentum to memorize more gradients and so provides better information for the optimization process. We extend the formulation in Equation 33 as: W i + 1 = W i + ๐ฆ i + 1 ( ๐ i ) , and ๐ฆ i + 1 = ฮฑ i + 1 ๐ฆ i โ ฮท t โ โ ( 2 ) ( ๐ i ; ๐ i , ๐ ) , W_{i+1}=W_{i}+\mathbf{m}_{i+1}\left(\bm{u}_{i}\right),\quad\textrm{and}\quad\mathbf{m}_{i+1}=\alpha_{i+1}\mathbf{m}_{i}-\eta_{t}\nabla\mathcal{L}^{(2)}(\bm{m}_{i};\bm{u}_{i},\mathds{1}), (50) where ๐ i = โ โ โ ( W i , ๐ i ) \bm{u}_{i}=\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right) and โ โ ( 2 ) โ ( โ
) \nabla\mathcal{L}^{(2)}(\cdot) is the internal objective of momentum (e.g., dot product similarity โจ ๐ โก ( ๐ i โค ) , ๐ โฉ \langle\bm{m}(\bm{u}_{i}^{\top}),\mathds{1}\rangle ). We refer to this variant as Deep Momentum Gradient Descent (DMGD). While it is clear from this example, it is worth emphasizing that the internal loss function and the model need to be designed carefully to obtain an effective momentum module.
Extension: Memory with Higher-order Feature Maps. One of the commonly used techniques to enhance the capacity of a memory is to use higher-order feature maps on the keys ( Kacham et al., 2024 ; Katharopoulos et al., 2020 ) . Using this technique on the momentum term, one can obtain: W i + 1 = W i + ๐ฆ i + 1 and ๐ฆ i + 1 = ฮฑ i + 1 โ ๐ฆ i โ ฮท t โ ๐ท i โ ฯ โ ( โ โ โ ( W i , ๐ i ) ) , W_{i+1}=W_{i}+\mathbf{m}_{i+1}\quad\textrm{and}\quad\mathbf{m}_{i+1}=\alpha_{i+1}\mathbf{m}_{i}-\eta_{t}\bm{P}_{i}\phi(\nabla\mathcal{L}\left(W_{i};\bm{x}_{i}\right)), (51) where ฯ โก ( โ
) \phi(\cdot) is a higher-order feature mapping (that may be learned through its internal objective).
Extension: Nonlinear Outputs. Building upon the associative memory perspective of the momentum, one common technique to enhance the representation power of memory module is to use non-linearity on top of its output ( Behrouz et al., 2025c ; Sun et al., 2024 ) . That is, we re-formulate Equation 50 as: W i + 1 = W i + ฯ ( ๐ฆ i + 1 ( ๐ i ) ) , and ๐ฆ i + 1 = ฮฑ i + 1 ๐ฆ i โ ฮท t โ โ ( 2 ) ( ๐ i ; ๐ i , ๐ ) , W_{i+1}=W_{i}+\sigma\left(\mathbf{m}_{i+1}\left(\bm{u}_{i}\right)\right),\quad\textrm{and}\quad\mathbf{m}_{i+1}=\alpha_{i+1}\mathbf{m}_{i}-\eta_{t}\nabla\mathcal{L}^{(2)}(\bm{m}_{i};\bm{u}_{i},\mathbf{I}), (52) where ฯ โก ( โ
) \sigma(\cdot) is an arbitrary non-linearity. As an example, we let ฯ โ ( โ
) = NewtonSchulz โ ( โ
) \sigma(\cdot)=\texttt{NewtonSchulz}(\cdot) , where Newton-Schulz โ ( โ
) \texttt{Newton-Schulz}(\cdot) is the iterative Newton-Schulz method ( Higham, 2008 ) , and ๐ โก ( โ
) \bm{m}(\cdot) be a linear layer; resulting in Muon ( Jordan et al., 2024 ) .
A Toy Example for Long Context in Optimizer. In Section 4.3 , we discussed that in complex setups, including continual learning with orthogonal tasks, we may need more complex momentum terms, either with higher capacity or better memory management. To better illustrate the potential gains of other momentum memory designs, we use a toy example of a time-varying curvature. Since the standard momentum acts as a low-pass filter, if the landscape changes at a high frequency, then standard momentum which aims to use the weighted average of past gradients, will be under the influence of irrelevant gradient terms, delaying the convergence. As an illustrative example, consider: ๐ โก ( r , ฮธ ) = r 2 + k ร ( r โ ฮธ + ฮฑ โ sin โก ( ฯ โ r ) ) 2 , \displaystyle\bm{\psi}(r,\theta)=r^{2}+k\times\left(r-\theta+\alpha\sin{\left(\omega r\right)}\right)^{2}, (53) and aims to optimize it using a standard momentum and our delta momentum. We start the optimization process from point ( r 0 , ฮธ 0 ) = ( โ 3.5 , 2 ) (r_{0},\theta_{0})=(-3.5,2) and continue until one of the algorithms converge to the optimal solution. The result is visualized in Figure 4 . The delta momentum finds the solution faster, mainly due to its gradient-dependent weight decay that helps the momentum term to decay or stop when it is needed.
### 4.5 Going Beyond Simple Gradient Descent and Momentum
Coming back to the discussion in Section 3.1 about the pre-training process and backpropagation being a form of associative memory, in this section, we aim to take advantage of NLโs viewpoint and present a more general form for gradient descent. As observed in Section 4.1 , backpropagation with gradient descent is an associative memory that aims to map the input data to the surprised caused by its predicted output โ y t โ โ ( W t , ๐ t ) \nabla_{y_{t}}\mathcal{L}(W_{t};\bm{x}_{t}) : W t + 1 = W t โ ฮท t + 1 โ โ W โ โ ( W t , ๐ t ) = W t โ ฮท t + 1 โ โ y โ โ ( W t , ๐ t ) โ ๐ t , where โ ๐ t โผ ๐ train , \displaystyle W_{t+1}=W_{t}-\eta_{t+1}\nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t})=W_{t}-\eta_{t+1}\nabla_{y}\mathcal{L}(W_{t};\bm{x}_{t})\otimes\bm{x}_{t},\qquad\text{where}\>\>\bm{x}_{t}\thicksim\mathcal{D}_{\text{train}}, (54) which from the associative memory perspective and proximal gradient viewpoint is equivalent to: W t + 1 = arg โก min W โ โจ W โ ๐ t , โ y t โ โ ( W t , ๐ t ) โฉ + 1 2 โ ฮท t โ โ W โ W t โ 2 2 . \displaystyle W_{t+1}=\arg\min_{W}\>\>\langle W\bm{x}_{t},\nabla_{y_{t}}\mathcal{L}(W_{t};\bm{x}_{t})\rangle+\frac{1}{2\eta_{t}}\>\|W-W_{t}\|^{2}_{2}. (55) This step aims at learning the negative of the gradient direction. The main drawback of the dot-product similarity as the inner objective is that its corresponding update rule and so learning algorithm treats each data sample (gradients) independent of the state, meaning that the state of the weights and so previous gradients do not affect the update term to the current state. While this design can be effective for nested problems with independent elements in their context (e.g., i.i.d. samples for training), it can be restrictive for context with highly dependent elements (e.g., tokens in a sequence). Defining ๐ฎ t = โ โ y t โ โ ( W t , ๐ t ) \mathbf{u}_{t}=-\nabla_{y_{t}}\mathcal{L}(W_{t};\bm{x}_{t}) , one can extend this process to more expressive objectives such as โ 2 \mathcal{L}_{2} regression loss: W t + 1 = arg โก min W โ 1 2 โ โ W โ ๐ t โ ๐ฎ t โ 2 2 + 1 2 โ ฮท t โ โ W โ W t โ 2 2 . \displaystyle W_{t+1}=\arg\min_{W}\>\>\frac{1}{2}\|W\bm{x}_{t}-\mathbf{u}_{t}\|^{2}_{2}+\frac{1}{2\eta_{t}}\|W-W_{t}\|^{2}_{2}. (56) For the cases that ๐ t \bm{x}_{t} is normalized (e.g. in normalized memory systems or in neural networks with normalization layers, โ ๐ t โ 2 = ฮป \|\bm{x}_{t}\|_{2}=\lambda ), and by defining ฮท t โฒ = ฮท t 1 + ฮท t \eta_{t}^{\prime}=\frac{\eta_{t}}{1+\eta_{t}} , we can use Sherman-Morrison lemma to get (see Appendix C for the details): W t + 1 \displaystyle W_{t+1} = W t โ ( ๐ โ ฮท t โฒ โ ๐ t โ ๐ t โค ) โ ฮท t โฒ โ โ W t โ โ ( W t , ๐ t ) \displaystyle=W_{t}\left(\mathbf{I}-\eta^{\prime}_{t}\bm{x}_{t}\bm{x}_{t}^{\top}\right)-\eta_{t}^{\prime}\nabla_{W_{t}}\mathcal{L}(W_{t};\bm{x}_{t}) = W t โ ( ๐ โ ฮท t โฒ โ ๐ t โ ๐ t โค ) โ ฮท t โฒ โ โ y t โ โ ( W t , ๐ t ) โ ๐ t , where โ ๐ t โผ ๐ train . \displaystyle=W_{t}\left(\mathbf{I}-\eta^{\prime}_{t}\bm{x}_{t}\bm{x}_{t}^{\top}\right)-\eta_{t}^{\prime}\nabla_{y_{t}}\mathcal{L}(W_{t};\bm{x}_{t})\otimes\bm{x}_{t},\qquad\text{where}\>\>\bm{x}_{t}\thicksim\mathcal{D}_{\text{train}}. (57) This new algorithm, which based on Delta rule ( Prados and Kak, 1989 ) we refer to as Delta Gradient Descent (DGD), updates the weights not only with respect to the current elements, but it also incorporates the previous state of weights, resulting in an adaptive decay term based on the current data sample. Next, we discuss a generalized viewpoint about the process of backpropagation with gradient descent, which later will help us to formulate Generalized Gradient Descent family of learning rules:
Based on the above interpretation, one can define backpropagation with gradient descent in a general form as any self-referential model that aims to compress training samples as keys and map them to self-generated values to better control its own learning process. Based on this definition, our above formulation of Appendix C , is only a simple instance that uses โ 2 \mathcal{L}_{2} regression loss; in general, however, one can define Generalized Gradient Descent (GGD) as:
###### Definition 5 (Generalized Gradient Descent (GGD) Learning Rule) .
Generalized Gradient Descent (GGD) learning rule is a self-referential associative memory that aims to compress data samples and map them to a set of self-generated keys: W t + 1 = arg โก min W โ โ ~ โ ( ๐ t , ๐ t ) + Ret โ ( W , { W i } i = t โ c + 1 t ) , \displaystyle W_{t+1}=\arg\min_{W}\>\>\tilde{\mathcal{L}}\left(\bm{x}_{t},\bm{u}_{t}\right)+\texttt{Ret}\left(W,\{W_{i}\}_{i=t-c+1}^{t}\right), (59) where ๐ฎ t \bm{u}_{t} is a self-generated value: ๐ t = ๐ W t โ ( ๐ t ) , \displaystyle\bm{u}_{t}=\bm{f}_{W_{t}}\left(\bm{x}_{t}\right), (60) for some function ๐ W t โ ( โ
) \bm{f}_{W_{t}}(\cdot) parameterized by W t W_{t} . Here, โ ~ โ ( โ
) \tilde{\mathcal{L}}\left(\cdot\right) measures the quality of the mapping, and Ret โ ( โ
) \>\texttt{Ret}(\cdot) ensures that the solution for the new instance is not far away from the current state.
Similarly, this formulation can be adapted for the momentum term, resulting in Generalized Momentum (GM). However, it is notable that the momentum itself, is a conventional associative memory and its keys and values are given, or more specifically are generated by a lower-frequency level. In Section 4.2 , we explored a special case of this formulation, where โ โก ( โ
) \mathcal{L}(\cdot) is L 2 L_{2} regression loss.
A Note on Optimizers in Continual Learning Setup. As discussed above, optimizers themselves are learning modules or associative memories that aim to compress the gradients into their parameters. These parameters are not necessarily trainable in the conventional terminology but indeed momentum-based optimizers store the knowledge about the loss landscape, helping them to better update the weights. When the โend of pretrainingโ happens for a neural learning module, the knowledge stored about the distribution of gradients/data, which are stored in the momentum term(s) is removed from the model and so continuing the training without recovering the momentum states can affect the modelโs ability to learn new capabilities. When the model is in continual learning setup, the knowledge about data is stored in the conventional parameters (optimized with backpropagation), but the knowledge about how the model optimizes itself and about the objective space are optimized in the lower-frequency levels of optimization (e.g., momentum terms).
## 5 Existing Architectures as Neural Learning Modules
Modern sequence models such as Transformers ( Vaswani et al., 2017 ) and recurrent models ( Katharopoulos et al., 2020 ; Schlag et al., 2021 ; Sun et al., 2024 ; Behrouz et al., 2025c ) are the backbones of recent advances in language models. Recently, the equivalency of such models with associative memories that aim to learn a mapping from keys to values from data have been studied in different settings and objectives ( Liu et al., 2024b ; Sun et al., 2024 ; Wang et al., 2025 ; Behrouz et al., 2025b ) . Particularly, we focus on the general framework of Miras ( Behrouz et al., 2025b ) , which defines associative memory as Definition 1 and optimizes the internal objective (called โattentional biasโ) with a choice of optimization algorithm on an arbitrary class of functions (i.e., memory architecture). While this formulation alone indicates that the well-known architectures are instances of nested systems of associative memory ( NSAM ), next, we review this equivalency for some learning rules and architectures.