Nested Learning: The Illusion of Deep Learning Architectures β part 3/14
Cleaned-text code points [31637, 47395). Bibliography omitted here; source.md includes it.
Associative memoryβthe ability to form and retrieve connections between eventsβis a fundamental mental process and is an inseparable component of human learning ( Terry, 2017 ) . Often in the literature, the concept of memorization and learning are used interchangeably; in neuropsychology literature, however, these two are clearly distinguished. More specifically, following neuropsychology literature ( Okano et al., 2000 ) , we build our terminology based on the following definition of memory and learning:
In this work, our goal is to first show that all the elements of a computational sequence model, including optimizers and neural networks, are associative memory systems that compress their own context flow . Broadly speaking, associative memory is an operator that maps a set of keys to a set of values. We follow the general definition of associative memory by Behrouz et al. (2025b) :
###### Definition 1 (Associative Memory) .
Given a set of keys π¦ β β d k \mathcal{K}\subseteq\mathbb{R}^{d_{k}} and values π± β β d v \mathcal{V}\subseteq\mathbb{R}^{d_{v}} , associative memory is an operator β³ β‘ ( β
) \mathcal{M}(\cdot) that maps the set of keys π¦ \mathcal{K} to values π± \mathcal{V} . To learn such mapping from the data, an objective β ~ β ( β
, β
) \tilde{\mathcal{L}}(\cdot;\cdot) measures the quality of the mapping and β³ \mathcal{M} can be computed by: β³ β = arg β‘ min β³ β ~ β ( β³ β‘ ( π¦ ) , π± ) . \displaystyle\mathcal{M}^{*}=\arg\min_{\mathcal{M}}\quad\tilde{\mathcal{L}}(\mathcal{M}(\mathcal{K});\mathcal{V}). (6)
While the operator itself is a memory and the mapping acts as a memorization process (i.e., memorizing the connections of events in the context), acquiring such effective operator based on the data, is a learning process. Notice that, here, keys and values can be any arbitrary event that memory aims to map them and are not limited to tokens. Later, we will discuss that given a context flow, keys and values might be tokens, gradients, sub-sequences, etc. Furthermore, while the term of associative memory is more common in neuroscience and neuropsychology literature, the above formulation is also closely related to data compression and low-dimensional representation. That is, one can interpret the optimization process in Equation 6 as the training process of a network β³ ( . ) \mathcal{M}(.) that aims to compress the mappings into its parameters, representing them in a lower dimensional space.
In sequence modeling, where keys and values are input tokens (e.g., tokenized text), the choice of objective and the optimization process for solving Equation 6 can result in distinct sequence modeling architectures (see Liu et al., 2024b and Behrouz et al., 2025b ) such as global/local softmax attention ( Vaswani et al., 2017 ) , or other modern recurrent models ( Katharopoulos et al., 2020 ; Sun et al., 2023 ; Behrouz et al., 2025c ) . This simple formulation of sequence models provides us with better understanding of their internal process and also a tool to simply compare their modeling power based on their objective and optimization process. In the following, using step-by-step examples, we discuss how this formulation can be applied to all components of a neural architecture (including its optimization process in pre-training) and in fact, how a model is an integrated system of multi-level, nested, potentially parallel memories, each of which with its own context flow.
A Simple Example of MLP Training. We start with a simple example, in which we aim to train a 1-layer MLP (parameterized with W W ) for task π― \mathcal{T} and on dataset π train = { x 1 , β¦ , π | π train | } \mathcal{D}_{\text{train}}=\{x_{1},\dots,\bm{x}_{|\mathcal{D}_{\text{train}}|}\} by optimizing the objective β β‘ ( β
, β
) \mathcal{L}(\cdot;\cdot) with gradient descent. In this case, the training process objective is to solve the following optimization problem: W β = arg β‘ min W β‘ β β‘ ( W , π train ) , \displaystyle W^{*}=\arg\min_{W}\>\>\mathcal{L}(W;\mathcal{D}_{\text{train}}), (7) whose optimization by (stochastic/online) gradient descent results in a weight update rule: W t + 1 = W t β Ξ· t + 1 β β W β β ( W t , π t + 1 ) β Surprise = W t β Ξ· t + 1 β β y t + 1 β β ( W t , π t + 1 ) β Surprise in the Output β π t + 1 , where β π t + 1 βΌ π train , \displaystyle W_{t+1}=W_{t}-\eta_{t+1}\underset{\text{Surprise}}{\underbrace{\nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t+1})}}=W_{t}-\eta_{t+1}\underset{\text{Surprise in the Output}}{\underbrace{\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1})}}\otimes\bm{x}_{t+1},\qquad\text{where}\>\>\bm{x}_{t+1}\thicksim\mathcal{D}_{\text{train}}, (8) where y t + 1 = W β x t + 1 y_{t+1}=Wx_{t+1} is the output of the model for input x t + 1 x_{t+1} and we used the simplifying notation: β y t + 1 β β ( W t , π t + 1 ) := β β β y | y = W β x t + 1 \nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}):=\frac{\partial\mathcal{L}}{\partial y}|_{y=Wx_{t+1}} . In this case, β W β β ( W t , π t + 1 ) \nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t+1}) is the surprise metric that shows how much the current input is different from previously observed data. Similarly, β y t + 1 β β ( W t , π t + 1 ) \nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}) is the surprised metric for the output (or more accurately local surprise signal in representation space that quantifies the mismatch between the current output and the structure the objective β β‘ ( β
, β
) \mathcal{L}(\cdot;\cdot) enforces)βmeasuring how much surprising the modelβs prediction is for this input. Given this formulation, one can let the surprise value of output be u t + 1 = β y t + 1 β β ( W t , π t + 1 ) u_{t+1}=\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}) and reformulate the backpropagation process as the solution to an optimization problem on finding an associative memory that maps input data points π train = { x t } t = 1 | π train | \mathcal{D}_{\text{train}}=\{x_{t}\}_{t=1}^{|\mathcal{D}_{\text{train}}|} to their corresponding u t + 1 = β y t + 1 β β ( W t , π t + 1 ) u_{t+1}=\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}) . That is, we let β³ ( β
) = W t β
\mathcal{M}(\cdot)=W_{t}\>\cdot parametrizes the memory, and use dot-product similarity to measure the quality of W t W_{t} βs mapping between x t + 1 x_{t+1} and β y t + 1 β β ( W t , π t + 1 ) \nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}) : W t + 1 = arg β‘ min W β β¨ W β π t + 1 , u t + 1 β© + 1 2 β Ξ· t + 1 β β W β W t β 2 2 = arg β‘ min W β β¨ W β π t , β y t + 1 β β ( W t , π t + 1 ) β© + 1 2 β Ξ· t + 1 β β W β W t β 2 2 . \displaystyle W_{t+1}=\arg\min_{W}\>\>\langle W\bm{x}_{t+1},u_{t+1}\rangle+\frac{1}{2\eta_{t+1}}\>\|W-W_{t}\|^{2}_{2}=\arg\min_{W}\>\>\langle W\bm{x}_{t},\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1})\rangle+\frac{1}{2\eta_{t+1}}\>\|W-W_{t}\|^{2}_{2}. (9) Therefore, this formulation translates the training phase of the model as a process of acquiring effective memory that maps data samples to their Local Surprise Signal (LSS) in representation spaceβmeasuring how surprising its corresponding output is. This gradient can be viewed as an error in the prediction (with gradient being zero when the loss is minimized). Later in Section 4 , we discuss the backpropagation process as an associative memory in more details, but as a preliminary takeaway from this simple example:
Accordingly, in this example, our model has a single gradient flow over the data samples, which is only active over dataset π train = { x 1 , β¦ , π | π train | } \mathcal{D}_{\text{train}}=\{x_{1},\dots,\bm{x}_{|\mathcal{D}_{\text{train}}|}\} and will be frozen for any other data samples afterwards (i.e., inference or test time).
In the above example, we can replace the gradient descent algorithm with its momentum-based variant, resulting in the update rule of: W t + 1 \displaystyle W_{t+1} = W t β π t + 1 , \displaystyle=W_{t}-\bm{m}_{t+1}, (10) π t + 1 \displaystyle\bm{m}_{t+1} = π t + Ξ· t + 1 β β W β β ( W t , π t + 1 ) = π t + Ξ· t + 1 β β y t + 1 β β ( W t , π t + 1 ) β π t + 1 . \displaystyle=\bm{m}_{t}+\eta_{t+1}\nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t+1})=\bm{m}_{t}+\eta_{t+1}\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1})\otimes\bm{x}_{t+1}. (11) In Equation 11 , given the previous state of Equation 10 (at time t t ), the value of β W β β ( W t , π t + 1 ) \nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t+1}) or similarly β y t + 1 β β ( W t , π t + 1 ) \nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1}) does not depend on the output of recurrence in Equation 11 and so can be pre-computed beforehand: Leting u t + 1 = β W β β ( W t , π t + 1 ) u_{t+1}=\nabla_{W}\mathcal{L}(W_{t};\bm{x}_{t+1}) , Equation 11 can be reformulated as: W t + 1 \displaystyle W_{t+1} = W t β π t + 1 , \displaystyle=W_{t}-\bm{m}_{t+1}, (12) π t + 1 \displaystyle\bm{m}_{t+1} = arg β‘ min π β β¨ π , β W t β β ( W t , π t + 1 ) β© + 1 2 β Ξ· t + 1 β β π β π t β 2 2 = arg β‘ min π β β¨ π β π t + 1 , β y t + 1 β β ( W t , π t + 1 ) β© + 1 2 β Ξ· t + 1 β β π β π t β 2 2 . \displaystyle=\arg\min_{\bm{m}}\>\>-\langle\bm{m},\nabla_{W_{t}}\mathcal{L}(W_{t};\bm{x}_{t+1})\rangle+\frac{1}{2\eta_{t+1}}\>\|\bm{m}-\bm{m}_{t}\|^{2}_{2}=\arg\min_{\bm{m}}\>\>-\langle\bm{m}\>\bm{x}_{t+1},\nabla_{y_{t+1}}\mathcal{L}(W_{t};\bm{x}_{t+1})\rangle+\frac{1}{2\eta_{t+1}}\>\|\bm{m}-\bm{m}_{t}\|^{2}_{2}. (13)
Given this formulation, one can interpret the momentum term as either: (1) a value-less associative memory that compress the gradients into its parameters, or (2) an associative memory that learns how to map data points to their corresponding LSS-value. Interestingly, this formulation reveals that gradient descent with momentum can be viewed as a two-level optimization procedure, where the memory is optimized by simple gradient descent algorithm -1 -1 -1 We use the term two- or multi-level to describe the optimization procedure . This differs from classical multi-level optimization, where the optimization problems are arranged hierarchically. .
Concluding the above examples, we observed that the training process of a 1-layer MLP with: (1) Gradient descent is a 1-level associative memory that learns how to map data points to their corresponding LSS-value; and (2) Gradient descent with momentum is a 2-level associative memory (or optimization process) that the inner-level learns to store gradient values into its parameters, and then the outer-level updates the slow weight (i.e., W t W_{t} ) with the value of the inner-level memory. While these are the most simple examples with respect to both architecture and optimizer algorithms, one might ask if similar conclusion can be made in more complex setups.
An Example of Architectural Decomposition. In the next example, we replace the MLP module in our previous example with a linear attention ( Katharopoulos et al., 2020 ) . That is, we aim to train a 1-layer linear attention for task π― \mathcal{T} and on a sequence of π train = { x 1 , β¦ , π | π train | } \mathcal{D}_{\text{train}}=\{x_{1},\dots,\bm{x}_{|\mathcal{D}_{\text{train}}|}\} by optimizing the objective β \mathcal{L} with gradient descent. Recalling the unnormalized linear attention formulation: π t = W π β π t , π t = W π β π t , π t = W π β π t , \displaystyle\bm{k}_{t}=W_{\bm{k}}\bm{x}_{t},\qquad\bm{v}_{t}=W_{\bm{v}}\bm{x}_{t},\qquad\bm{q}_{t}=W_{\bm{q}}\bm{x}_{t}, (14) β³ t = β³ t β 1 + π t β π t β€ , \displaystyle\mathcal{M}_{t}=\mathcal{M}_{t-1}+\bm{v}_{t}\bm{k}_{t}^{\top}, (15) y t = β³ t β π t . \displaystyle y_{t}=\mathcal{M}_{t}\bm{q}_{t}\>. (16) As discussed in earlier studies ( Liu et al., 2024b ; Behrouz et al., 2025b ) , the recurrence in Equation 15 can be reformulated as the optimization process of a matrix-valued associative memory β³ t β ( β
) \mathcal{M}_{t}(\cdot) to compress the mappings of keys and values into its parameters. Specifically, in Definition 1 , if we let β ~ β ( β³ t β 1 , π t , π t ) := β β¨ β³ t β 1 β π t , π t β© \tilde{\mathcal{L}}(\mathcal{M}_{t-1};\bm{k}_{t},\bm{v}_{t}):=-\langle\mathcal{M}_{t-1}\bm{k}_{t},\bm{v}_{t}\rangle and aim to optimize the memory with gradient descent, the memory update rule is: (Note that β β ~ β ( β³ t β 1 , π t , π t ) = β π t β π t β€ \nabla\tilde{\mathcal{L}}(\mathcal{M}_{t-1};\bm{k}_{t},\bm{v}_{t})=-\bm{v}_{t}\bm{k}_{t}^{\top} and we let learning rate Ξ· t = 1 \eta_{t}=1 ) β³ t + 1 = arg β‘ min β³ β β¨ β³ β π t + 1 , π t + 1 β© + 1 2 β β β³ β β³ t β 2 2 \displaystyle\mathcal{M}_{t+1}=\arg\min_{\mathcal{M}}\>\>-\langle\mathcal{M}\bm{k}_{t+1},\bm{v}_{t+1}\rangle+\frac{1}{2}\|\mathcal{M}-\mathcal{M}_{t}\|^{2}_{2}\quad (17) β \displaystyle\Rightarrow\>\> β³ t + 1 = β³ t β β β ~ β ( β³ t , π t + 1 , π t + 1 ) = β³ t + π t + 1 β π t + 1 β€ , \displaystyle\mathcal{M}_{t+1}=\mathcal{M}_{t}-\nabla\tilde{\mathcal{L}}(\mathcal{M}_{t};\bm{k}_{t+1},\bm{v}_{t+1})=\mathcal{M}_{t}+\bm{v}_{t+1}\bm{k}_{t+1}^{\top}, (18) which is equivalent to the update rule of an unnormalized linear attention in Equation 15 . Also, as we observed in the first example, training a linear layer with gradient descent can be viewed as a 1-level optimization of an associative memory ( Equation 8 ) and so the general training/updating process of projection layers (i.e., W π , W π , W_{\bm{k}},W_{\bm{v}}, and W π W_{\bm{q}} ) is itself an optimization process of associative memory. Therefore, training a linear attention with gradient descent can be seen as a two-level optimization process, where the outer-loop (also known as training process) optimizes the projection layers with gradient descent, while the inner-loop optimizes the inner memory of β³ t \mathcal{M}_{t} with gradient descent.
In the examples we discussed so far, we have two associative memories, each of which has their own optimization process and gradient flow. That is, in the optimization of the outer-level parameters of W π , W π , W_{\bm{k}},W_{\bm{v}}, and W π W_{\bm{q}} there is no gradient with respect to parameter β³ β‘ ( β
) \mathcal{M}(\cdot) and so there is no backpropagation through it. Similarly, in the inner-level, there is no backpropagation through projection layers and they are considered frozen. Furthermore, it is notable that in this example, the above formulation is also closely connected to FWPs perspective of linear attentions ( Schlag et al., 2021 ) , where projections are considered slow weights, and memory update in Equation 15 is the fast weight update rule.
Architectural Decomposition with More Levels. In both examples above, we discussed how they can be viewed as a 2-level optimization process (coinciding with their FWPs interpretations). In practice, however, we may need to use more powerful optimization process, and/or more powerful recurrent update rules for memory. As a simple example, assume we use gradient descent with momentum to train a linear attention model. As we saw above, the linear attention component can be decomposed into two nested optimization processes. Similarly, the model here can be represented as a 2-level optimization problem, where (1) the inner level optimizes the memory to compress the context using gradient descent ( Equation 17 ), and (2) the outer level optimizes the projection layers with gradient descent with momentum. Interestingly, we saw that βgradient descent with momentumβ algorithm itself can be viewed as a 2-level optimization process where the momentum term itself is an associative memory that compress the past gradients into its parameters.
### 3.2 Nested Optimization Processes
In the previous section, we provided examples to demonstrate how one can decompose a machine learning model into a set of nested or multi-level optimization procedures. Next, we first present a formalization of nested learning problems and then define Neural Learning Moduleβan integrated computational system that learns from data.