Nested Learning: The Illusion of Deep Learning Architectures — part 4/14

Cleaned-text code points [47395, 62819). Bibliography omitted here; source.md includes it.

In previous sections, we decomposed the model into a set of optimization process. However, it is still unclear if we can define a hierarchy (or order) over these processes, and uniquely represent the model in this format. Inspired by the hierarchy of brain waves that indicates the information processing frequency rate of each part (discussed in Section 1 ), we use the update rate of each optimization process to order the components in multiple levels. To this end, we let the one update step over one data point to be the unit of time, and define the update frequency rate of each component as:

###### Definition 2 (Update Frequency) .

For any component of A A , which can be a parametric component (e.g., learnable weights or momentum term in gradient descent with momentum) or a non-parametric component (e.g., attention block), we define its frequency, denoted by f A f_{A} , as its number of updates per unit of time.

Given the above update frequency, we can order the components of a machine learning algorithm based on operator ( ⋅ ≻ ⋅ ) (\cdot\succ\cdot) . We say A A is faster than B B and denote A ≻ B A\succ B if: (1) f A > f B f_{A}>f_{B} , or (2) f A = f B f_{A}=f_{B} but the computation of the B B ’s state at time t t requires the computation of A A ’s state at time t t . In this definition, when A ⊁ B A\nsucc B and B ⊁ A B\nsucc A , we let A ​ = 𝑓 ​ B A\overset{f}{=}B , which indicates that A A and B B has the same frequency update, but their computation is independent of each other (Later, we provide an example of this cases in AdamW optimizer). Based on the above operator, we sort the components into an ordered set of “ levels ”, where (1) components in the same level have the same frequency update, and (2) the higher the level is, the lower its frequency. Given the above formulations of levels and update frequency, we next formally define nested learning:

###### Definition 3 (Nested System) .

A (ordered) nested system is a system with K K (ordered) levels such that each level k k , 1 ≤ k ≤ K 1\leq k\leq K , consists of a set of optimization problems { ( ℒ i ( k ) , 𝒞 i ( k ) , 𝚯 i ( k ) ) } i = 1 N k \{(\mathcal{L}^{(k)}_{i},\mathcal{C}^{(k)}_{i},\bm{\Theta}^{(k)}_{i})\}_{i=1}^{N_{k}} , where ℒ i ​ ( ⋅ , ⋅ ) \mathcal{L}_{i}(\cdot;\cdot) is the optimization objective in the i i -th problem, 𝒞 i \mathcal{C}_{i} is its context (the data that is optimized on), 𝚯 i \bm{\Theta}_{i} is the feasible set of its parameters, and each parameter is optimized using gradient descent: 𝜽 i t + 1 ( k ) = arg ⁡ min 𝚽 i ( k ) ​ ⟨ 𝚽 i ( k ) ​ 𝒙 t + 1 , − ∇ ℒ i ( k ) ​ ( 𝜽 i t ( k ) , 𝒙 t + 1 ) ⟩ + 1 2 ​ η i t + 1 ( k ) ​ ‖ 𝚽 i ( k ) − 𝜽 i t ( k ) ‖ 2 2 where ​ 𝒙 t + 1 ∼ 𝒞 i ( k ) , and ​ 𝚽 i ( k ) ∈ 𝚯 i ( k ) . \displaystyle{\bm{\theta}_{i}}^{(k)}_{t+1}=\arg\min_{\bm{\Phi}^{(k)}_{i}}\>\langle\bm{\Phi}^{(k)}_{i}\bm{x}_{t+1},-\nabla\mathcal{L}^{(k)}_{i}(\bm{\theta}^{(k)}_{i_{t}};\bm{x}_{t+1})\rangle+\frac{1}{2{\eta_{i}}^{(k)}_{t+1}}\>\|\bm{\Phi}^{(k)}_{i}-{\bm{\theta}_{i}}^{(k)}_{t}\|^{2}_{2}\quad\text{where}\>\>\bm{x}_{t+1}\thicksim\mathcal{C}^{(k)}_{i},\>\>\>\text{and}\>\>\>{\bm{\Phi}_{i}}^{(k)}\in\bm{\Theta}^{(k)}_{i}\!\!. (19)

Notice that each optimization process has its own gradient flow, and therefore sometimes we refer to them as a box of gradient flow corresponding to an optimization problem. Through this paper, we further generalize our definition of nested systems, and allow finding non-parametric solutions for some boxes (i.e., optimization problems).

The above definition provides a general flexible definition for a nested system that does not specify if there is any dependence among different boxes (i.e., a box can determines the context or the parameter space of another box). In the next sections, we discuss how knowledge/information can be transferred between different levels or boxes. Throughout this paper, we focus on the Nested Systems of Associative Memories ( NSAM ) that is a nested system, in which each optimization process is an associative-memory. More formally,

###### Definition 4 (Nested System of Associative Memories) .

A nested system of associative memory ( NSAM ) is a system with K K (ordered) levels such that each level k k , 1 ≤ k ≤ K 1\leq k\leq K , consists of a set of optimization problems { ( ℒ i ( k ) , 𝒞 i ( k ) , 𝚯 i ( k ) ) } i = 1 N k \{(\mathcal{L}^{(k)}_{i},\mathcal{C}^{(k)}_{i},\bm{\Theta}^{(k)}_{i})\}_{i=1}^{N_{k}} , where 𝒞 i = { ( 𝐤 j ( i ) , 𝐯 j ( i ) ) } j = 1 L i \mathcal{C}_{i}=\{(\bm{k}^{(i)}_{j},\bm{v}^{(i)}_{j})\}_{j=1}^{L_{i}} is a set of key-value pairs, ℒ i ​ ( ⋅ , ⋅ , ⋅ ) \mathcal{L}_{i}(\cdot;\cdot,\cdot) measures the quality of memory learned mappings in the i i -th problem, 𝚯 i \bm{\Theta}_{i} is the set of feasible memory parameters with each parameter optimized using gradient descent: 𝜽 i t + 1 ( k ) = arg ⁡ min 𝚽 i ( k ) ​ ⟨ 𝚽 𝒊 ( 𝒌 ) ​ 𝒌 t + 1 ( i ) , − ∇ ℒ i ( k ) ​ ( 𝜽 i t ( k ) , 𝒌 t + 1 ( i ) , 𝒗 t + 1 ( i ) ) ⟩ + 1 2 ​ η i t + 1 ( k ) ​ ‖ 𝚽 i ( k ) − 𝜽 i t ( k ) ‖ 2 2 , \displaystyle{\bm{\theta}_{i}}^{(k)}_{t+1}=\arg\min_{\bm{\Phi}^{(k)}_{i}}\>\>\langle\bm{\Phi_{i}^{(k)}}\>\bm{k}^{(i)}_{t+1},-\nabla\mathcal{L}^{(k)}_{i}(\bm{\theta}^{(k)}_{i_{t}};\bm{k}^{(i)}_{t+1},\bm{v}^{(i)}_{t+1})\rangle+\frac{1}{2{\eta_{i}}^{(k)}_{t+1}}\>\|\bm{\Phi}^{(k)}_{i}-{\bm{\theta}_{i}}^{(k)}_{t}\|^{2}_{2}, (20) where ( 𝐤 t + 1 ( i ) , 𝐯 t + 1 ( i ) ) ∼ 𝒞 i ( k ) (\bm{k}^{(i)}_{t+1},\bm{v}^{(i)}_{t+1})\thicksim\mathcal{C}^{(k)}_{i} and 𝚽 i ( k ) ∈ 𝚯 i ( k ) {\bm{\Phi}_{i}}^{(k)}\in\bm{\Theta}^{(k)}_{i} .

Given a query 𝒒 \bm{q} , for each associative memory ℳ i ( k ) \mathcal{M}_{i}^{(k)} , we use ℳ i ( k ) ​ ( 𝒒 ) \mathcal{M}_{i}^{(k)}(\bm{q}) to refer to the forward pass process (i.e., retrieval process) of the memory. While our formulation of NSAM can simply be defined by any optimization process beyond gradient descent, and initially, it may seem that the strict condition of optimization by gradient descent can limit the modeling power of the definition, through this paper, we show that modern architectures alongside certain well-known optimization algorithms can be viewed as instances of NSAM . We then build on top of this intuition and discuss how to go further with stacking multiple levels and design models with enhanced continual learning capabilities.

The new dimension of stacking multiple levels is an important characteristic of nested learning, where the depth of computation can be enhanced with increasing the number of levels. Depending on the design, context, and the type of knowledge transfer, this depth of computation can itself be viewed as different concepts such as: higher-order in-context learning ability, latent computation (e.g., Loop Transformers), multiple memory systems, and more expressive optimizers. Later, we discuss all such implications, but next, we use a simple example that connects MLP layers in Transformer architectures with linear or deep memory blocks.

Example of an MLP Layer vs. Linear Attention. Let us compare two models: (1) a Transformer architecture, and (2) the same backbone but with replacing the MLP block with a linear attention mechanism (sharing the keys and values from previous layer), in which the initial state of its memory is meta-learned (similar to Behrouz et al. (2025c) or Sun et al. (2024) ). We refer to the second variant as Adaptive Transformer or AdaTransformer. Both models are optimized on Next Token Prediction (NTP) objective with gradient descent. As discussed in the initial examples, and also illustrated in Figure 3 , both models have two levels, and for the sake of clarity, we use red (resp. blue ) to highlight computations/weight in the first level (resp. second level). More formally, let X = { x i } i = 1 T X=\{x_{i}\}_{i=1}^{T} be an input sequence of tokens, the output of both blocks are computed as (For the sake of simplicity, we assume MLP ( ⋅ ) = ⋅ W MLP \texttt{MLP}(\cdot)=\cdot\>W_{\texttt{MLP}}\> , and remove normalizations):

𝒌 t = 𝒙 t ​ W 𝒌 , 𝒗 t = 𝒙 t ​ W 𝒗 , 𝒒 t = 𝒙 t ​ W 𝒒 , \displaystyle\bm{k}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{k}}},\qquad\bm{v}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{v}}},\qquad\bm{q}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{q}}}, 𝒚 attn = Attn ​ ( 𝒌 , 𝒗 , 𝒒 ) , \displaystyle\bm{y}_{\texttt{attn}}={\color[rgb]{0.1641,0.2891,0.4023}\texttt{Attn}}\left(\bm{k},\bm{v},\bm{q}\right), 𝒚 block = MLP ​ ( 𝒚 attn ) = 𝒚 attn ​ W MLP , \displaystyle\bm{y}_{\texttt{block}}={\color[rgb]{0.6094,0.3242,0.3242}\texttt{MLP}}\left(\bm{y}_{\texttt{attn}}\right)=\bm{y}_{\texttt{attn}}\>{\color[rgb]{0.6094,0.3242,0.3242}W_{\texttt{MLP}}}, (Transformer Block)

𝒌 t = 𝒙 t ​ W 𝒌 , 𝒗 t = 𝒙 t ​ W 𝒗 , 𝒒 t = 𝒙 t ​ W 𝒒 , \displaystyle\bm{k}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{k}}},\qquad\bm{v}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{v}}},\qquad\bm{q}_{t}=\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{\bm{q}}}, 𝒚 attn = Attn ​ ( 𝒌 , 𝒗 , 𝒒 ) , \displaystyle\bm{y}_{\texttt{attn}}={\color[rgb]{0.1641,0.2891,0.4023}\texttt{Attn}}\left(\bm{k},\bm{v},\bm{q}\right), 𝒚 block = 𝒚 attn ​ W LinAttn . \displaystyle\bm{y}_{\texttt{block}}=\bm{y}_{\texttt{attn}}\>{\color[rgb]{0.1641,0.2891,0.4023}W_{\texttt{LinAttn}}}. (AdaTransformer Block)

The formulation of both blocks seem to be very similar and the only difference comes from the level of W MLP W_{\texttt{MLP}} and W LinAttn W_{\texttt{LinAttn}} weights. That is, while W MLP W_{\texttt{MLP}} is in the first level and so is persistent with respect to the context, W LinAttn W_{\texttt{LinAttn}} is adaptive and is updated in-context by ( ℳ ⁡ ( ⋅ ) \mathcal{M}(\cdot) is parametrized by W LinAttn {W_{\texttt{LinAttn}}} ): ℳ t = ℳ t − 1 + 𝒗 t ​ 𝒌 t ⊤ . \displaystyle{\color[rgb]{0.1641,0.2891,0.4023}\mathcal{M}_{t}=\mathcal{M}_{t-1}+\bm{v}_{t}\bm{k}^{\top}_{t}}. (21) In earlier variants of linear attention, the initial state of ℳ ⁡ ( ⋅ ) \mathcal{M}(\cdot) or equivalently W LinAttn {W_{\texttt{LinAttn}}} is considered as zero matrix, ℳ 0 = 𝟎 \mathcal{M}_{0}=\mathbf{0} . Similar to more advanced design choices ( Sun et al., 2024 ; Behrouz et al., 2025c ) , however, this initial state can be meta-learned to adapt fast to a context. In this setting, the initial state of ℳ 0 = W LinAttn init {\color[rgb]{0.6094,0.3242,0.3242}\mathcal{M}_{0}=W_{\texttt{LinAttn}_{\texttt{init}}}} is optimized in the first level with NTP objective, while given the context, W LinAttn {\color[rgb]{0.1641,0.2891,0.4023}W_{\texttt{LinAttn}}} is optimized in the second level as an associative memory with dot-product objective ( Behrouz et al., 2025b ) .

The above example is also valid when using more advanced and deep MLP blocks in Transformer architecture (such as SwiGLU ( Shazeer, 2020 ) ) and compared it with its recurrent memory counterpart ( Behrouz et al., 2025a ) . Furthermore, this simple example implies that the current perspective on hybrid architectures as the combination of expressive softmax attention and efficient recurrent models is somewhat misleading and it follows the conventional Transformer backbone design but with additional in-context learning capabilities for MLP blocks. We discuss this further in Section 6 and Section 7 . As a takeaway of the discussion in this subsection about the concept of nested systems and nested learning:

As discussed earlier, it is common in the literature to separate architectures from their optimization processes and to treat them as independent design choices, with the aim of combining algorithms that achieve the greatest expressive power in each aspect. In practice, however, a Transformer architecture ( Vaswani et al., 2017 ) that is optimized with stochastic gradient descent can learn a very different solution than the same architecture when Adam optimizer is used ( Kingma and Ba, 2014 ) . Accordingly, when interacting with such machine learning algorithms, we observe that despite the similarity in the architecture axes, the overall trained models show different predictions or generates different outputs. From the NL’s viewpoint, however, a machine learning algorithm is represented as an interconnected system of optimization problems and model’s actions, predictions, and generation of output depends on this system as a whole, not necessarily each of its sub-components. To this end, we define the term of neural learning module to refer to this representation of a model, where architecture and optimization process jointly determine the model and its outputs. While such joint representation might not seem significant in the current machine learning pipelines, where there is a training and then test phases, it becomes more important in the continual setup we advocate for, where there is no training/test phases (see more discussions in Section 8 ).

Neural Learning Modules are Inter-connected Systems. Based on the definition of neural learning module, one important question is how the architecture and optimization process are interconnected systems and how they can affect each other. Recall the general formulation for training a neural network: Given a task 𝒯 \mathcal{T} , its corresponding data distribution p ⁡ ( 𝒯 ) p(\mathcal{T}) , a model f ⁡ ( ⋅ , ⋅ ) f(\cdot;\cdot) parameterized by Φ 𝒯 \Phi_{\mathcal{T}} , and an objective ℒ ⁡ ( ⋅ , ⋅ ) \mathcal{L}(\cdot;\cdot) we aim to learn parameters Φ 𝒯 ∗ \Phi^{*}_{\mathcal{T}} such that: Φ 𝒯 ∗ = arg ⁡ min Φ ​ 𝔼 𝒙 , 𝒚 ∼ p ⁡ ( 𝒯 ) ​ [ ℒ ⁡ ( Φ , 𝒙 , 𝒚 ) ] . \displaystyle\Phi^{*}_{\mathcal{T}}=\arg\min_{\Phi}\>\>\mathbb{E}_{\>\bm{x},\bm{y}\sim p(\mathcal{T})}\left[\mathcal{L}\left(\Phi;\bm{x},\bm{y}\right)\right]. (22) In practice, we optimize the above problem based on a given dataset 𝒟 train \mathcal{D}_{\text{train}} and an optimization algorithm such as stochastic gradient descent: Φ t + 1 = Φ t − η t + 1 ​ ∇ Φ t ℒ ​ ( Φ t , 𝒙 t + 1 , 𝒚 t + 1 ) , where ​ ( 𝒙 t + 1 , 𝒚 t + 1 ) ∼ 𝒟 train . \displaystyle\Phi_{t+1}=\Phi_{t}-\eta_{t+1}\nabla_{\Phi_{t}}\mathcal{L}(\Phi_{t};\bm{x}_{t+1},\bm{y}_{t+1}),\qquad\text{where}\>\>(\bm{x}_{t+1},\bm{y}_{t+1})\thicksim\mathcal{D}_{\text{train}}. (23) One interpretation for the optimization process of model f ⁡ ( ⋅ , ⋅ ) f(\cdot;\cdot) in Equation 23 is to see the model as the data generator for the optimization process in Equation 23 . That is, as discussed in the first example in Section 3.1 , and later we will show in Section 4 , the optimization process is an associative memory that aims to compress the patterns between the training data and its gradients (or surprise) and so the dataset for internally training such memory (i.e., the gradients of the model) is generated by the model. Therefore, the type of the model can result in generating dataset (i.e., gradients) with different patterns and distributions over time. The effect of the optimization process and this data generation also fed back at the model itself, where the next state of the parameters in the model are determined by the optimization algorithms. As we will discuss in Section 4 , looking at optimizers as associative memories on the gradients of the model implies that each optimizer has some special traits such as better memory management, higher compression, etc. Therefore, the choice of such algorithms requires understanding the generated gradients and also changes of the model in the parameter space.

### 3.3 Knowledge Transfer Between Levels