Hierarchical Associative Memory, Parallelized MLP-Mixer, and Symmetry Breaking

Let us first briefly review the two-layer Hopfield network proposed in [12]. In this system, the dynamical variables are composed of $N_{v}$ visible neurons and $N_{h}$ hidden neurons both continuous,

where the argument $t$ can be thought of as “time”. The interaction matrices between them,

are basically supposed to be symmetric: $\xi^{\prime}=\xi^{\top}$, see Fig. 1. With the relaxing time constants of the two groups of neurons $\tau_{v}$ and $\tau_{h}$, the dynamics of the system is described by the following differential equations,

where the activation functions $f$ and $g$ are determined through Lagrangians $L_{h}:\mathbb{R}^{N_{h}}\to\mathbb{R}$ and $L_{v}:\mathbb{R}^{N_{v}}\to\mathbb{R}$, such that

The canonical energy function for this system is given by

One can easily find that this energy function monotonically decreases along the trajectory of the dynamical equations to define an associative memory model,

provided that the Hessians of the Lagrangians are positive semi-definite. In addition to this, if the overall energy function is bounded from below, the trajectory is guaranteed to converge to a fixed-point attractor state, which corresponds to one of the local minima of the energy function. Such fixed points and the process of convergence are thought of as associative memories and memory retrieval of an associative memory model. The formulation of neural networks in terms of Lagrangians and the associated energy functions enables us to easily experiment with different choices of the activation functions and different architectural arrangements of neurons.

Suppose we have a fixed interaction matrix $\xi_{\mu i}$, then the system is defined by the choice of Lagrangians $L_{h}$ and $L_{v}$. Tang and Kopp demonstrated that the specific choice of Lagrangians called “model C” in [12] essentially reproduces the mixing layers in the MLP-Mixer [14], which is given by the following Lagrangians:

where $F$ will be specified below, and $\bar{v}=\sum_{i}v_{i}/N_{v}$. For these Lagrangians, the activation functions are²²2$g_{i}$ can include learnable parameters such as $\bm{\alpha}$ and $\bm{\beta}$ as in [34] by slightly modifying the Lagrangian $L_{v}$, see Appendix A.

We now consider the adiabatic limit, $\tau_{v}\gg\tau_{h}$, which means that the dynamics of the hidden neurons is much faster than that of the visible neurons, i.e., we can take $\tau_{h}\to 0$:

By substituting this expression into the other dynamical equation, we find

Notice that we can put an arbitrary coefficient $\alpha$, which can be even zero, in front of the decay term, since for this choice of Lagrangian $L_{v}$ its Hessian has a zero mode:

If we take $\alpha=0$ and discretize the differential equation by taking $\Delta t=\tau_{v}$, then we obtain the update rule for the visible neurons as

where we defined $\sigma:=F^{\prime}$. If it is chosen as $\sigma=\operatorname{GELU}$, this update rule is identified with the token- and channel-mixing blocks in the mixing layers discussed in [4], as in Fig. 1.

As a simple extension in the context of hierarchical Hopfield network, we consider a structure with three layers: one visible layer with $N_{v}$ neurons and two hidden layers with $N_{s}$ and $N_{c}$ neurons each. The crucial point is that the set of the visible neurons lies in between the two hidden layers as depicted in Fig. 2, unlike the original configuration discussed in [13]. The dynamics of this system is then described by the following differential equations,

where $\xi^{(A,B)}$ denotes the interaction from $B$ neurons to $A$ neurons: $\xi^{(A,B)}\in\mathbb{R}^{N_{A}\times N_{B}}$. The activation functions are again determined through the corresponding Lagrangians,

In this system, the canonical energy function is given by [13]

where $s+1:=v$ and $v+1:=c$, and the arguments for $g^{A}$ and $L^{A}$ are omitted for simplicity. We have moved on to the matrix notation for brevity, and all the products here after imply matrix multiplication. Under the symmetric condition on the interaction matrices,

and ensuring the Hessians of the Lagrangians are positive semi-definite, the above dynamical equations confirm that

This fact demonstrates that we have a well-defined Lyapunov function for this system, which monotonically decreases along the trajectory of the dynamical equations.

We now take the specific Lagrangian for the visible neurons as

Then, the dynamical equations become

It is not necessary to specify the activation functions $g^{s}$ and $g^{c}$ at this stage. Following the discussion in Sec. 3.2, the discrete update rule for $x^{v}$ neurons will lead to a new type of MLP-Mixer model, which we will discuss in the next section. The flexible activation functions $g^{s}$ and $g^{c}$ suggest that the system may form a class of MetaFormers in turn. Thus, these dynamical equations represent a hierarchical associative memory model, serving as a continuous MetaFormer endowed with an energy function. This fact motivates us to refer to this system as the Energy MetaFormer.

Before getting into the discussion of deep neural network, it is informative to demonstrate the functionality of the Energy MetaFormer as an associative memory. To gain intuition, in this subsection we instantiate a model by taking a specific combination of Lagrangians and train the model on a denoising task using Eqs. (23). For the numerical demonstration, we utilize the HAMUX framework [35]³³3https://github.com/bhoov/barebones-hamux to implement the Energy MetaFormer.

Here, we set the dimensions of each layer to $N_{v}=784$ and $N_{s}=N_{c}=900$, and take the hidden Lagrangians as $\frac{1}{2}\max(x,0)^{2}$, which results in the $\operatorname{ReLU}$ activation function: $g^{s}=g^{c}=\operatorname{ReLU}$. The energy function for this system is determined by the Lagrangians,

The model is trained on 60,000 training samples from the MNIST dataset [36]. The mini-batch size is set to 512. We add Gaussian noise to the training batches before inputting them into the model. The input (initial state of visible neurons) evolves by the dynamical equations (23) and the model eventually outputs the state at one of the minima of the energy function Eq. (24). We compute the mean-squared error as the objective function between the outputs and the true images. The model is trained for 100 epochs with the Adam [37] optimizer (using Optax [38] defaults: $\beta_{1}=0.9$ and $\beta_{2}=0.999$) at a constant learning rate of $10^{-4}$.

Figure 2 depicts the receptive fields (a part of the weight vectors $\xi^{(s,v)}_{\alpha}$ and $\xi^{(c,v)}_{\beta}$) of the Energy MetaFormer learned by the training, which roughly correspond to the minima of the energy function adapted to the MNIST training samples. We can also verify the energy descent of the model. Randomly initialized neuron states of the trained model evolve according to the dynamical equations (23) to converge to one of the minima of the energy function, as illustrated in Fig. 2. These results demonstrate that the Energy MetaFormer has successfully learned to embed the training dataset into the network (weight matrices) and functions effectively as an associative memory model.

To derive the model, we essentially follow the three-layer hierarchical Hopfield network setup discussed in the previous section. A key point here is to introduce a two-dimensional structure into the layer of visible neurons, as shown in Fig. 3. The set of visible neurons can now be considered as a matrix rather than a vector, $x^{v}\in\mathbb{R}^{N_{v_{s}}\times N_{v_{c}}},~{}N_{v}=N_{v_{s}}\times N_{v_{c}}$, and interacts with the hidden neurons $x^{s}$ and $x^{c}$ along relatively perpendicular directions.

We continue to study the system Eq. (23), while it should be noted that $x^{v}$ has two indices, e.g.,

etc, where $i=1,\dots,N_{v_{s}}$ and $I=1,\dots,N_{v_{c}}$.

We now consider the adiabatic limit, $\tau_{v}\gg\tau_{s},\,\tau_{c}$, then the dynamical equations for $x^{s}$ and $x^{c}$ neurons are reduced to

By substituting these expressions to the other one and by following the same discussion as in Sec. 3.2, the differential equation for $x^{v}$ eventually becomes the update rule for the visible neurons:

Furthermore, as for the comparison with the vanilla MLP-Mixer (or more generally, MetaFormers), writing symmetric matrices

we truncate the interaction matrices between the visible and the hiddens as

With these expressions, we obtain the update rule for the visible neurons as shown in Fig. 3,

Each term on the right-hand side reads the skip-connection, the token-mixing, and the channel-mixing modules, respectively. The activation functions $g^{s}$ and $g^{c}$ are not specified at this stage. In practice, we set them to $g^{s}=g^{c}=\operatorname{GELU}$ in experiments to compare the proposed models with the vanilla MLP-Mixer.

Our single update rule (32) contains all the components of the mixing layer: token- and channel-mixing modules, layer normalization, and skip-connection, unlike Eq. (14) in Sec. 3.2. The two main differences from an ordinary mixing layer arise from our construction via the correspondence with the hierarchical Hopfield network: the layer normalization is simultaneously applied along both token and channel axes as in Eq. (25), and the layer processes features for token- and channel-mixing modules in a parallel manner.⁴⁴4Recently, a parallel Transformer block is occasionally adopted for some notable models (e.g., [39, 40]), and it shows comparable performance empirically. A model consisting of stacked associative memory models Eq. (32) as mixing layers becomes an MLP-Mixer model composed of parallelized token- and channel-mixing modules.

We incorporate the update rule Eq. (32) into the mixing layers, which results in a parallelized MLP-Mixer as a stack of associative memory models. To investigate the trainability of the proposed model and the implications from the correspondence with the Hopfield network side, we conduct several empirical studies to compare our model with the vanilla MLP-Mixer (VanillaMixer). The statistics of the results are obtained from ten trials with random initialization across all the experiments in this subsection. The experimental details are provided in Appendix B.

Network architecture. We consider two cases of the model: parallelized MLP-Mixer with symmetric weights as in Eqs. (28) and (29), which we refer to as SymMixer, and parallelized MLP-Mixer without weight constraint for comparison, which we simply refer to as ParaMixer. We incorporate the proposed module Eq. (32) with the mixing layers, and keep the rest of the networks exactly identical to VanillaMixer. From this construction, Para/Sym-Mixer actually do not have additional hyperparameters compared to VanillaMixer. The main difference between the proposed models and the ordinary MLP-Mixer is the normalization part within the mixing layers. As already mentioned, the layer normalization is applied over both the token and channel axes. This stems from the symmetry of the two directions inherent in the visible neurons as in Fig. 3 and is thus natural from the correspondence with the Hopfield network side. For control experiments, we implement this symmetric layer normalization for all the models, and examine some ablation studies below. We utilize PyTorch Image Models timm [41] to implement the models in all the experiments in this section.

Trainability. To study the trainability aspect of Para/Sym-Mixer, we perform scratch training of the models on a classification task with CIFAR-10/100 [42]. The CIFAR-10/100 datasets each consist of 60,000 natural images of size $32\times 32$. The ground-truth object category labels are attached to each image, and the number of categories is 10/100, with 6,000/600 images per class. There are 50,000 training images and 10,000 test images. For the training setting, we basically follow the previous works [3, 43]. The images are resized to $224\times 224$, and the AdamW [44] optimizer is used. We set the base learning rate to $\frac{\text{batch size}}{512}\times 5\times 10^{-4}$, and the mini-batch size is 384. As regularization methods, we employ label smoothing [45] and stochastic depth [46]. For data augmentation, we apply cutout [47], cutmix [48], mixup [49], random erasing [50], and randaugment [51]. For more details, see Appendix B.2. We train all the Mixer models under the exact same training configuration for fair comparison. All the trainings are conducted with four V100 GPUs.

Table 1 shows the results. The performance of ParaMixer is competitive with VanillaMixer, whereas SymMixer exhibits significant performance drops. This is quite counterintuitive since the difference between ParaMixer and SymMixer is only the symmetry constraint on the weight matrices. From previous studies [52, 53], it may be expected that many modern neural network architectures exhibit robust performance under certain symmetric constraints on weights. The enormous performance drop from ParaMixer to SymMixer suggests that symmetric weights and their symmetry breaking have an important effect on the capability of mixing layers. We will investigate this aspect in the next section.

Iterative mixing layers. Since the proposed mixing layer Eq. (32) is viewed as a lump of an associative memory model, it is natural to consider iterative updates of neurons. We use the trained Mixer models for CIFAR-10 and iteratively update the neuron states of the last mixing layer, which is regarded as the linear classifier of the network. The results are shown in Fig. 4. The inputs are 1,000 train and test samples of CIFAR-10 chosen at random. From these results, we observe that the classification performance of ParaMixer significantly degrades as the number of iterations of the mixing layer increases, while SymMixer maintains its performance over long iterations. This suggests that SymMixer stores associative memories of features during training and that the symmetric weights can properly retrieve them in the inference phase. We have also plotted the results for VanillaMixer for comparison, but it is difficult to make any sense since the mixing layers of VanillaMixer have nothing to do with the Hopfield network side.

As another aspect of the iterative mixing layers, we also perform scratch training of the Mixer models with a multiple number of iterations for all the mixing layers. Here, we adopt the number of iterations for all the mixing layers from one to four, and keep the rest of the training configurations intact. Table 2 shows the results. The results for a single iteration indicate the baseline from Table 1. We observe that the performance of Para/Sym-Mixer monotonically improves with a large number of iterations, while the performance improvement of VanillaMixer stops immediately. In particular, there is a clear margin between ParaMixer and VanillaMixer for a large number of iterations.

Ablation study. We here consider some ablation studies to examine the functions of the proposed models. Table 3 shows the overall results. We did not include bias terms in the two-layer MLPs of token- and channel-mixing blocks in Para/Sym-Mixer due to the correspondence with the Hopfield network side, while VanillaMixer does have. Table 3 tells us that adding the bias terms does not affect performance much, as ParaMixer with bias terms and VanillaMixer differ only in the parallel or serial mixing blocks, as depicted in Figs.3 and 1. As in Table 1, we have observed enormous performance drops in SymMixer. One might expect that the small number of parameters causes such a performance drop, since SymMixer has only roughly half the number of parameters as ParaMixer. (The number of parameters for each model is provided in Appendix B.2.) We train ParaMixer with half widths (Para (hw)) and SymMixer with double widths (Sym (dw)) to match the parameter number condition, and keep the remaining configurations unchanged.

Table 3 shows that there is still a large gap between Para (hw) and Sym (dw). This result indicates that the symmetry condition imposed on the weight matrices of mixing layers contributes more crucially to the classification performance than the number of parameters. In addition, this trend is observed not only between Para/Sym-Mixer but also in VanillaMixer (Table 3). This may imply that the symmetry breaking of weight matrices in the mixing modules of MetaFormers could generically be crucial to ensuring performance. Finally, to examine the effect of the unusual layer normalization within the Mixer models, we replace the normalization layer with the ordinary channel-only layer normalization. From Table 3, we see that all the Mixer models achieve slightly better performance. This is probably because our layer normalization normalizes the features over both token and channel axes, causing the sequence of token vectors to be overly normalized, thus suppressing individuality.

Hopfield networks, as associative memory models, are generally supposed to have symmetric weight matrices between neuron layers, as discussed in the previous sections. In order to examine the effect of symmetry breaking of weight matrices, we propose a minimal extension of the setup of the three-layer hierarchical Hopfield network. In doing so, it is natural to extend the original energy function Eq. (19) to the following interaction-symmetric form,

where the arguments for $g^{A}$ and $L^{A}$ are again omitted. If $\xi^{(v,s)}=\left(\xi^{(s,v)}\right)^{\top}$ and $\xi^{(v,c)}=\left(\xi^{(c,v)}\right)^{\top}$, this energy function is nothing but the original one.

We then extend the hidden-to-visible interactions to include a symmetry breaking term,

Taking the corresponding Lagrangians for Para/Sym-Mixer, and considering the adiabatic limit, $\tau_{v}\gg\tau_{s},\tau_{c}$, the total energy with symmetry breaking terms reads:

It is understood that $x^{s}$ and $x^{c}$ above are shorthands for Eq. (26). We refer to this extended energy function as the pseudo energy function. We defined the Lagrangians for the hidden neurons as $L^{s}=L^{c}=\sum_{a}\operatorname{\Phi_{\text{GELU}}}(x^{A}_{a})$, where $\operatorname{\Phi_{\text{GELU}}}$ is the primitive function of $\operatorname{GELU}$:

With these Lagrangians and the pseudo energy function, the time evolution of the visible neurons $x^{v}$ is governed by the differential equation,

where we set $\tau_{v}=1$. In this equation, the weight matrices from hidden neurons to visible neurons generically involve the symmetry breaking term as in Eq. (34). We conduct an empirical study of the effects of symmetry breaking on associative memories by solving this dynamical equation using the torchdiffeq framework [54, 55].

Let us first take a simple example of $N_{v_{s}}=2$ and $N_{v_{c}}=1$, for which the pseudo energy function Eq. (35) can be plotted in 2D/3D. A random initialization of the weight matrices with the symmetric condition, $\tilde{\xi}^{(v,\cdot)}=0$, yields an energy landscape as shown in Figs. 5 and 5. The energy function becomes symmetric with respect to $x^{v}_{1}+x^{v}_{2}=0$, as in this simple case the Lagrangian for $x^{v}$ reduces to $L^{v}=\left\lvert x^{v}_{1}-x^{v}_{2}\right\rvert/\sqrt{2}$. The minima of this energy function form two zero modes along the $x^{v}_{1}=x^{v}_{2}$ line. Any solution trajectories in this system approach the degenerate minima of the energy function (Fig. 5). By turning on the symmetry breaking terms $\tilde{\xi}^{(v,\cdot)}$, the degeneracy of the attractors is resolved and solution trajectories can retrieve different patterns in turn (Fig. 6). It is noted that for the dynamical equation (37) with the symmetry breaking terms, the pseudo energy function does not necessarily decrease along the trajectories. These observations imply that the associative memories the system stores degenerate for the symmetric weights, leading to an extremely limited effective memorization capability.

We can gain more insights into this system by increasing the number of neurons. For $N_{v_{s}}=4$, $N_{v_{c}}=8$, $N_{s}=20$, and $N_{c}=160$, we obtain Fig. 7, although the full landscape of the pseudo energy function cannot be visualized. We observe that $N_{v_{s}}$ roughly corresponds to the number of attractors, $N_{v_{c}}$ lifts the total energy, and a large total number of neurons makes the convergence slower. These observations are natural, as $N_{v_{s}}$ corresponds to the number of tokens and $N_{v_{c}}$ to the number of channels in terms of Mixer models of deep neural networks.

Having knowledge from the Hopfield network side, we consider the symmetry breaking of SymMixer by stacking the discrete version of the previous subsection as mixing layers. We refer to the parallelized MLP-Mixer with asymmetric weights as the AsymMixer, whose mixing layers consist of a symmetry-breaking extension of Eq. (32),

To perturbatively add the symmetry breaking and study its effect on the trainability of the model, we train the AsymMixer with the following custom loss:

where $\mathcal{L}_{\text{CE}}$ is the cross-entropy loss function, $L$ is the number of layers, and $\left\lVert\cdot\right\rVert_{F}$ denotes the Frobenius norm.

Figure 8 shows the results. We use CIFAR-10 to train the AsymMixers from scratch with $\lambda=\left\{0,10^{-r}\right\}$ for $r=0,\dots,6$. For $\lambda=1$ ($r=0$), the symmetry breaking terms $\tilde{W}$ are almost not allowed to have non-trivial entries due to the penalty term in the loss function, for which the trained AsymMixer boils down to SymMixer. For $\lambda=0$, in contrast, $\tilde{W}$ have no restrictions to learn and thus the model weights are trained in a similar manner as ParaMixer. Figure 8 provides reasonable results for the trainability of AsymMixers, interpolating the performance of SymMixer and ParaMixer. The discrepancy between the performance of AsymMixer with $\lambda=0$ and ParaMixer is due to the architecture designs; we assume that the symmetry breaking terms in AsymMixer are included only in the hidden-to-visible interactions for simplicity. These observations demonstrate consistent results with the previous subsection, indicating that symmetric weights in Mixer models have only a small amount of effective associative memories, and that symmetry breaking plays a crucial role in assuring performance.