Nested Learning: The Illusion of Deep Learning Architectures β€” part 5/14

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

So far, we mainly focused on the concept of nested learning and how optimization problems are located in different levels. It is, however, still unclear that how nested optimization problems (in different levels) can affect each other, or in general how they can contribute to the output of the system, and so be interconnected. In this section, we discuss several potential knowledge transfer methods between components in different levels. For the sake of clarity, we discuss the knowledge transfer between two levels and blocks, i.e., ℬ ( 0 ) = ( β„’ ( 0 ) , π’ž ( 0 ) , 𝚯 ( 0 ) ) \mathcal{B}^{(0)}=(\mathcal{L}^{(0)},\mathcal{C}^{(0)},\bm{\Theta}^{(0)}) and ℬ ( 1 ) = ( β„’ ( 1 ) , π’ž ( 1 ) , 𝚯 ( 1 ) ) \mathcal{B}^{(1)}=(\mathcal{L}^{(1)},\mathcal{C}^{(1)},\bm{\Theta}^{(1)}) with corresponding memories β„³ ( 0 ) ​ ( β‹… ) \mathcal{M}^{(0)}(\cdot) and β„³ ( 1 ) ​ ( β‹… ) \mathcal{M}^{(1)}(\cdot) , respectively:

Direct Connection of Levels (Parametric). The first type of knowledge transfer is to directly incorporate the weights in different levels or blocks. To this end, the forward pass or the retrieval process from the lower-frequency (i.e., higher-level) memory system is also conditioned on the parameters of the higher-frequency (i.e., lower-level) memory: β„³ ( 0 ) ​ ( β‹… ) := β„³ ( 0 ) ​ ( β‹… , 𝚯 ( 1 ) ) . \displaystyle\mathcal{M}^{(0)}(\cdot):=\mathcal{M}^{(0)}(\cdot\>;\>\bm{\Theta}^{(1)}). (24) In a more specific formulation, and as a special variant of the above formulation, we can condition the output of β„³ ( 0 ) ​ ( β‹… ) \mathcal{M}^{(0)}(\cdot) based on the output (or forward pass) of the higher-frequency memory: β„³ ( 0 ) ​ ( β‹… ) := β„³ ( 0 ) ​ ( β‹… , β„³ ( 1 ) ​ ( β‹… ) ) , \displaystyle\mathcal{M}^{(0)}(\cdot):=\mathcal{M}^{(0)}(\cdot\>;\>\mathcal{M}^{(1)}(\cdot)), (25) where we slightly abused notation by hiding the dependence on the second argument. As an example of this type of knowledge transfer, in linear Transformer (or FWP) ( Katharopoulos et al., 2020 ; Schlag et al., 2021 ) , where the initial memory state is zero, the stored knowledge of a lower level (fast weight) directly affect the output of the model in another level. That is, one can re-write the forward pass (memory retrieval) as: π’š t = β„³ t ​ 𝒒 t = β„³ t ​ [ 𝒙 t ​ W q ] ⏟ Forward pass of the lower-frequency memory ⏞ Forward pass of the higher-frequency memory \displaystyle\bm{y}_{t}={\color[rgb]{0.1641,0.2891,0.4023}\mathcal{M}_{t}}\bm{q}_{t}=\overset{\underset{\text{\small{\color[rgb]{0.1641,0.2891,0.4023}higher-frequency} memory}}{\text{\small Forward pass of the}}}{\overbrace{{\color[rgb]{0.1641,0.2891,0.4023}\mathcal{M}_{t}}\underset{\underset{\text{\small{\color[rgb]{0.6094,0.3242,0.3242}lower-frequency} memory}}{\text{\small Forward pass of the}}}{\underbrace{\left[\bm{x}_{t}{\color[rgb]{0.6094,0.3242,0.3242}W_{q}}\right]}}{}}} (26)

Direct Connection of Levels (Non-Parametric). Another form of direct connection between levels is a non-parametric variant of the above formulation, where the block ℬ ( 1 ) \mathcal{B}^{(1)} is optimized by finding non-parametric solution. Therefore, the forward pass of the low-frequency memory is conditioned on the context and output of higher-frequency memory: β„³ ( 0 ) ​ ( β‹… ) := β„³ ( 0 ) ​ ( β‹… , π’ž ( 1 ) ) , or similarly, β„³ ( 0 ) ​ ( β‹… ) := β„³ ( 0 ) ​ ( β‹… , β„³ ( 1 ) ​ ( β‹… , π’ž ( 1 ) ) ) . \displaystyle\mathcal{M}^{(0)}(\cdot):=\mathcal{M}^{(0)}(\cdot\>;\>\mathcal{C}^{(1)}),\qquad\text{or similarly, }\qquad\mathcal{M}^{(0)}(\cdot):=\mathcal{M}^{(0)}(\cdot\>;\>\mathcal{M}^{(1)}(\cdot;\mathcal{C}^{(1)})). (27) As an example of this variant, one can refer to Transformers and softmax attention module ( Vaswani et al., 2017 ) . There is an important characteristic for both above variants: there is no backpropagation through any state of the blocks in two different levels and knowledge transfers through direct conditioning the output of one level on the other’s output/parameters. Therefore, in this process the state of each block is treated as hyperparameter for the other.

Knowledge Transfer via Backpropagation. Another form of knowledge transfer is through backpropagation, where there is a gradient flow between blocks in different levels. The forward pass of this design is the same as the forward pass discussed above. The backward pass, however, is the main difference, where in the above two cases the state of each associative memory is considered as hyperparameters of the other but here both states are optimized in the same gradient flow. Therefore, for a simple case of two blocks in two levels, we have: β„³ ( 0 ) ​ ( β‹… ) := β„³ ( 0 ) ​ ( β‹… , β„³ ( 1 ) ​ ( β‹… ) ) \displaystyle\>\>\>\>\mathcal{M}^{(0)}(\cdot):=\mathcal{M}^{(0)}(\cdot\>;\>\mathcal{M}^{(1)}(\cdot)) (Forward Pass) { 𝚯 t + 1 ( 1 ) = 𝚯 t ( 1 ) βˆ’ Ξ· t + 1 ( 1 ) ​ 𝜹 1 ​ 𝒙 ^ t + 1 ⊀ , 𝚯 t + 1 ( 0 ) = 𝚯 t ( 0 ) βˆ’ Ξ· t + 1 ( 0 ) ​ 𝜹 0 ​ 𝒙 t + 1 ⊀ , \displaystyle\begin{cases}\>\bm{\Theta}^{(1)}_{t+1}=\bm{\Theta}^{(1)}_{t}-\eta^{(1)}_{t+1}\>\bm{\delta}_{1}\>\hat{\bm{x}}_{t+1}^{\top}\>,\\ \>\bm{\Theta}^{(0)}_{t+1}=\bm{\Theta}^{(0)}_{t}-\eta^{(0)}_{t+1}\>\bm{\delta}_{0}\>\bm{x}_{t+1}^{\top}\>,\end{cases} (Backward Pass) where 𝜹 t + 1 ( 0 ) = 𝑱 Ο• ( 0 ) ​ ( 𝒙 t + 1 ) ⊀ ​ ( 𝚯 t + 1 ( 1 ) ​ 𝜹 t + 1 ( 1 ) ) \bm{\delta}^{(0)}_{t+1}=\bm{J}_{\phi^{(0)}}\left(\bm{x}_{t+1}\right)^{\top}\left(\bm{\Theta}^{(1)}_{t+1}\bm{\delta}^{(1)}_{t+1}\right) , 𝒙 ^ t + 1 = Ο• ( 0 ) ​ ( β„³ ( 0 ) ​ ( 𝒙 t ) ) \hat{\bm{x}}_{t+1}=\phi^{(0)}(\mathcal{M}^{(0)}(\bm{x}_{t})) , Ο• ( 0 ) ​ ( β‹… ) \phi^{(0)}(\cdot) is non-linearity, and 𝑱 Ο• ( 0 ) ​ ( β‹… ) \bm{J}_{\phi^{(0)}}(\cdot) is the Jacobian. In this design, two blocks are in the same flow of gradient, but are updated based on different frequencies. We provide an example this design in Section 7 , when we discuss continuum memory systems.

Knowledge Transfer via Initialization. Model Agnostic Meta-Learning (MAML) ( Finn et al., 2017 ) , is one of the most popular form of meta-learning (or learning to learn) that aims to learn a global initial point for a model so it could learn fast a new task. From the nested learning perspective, there are two nested optimization processes, in which the inner problem iterates over its own context and based on its internal objective, the higher-level problem measures its own learned weight as the initial point of the inner problem. More formally, we have 𝚯 0 ( 1 ) = arg ⁑ min Ξ¦ ​ 𝔼 π’ž ∼ π’ž ( 0 ) ​ [ β„“ ⁑ ( β„³ ( 1 ) ​ ( β‹… , Ξ¦ ) , π’ž ) ] , \displaystyle\bm{\Theta}^{(1)}_{0}=\underset{\Phi}{\arg\min}\;\;\mathbb{E}_{\mathcal{C}\sim\mathcal{C}^{(0)}}\Biggl[\bm{\ell}(\mathcal{M}^{(1)}\left(\cdot;\Phi\right),\mathcal{C})\Biggr], (28) where the higher level block learns the best initial value over all possible contexts the lower level problem might have. As discussed earlier, any MAML-based learning model is an instance of this case, but as a more specific example, we refer to the example of (MLP Layer vs. Linear Attention), which is discussed in Section 3.2 and Figure 3 .

Connections with Generation. One of the most common form of knowledge transfer is through generating weights or context. That is, one lower-frequency (resp. higher-frequency) block generates the weight of a higher-frequency (resp. lower-frequency) block. More formally,

{ 𝚯 ( 1 ) = π’ˆ ⁑ ( β„³ ( 0 ) , β„’ ( 0 ) , π’ž ( 0 ) , 𝚯 ( 0 ) ) , or 𝚯 ( 0 ) = π’ˆ ⁑ ( β„³ ( 1 ) , β„’ ( 1 ) , π’ž ( 1 ) , 𝚯 ( 1 ) ) , \displaystyle\hskip-8.61108pt\begin{cases}\bm{\Theta}^{(1)}\!=\bm{g}\left(\mathcal{M}^{(0)};\mathcal{L}^{(0)}\!\!,\mathcal{C}^{(0)}\!,\bm{\Theta}^{(0)}\right),&\!\!\!\!\text{or \>}\\ \bm{\Theta}^{(0)}\!=\bm{g}\left(\mathcal{M}^{(1)};\mathcal{L}^{(1)}\!\!,\mathcal{C}^{(1)}\!,\bm{\Theta}^{(1)}\right),\end{cases} (Weight Generation)

{ π’ž ( 1 ) = π’ˆ ⁑ ( β„³ ( 0 ) , β„’ ( 0 ) , π’ž ( 0 ) , 𝚯 ( 0 ) ) , or π’ž ( 0 ) = π’ˆ ⁑ ( β„³ ( 1 ) , β„’ ( 1 ) , π’ž ( 1 ) , 𝚯 ( 1 ) ) . \displaystyle\begin{cases}\mathcal{C}^{(1)}\!=\bm{g}\left(\mathcal{M}^{(0)};\mathcal{L}^{(0)}\!\!,\mathcal{C}^{(0)}\!,\bm{\Theta}^{(0)}\right),&\!\!\!\!\text{or}\\ \mathcal{C}^{(0)}\!=\bm{g}\left(\mathcal{M}^{(1)};\mathcal{L}^{(1)}\!\!,\mathcal{C}^{(1)}\!,\bm{\Theta}^{(1)}\right).\end{cases} (Context Generation)

There are two important examples of the above form for knowledge transfer: (1) Hypernetworks: where the weights of a targeted neural network is generated by another (generator) network. (2) Optimization process: where the architecture generates the input for the optimizer. That is, the context (or input data) of an optimizer is the gradients that are generated by the architecture. For more discussion on this topic, see Section 4 . Note that this example is not necessarily about β€œ learned optimizers ” and it is valid for the commonly used optimization process and algorithms such as gradient descent, Adam ( Kingma and Ba, 2014 ) , AdaGrad ( Duchi et al., 2011 ) , etc.

A Note on Designing Neural Learning Modules. In the above, we discussed only some examples of possible knowledge transfer methods and also the potential connections of different levels. The formulation of NL and neural learning module, however, is general and so is not limited to the above specific set of methods. Accordingly, to design a neural learning module from a nested learning perspective, there are two important steps and design choices:

It is notable that with different choices of knowledge transfer, some learning paradigms can be seen as a part of a neural learning model: E.g., (1) Meta learning, when two blocks in two levels transfer their knowledge with one level meta-learns the other; more specifically, (2) Model Agnostic Meta Learning (MAML) ( Finn et al., 2017 ) , when knowledge transfer is though learning the initialization; (3) Hypernetworks, when one higher-frequency block generates the weights for the other lower-frequency block; (4) Learned optimizers, when the knowledge transfer is through data generation (i.e., one high-frequency block generates the gradient for the other lower-frequency block).

## 4 Optimizers as Learning Modules

In this section, we start with viewing backpropagation process and optimizing a neural network from the associative memory and data compression perspective. Then, we discuss how variants such as momentum-based optimizers are instances of nested associative memory systems. Finally, we discuss alternative methods leading to deep optimizers with higher expressive power from the associative memory perspective.

### 4.1 Backpropagation as an Associative Memory

Updating the weights of a neural network through backpropagation ( Linnainmaa, 1970 ; Rumelhart et al., 1986 ) has been the critical component of training large-scale deep neural networks. Intuitively, in this optimization process, first, the error of the model’s output with respect to target is calculated, and then each layer is updated based on its contribution to this error. This section aims to explain this process through the lens of associative memory and discuss how it fits within the nested learning paradigm. For the sake of clarity and simplicity, we assume a deep MLP model, but all the derived formulations in the following can simply be adapted to other architectures as well. Given an MLP with L L layers parameterized with { W β„“ β‹… + 𝒃 β„“ } β„“ = 1 L \{W_{\ell}\cdot+\>\bm{b}_{\ell}\}_{\ell=1}^{L} , the required gradients in backpropagation are computed as: βˆ‚ β„’ βˆ‚ W β„“ = 𝜹 β„“ ​ 𝒙 ^ β„“ βˆ’ 1 ⊀ , and 𝜹 β„“ = 𝑱 Ο• β„“ ​ ( 𝒛 β„“ ) ⊀ ​ ( W β„“ + 1 ⊀ ​ 𝜹 β„“ + 1 ) ⏟ local output surprise for layer β„“ , \displaystyle\frac{\partial\mathcal{L}}{\partial W_{\ell}}=\bm{\delta}_{\ell}\>\hat{\bm{x}}_{\ell-1}^{\top},\qquad\text{and}\qquad\bm{\delta}_{\ell}=\underset{\text{local output surprise for layer $\ell$}}{\underbrace{\bm{J}_{\phi_{\ell}}\left(\bm{z}_{\ell}\right)^{\top}\left(W_{\ell+1}^{\top}\bm{\delta}_{\ell+1}\right)}}, (29) where 𝒛 β„“ = W β„“ ​ 𝒙 ^ β„“ βˆ’ 1 + 𝒃 β„“ \bm{z}_{\ell}=W_{\ell}\>\hat{\bm{x}}_{\ell-1}+\bm{b}_{\ell} is pre-activation, and so 𝒙 ^ β„“ = Ο• β„“ ​ ( 𝒛 β„“ ) \hat{\bm{x}}_{\ell}=\phi_{\ell}\left(\bm{z}_{\ell}\right) is the output of β„“ \ell -th layer, Ο• β„“ ​ ( β‹… ) \phi_{\ell}(\cdot) is its non-linearity, and 𝑱 Ο• β„“ ​ ( β‹… ) \bm{J}_{\phi_{\ell}}(\cdot) is the Jacobian. Therefore, the update of the β„“ \ell -th layer with gradient descent is computed as: W β„“ t + 1 = W β„“ t βˆ’ Ξ· β„“ t + 1 𝜹 β„“ 𝒙 ^ β„“ βˆ’ 1 ⊀ . \displaystyle{W_{\ell}}_{{}_{t+1}}={W_{\ell}}_{{}_{t}}-\eta_{\ell_{t+1}}\>\bm{\delta}_{\ell}\>\hat{\bm{x}}_{\ell-1}^{\top}\>. (30) Here, 𝒙 ^ β„“ βˆ’ 1 \hat{\bm{x}}_{\ell-1} is the input of the layer and 𝜹 β„“ \bm{\delta}_{\ell} measures the local error signal for layer β„“ \ell or equivalently is a metric that measures the surprise of layer β„“ \ell ’s output given its input. Similar to our example in Section 3.1 , we can write Equation 30 as: W β„“ t + 1 = arg min W ⟨ W 𝒙 ^ β„“ βˆ’ 1 , 𝜹 β„“ ⟩ + 1 2 ​ Ξ· β„“ t + 1 βˆ₯ W βˆ’ W β„“ t βˆ₯ F 2 , \displaystyle{W_{\ell}}_{{}_{t+1}}=\arg\min_{W}\>\langle W\hat{\bm{x}}_{\ell-1},\bm{\delta}_{\ell}\rangle+\frac{1}{2\eta_{\ell_{t+1}}}\|W-{W_{\ell}}_{{}_{t}}\|_{F}^{2}, (31) which is an associative memory module that aim to map the input of each layer 𝒙 ^ β„“ βˆ’ 1 \hat{\bm{x}}_{\ell-1} to its local error signal, 𝜹 β„“ \bm{\delta}_{\ell} (see Definition 1 ). That is, the above formulation implies that the training process of a neural network with gradient descent and backpropagation can be viewed as a compression process , in which each layer stores the mappings between its input and the corresponding local error signal. Later in Β§ 4.5 , we discuss how this viewpoint helps with designing more expressive learning rules for backpropagation.

Backpropagation β‰  \neq Linear Attention. A common misinterpretation for Equation 30 is to assume 𝜹 β„“ \bm{\delta}_{\ell} is a pre-computed term and so backpropagation (at least on a linear layer) recovers Hebbian-rule, resulting in the equivalency of the optimization process and performing linear attention on gradients. Our formulation, however, shows that the update rule in backpropagation is a self-referential process ( Schmidhuber, 1993 ) , where the values of the associative memory is generated by itself, making it a more complex associative memory than a simple linear attention on gradients (see Section 4.5 ).

### 4.2 Momentum-based Optimizers as Associative Memories