Heterogenous Memory Augmented Neural Networks
It has been shown that semi-parametric methods, which combine standard neural
networks with non-parametric components such as external memory modules and
data retrieval, are particularly helpful in data scarcity and
out-of-distribution (OOD) scenarios. However, existing semi-parametric methods
mostly depend on independent raw data points - this strategy is difficult to
scale up due to both high computational costs and the incapacity of current
attention mechanisms with a large number of tokens. In this paper, we introduce
a novel heterogeneous memory augmentation approach for neural networks which,
by introducing learnable memory tokens with attention mechanism, can
effectively boost performance without huge computational overhead. Our
general-purpose method can be seamlessly combined with various backbones (MLP,
CNN, GNN, and Transformer) in a plug-and-play manner. We extensively evaluate
our approach on various image and graph-based tasks under both in-distribution
(ID) and OOD conditions and show its competitive performance against
task-specific state-of-the-art methods. Code is available at
\url{https://github.com/qiuzh20/HMA}.
Zihan Qiu, Zhen Liu, Shuicheng Yan, Shanghang Zhang, Jie Fu
# Heterogenous Memory Augmented Neural Networks

###### Abstract

It has been shown that semi-parametric methods, which combine standard neural networks with non-parametric components such as external memory modules and data retrieval, are particularly helpful in data scarcity and out-of-distribution (OOD) scenarios. However, existing semi-parametric methods mostly depend on independent raw data points - this strategy is difficult to scale up due to both high computational costs and the incapacity of current attention mechanisms with a large number of tokens. In this paper, we introduce a novel heterogeneous memory augmentation approach for neural networks which, by introducing learnable memory tokens with attention mechanism, can effectively boost performance without huge computational overhead. Our general-purpose method can be seamlessly combined with various backbones (MLP, CNN, GNN, and Transformer) in a plug-and-play manner. We extensively evaluate our approach on various image and graph-based tasks under both in-distribution (ID) and OOD conditions and show its competitive performance against task-specific state-of-the-art methods. Code is available at [https://github.com/qiuzh20/HMA](https://github.com/qiuzh20/HMA).

## 1 Introduction

Semi-parametric methods, which parametrize the mapping from an input domain \(\mathcal{X}\) to an output domain \(\mathcal{Y}\) with both neural net parameters and non-parametric data, are widely used in a variety of tasks including but not limited to meta-learning [42, 53], energy-based models [2] and planning [27]. With a non-parametric component in the neural net architecture, the model may better incorporate priors like distances between data points and characterize attributes such as data uncertainty. A typical design is to retrieve data associated with the current input with k-nearest neighbors (kNN). For instance, kNN-LM [29] proposes to perform inference with data retrieval from a small batch of inputs, similar to the concept of support set used in few-shot learning [54]. In practice, the retrieval strategies and parameters, such as the number of nearest neighbors, must be carefully designed for performance and low computational overhead [16, 47].

Instead of using full datasets, one may augment the neural net architecture with _external_ memory. With an external memory module, one can store a tiny amount of data points or latent features obtained through dataset distillation [60], random selection [31], or learned policies [47]. Especially with the attention mechanism employed by NPT [31], a model with external memory can utilize non-parametric components more flexibly than traditional methods such as Gaussian processes [12] due to the learnable dependencies instead of a fixed kernel function. Still, it suffers from the same scalability problem as the attention has to operate directly on a potentially large data set or features.

Inspired by some recent work that effectively learns a tiny amount of synthetic data for teaching student networks [60, 77], we propose Heterogeneous Memory Augmentation (HMA), a general approach for learning dependencies between data and augmenting networks without high computational overhead for various downstream tasks. HMA sequentially employs on the feature space of some backbone network (a) a real memory augmentation module (RMA), with classical memory augmentation methods following [42, 22], and (b) a synthetic memory augmentation module (SMA). Specifically, SMA learns compressed memory entries that encode _dataset-relevant_ information and leverages attention across datapoints (ABD) which consider the cross-relationship between the input batch in a manner similar
to NPT. These learned memory entries assist SMA in working effectively without relying on large batch sizes or additional selection methods. Notably, our HMA is architecture-agnostic and can be plugged into almost any backbone architecture in a few lines of code, without any additional training phases or losses.

In summary, our contributions are:

* We introduce heterogeneous memory augmentation (HMA), a better semi-parametric module that leverages both synthetic (learnable) and real (past data) memory. Compared to previous methods, our HMA can better capture the data dependencies for semi-parametric inference.
* The key component, synthetic memory augmentation (SMA), utilizes attention mechanisms to perform non-parametric inference on top of learned memory slots and inputs. The use of attention across datapoints allows the module to better capture the data dependencies 1) between the input batch and the memory slots as well as 2) between individual datapoints in the input batch, compared to traditional non-parametric techniques like Gaussian processes.
* Our model-agnostic HMA directly operates on the feature space and, therefore, can be incorporated into various backbones in a plug-in-play fashion.
* With learnable memory slots in SMA, HMA is more tailored to downstream tasks. Extensive experiments show its effectiveness in various in-distribution and out-of-distribution tasks.

## 2 Related Work

### Semi-parametric Methods

Retrieval methods leverage the similarity between data points to provide information for prediction. For input \(x\), the prediction probability is \(p=\lambda p_{\text{kNN}}+(1-\lambda)p_{\text{BaseModel}}\) to improve its ability to handle long-tail data and robustness [29]. During testing, the model can directly utilize similar training data for predictions. These characteristics make the retrieval method bring substantial improvements in various tasks of NLP [28, 71, 37] and graph tasks [59]. GNN-LM [40] also builds a graph by combining the samples with their corresponding retrieval data and further enhances the sample features through graph aggregation. Non-Parametric Classifier [44, 69] stores the training data's features and labels, enhancing similarity-based classification. Prototypical Networks [53, 45, 55] employ specifically designed objectives to learn a metric space. Each class's non-parametric "prototype representations" are easily obtained and used for few-shot and zero-shot tasks in this space. Compared to these methods, HMA learns synthetic memory slots directly with gradient descent updates, which can be more tailored for downstream tasks.

Figure 1: The overall framework of HMA. For inputs from some backbone, we concatenate them with label embeddings from **known** token similar to CLS. In RMA, a momentum memory queue and attention-based reader provide inputs with cross-batch information. During training, the real memory is updated with features from the momentum encoder and **true** label embeddings. In SMA, attention between datapoints is used to pass information among inputs and synthetic memory slots.
### Attention

Attention [58] can integrate various information and offer a versatile mechanism to uncover and exploit relationships between data. BatchFormer [25] performs attention in batches, allowing the model weights to learn the relationships. NPT [31] deploys Attention between Attributes and Attention between Datapoints to incorporate the learning of representation and relationships between data. However, ABD in NPT leads to a computational complexity of \(O(n^{2})\), which makes it difficult to incorporate a large amount of data and further hinders the feature learning process. To address the efficiency issue, [50] presents inducing points to avoid expensive full self-attention.

### Neural Memory

Memory mechanisms can be used to discover and exploit the relationships between data through contrastive learning [22, 48]. Such relationships are also pervasive in meta-learning. [70] proposes a learnable memory to capture the meta-class information of semantic segmentation during the base class training procedure. Unlike the works mentioned above that use memory in a no-parametric fashion, [54] utilizes relevant features in the input to recall memory and uses them to modify the network parameters. Meta Networks [42] and Neural Stored-program Memory [35] directly store network parameters in different training stages in memory.

Soft prompts [39, 38, 52] can also be regarded as a form of learnable memory that captures task-relevant information during fine-tuning and aids in leveraging the knowledge within pre-trained models. L2P [64, 63] stores a series of soft prompts and reads them using key-value matching, resulting in remarkable performance during continued learning.

## 3 Preliminaries

**Attention mechanism**[58] is a fundamental component in deep learning models that enables the model to focus on relevant information selectively. It is commonly formulated as follows:

\[\text{Att}(Q,K,V)=\text{softmax}(QK^{T}/\sqrt{d_{k}})V, \tag{1}\]

where \(Q,K,V\) are stacked matrixes of queries, keys, and values. Typically, multi-head attention (MultiHA) is used for large model capacity. Using multiple attention modules in parallel, the model can capture different aspects of the input data and combine their results. The outputs of the attention heads are concatenated to yield the final output:

\[\text{MultiHA}(Q,K,V)=\text{concat}(\text{head}_{1},\cdots,\text{head}_{h}), \text{ where }\text{head}_{i}=\text{Att}(QW_{i}^{Q},KW_{i}^{K},VW_{i}^{V})\]

where \(W_{i}^{Q},W_{i}^{K}\) and \(W_{i}^{V}\) are learnable matrices for the \(d_{k}\)-dimensional queries, keys and values for \(i\)-th head. By \(AB\) with \(A\in\mathbb{R}^{n\times a\times b}\) and \(B\in\mathbb{R}^{n\times b\times c}\), we mean batched matrix multiplication in Einstein summation notation einsum\((A,B,\) nab,nbc-nac\()\). The output of MHA is computed with additional components, including feed-forward layer (FF), layer normalization (LN) and residual connection:

\[\text{MHA}(Q,K,V)=H+\text{FF}(\text{LN}(H)),\text{ where }H=V+\text{ MultiHA}(\text{LN}(Q),\text{LN}(K),\text{LN}(V)) \tag{2}\]

**Non-parametric transformer (NPT)**[31] extends the standard transformer in a semi-parametric way by leveraging the dependencies between datapoints with the attention among large batched inputs. Specifically, NPT first encodes the input batch \(H\in\mathbb{R}^{bs\times d}\) and reshape the output into \(H^{\prime}=\mathbb{R}^{1\times bs\times(d*h)}\), and then applies attention across datapoints (ABD) by treating the input batch dimension as the token dimension:

\[H^{\prime}_{\text{ABD}}=\text{ABD}(H)=\text{MHA}(H^{\prime},H^{\prime})\in \mathbb{R}^{1\times bs\times(d*h)} \tag{3}\]

Once information has been exchanged between datapoints in the input batch, the ABD output \(H^{\prime}_{\text{ABD}}\in\mathbb{R}^{bs\times 1\times(d*h)}\) is reshaped back into \(H_{\text{ABD}}\in\mathbb{R}^{bs\times d\times h}\) to apply a self-attention operation on the feature dimension \(d\), which we term attention across attributes (ABA) to distinguish it from ABD:

\[\text{ABA}(H_{\text{ABD}})=\text{MHA}(H_{\text{ABD}},H_{\text{ABD}},H_{\text {ABD}})\in\mathbb{R}^{bs\times d\times h} \tag{4}\]

ABD and ABA can be iteratively applied for several times before feeding the resulted (and reshaped) feature into a classifier or some other neural net.

## 4 Heterogenous Memory Augmentation

HMA performs two consecutive steps on the extracted feature batch 4.1: 1) real memory augmentation (RMA) 4.2, and 2) synthetic memory augmentation (SMA) 7.4.2. As semi-parametric methods heavily rely on the quality of non-parametric components, we employ RMA to _expand the non-parametric part implicitly_. Subsequently, we combine it with synthetic memory and apply attention between datapoints. The final output feature from SMA is fed into a linear projection head to compute logits for classification. We summarize HMA in the supplementary material.
### Feature Extraction

Given input \(X\) and label \(Y\) with batch size \(bs\), We compute the feature \(h^{1}=E(X)\in\mathbb{R}^{bs\times d_{1}}\), where \(E(\cdot)\) can be any backbone, and \(d_{1},d_{2}\) is the feature embedding dimension. In line with the setup of NPT, we incorporate labels as attributes. However, directly adding the embeddings of the true labels would leak information. To address this, for an n-classification problem, we replace the true labels \(Y\) with an additional label, denoted as n+1, and treat it as the CLS token as common in transformer literature [14]. We compute label embedding \(e=E_{\text{label}}(\textbf{n+1})\in\mathbb{R}^{bs\times d_{2}}\), where \(d_{2}\) is the dimension of the label embedding. After concatenating, we obtain aggregated features \(C^{1}=[h^{1};e]\in\mathbb{R}^{bs\times d},d=d_{1}+d_{2}\). This concatenated representation \(C^{1}\) is used for later processes.

### Real Memory Augmentation (RMA)

The properties of non-parametric components greatly influence the effectiveness of semi-parametric methods. NPT uses large batch sizes to reduce the influence of randomly selected non-parametric components. Retrievers ensure reference data quality through top-k selection. We aim to develop a method that effectively provides information beyond individual batches.

Drawing inspiration from MoCo [22], which employs a momentum-updated memory queue to provide stable and extensive contrastive samples, we introduce a real memory buffer \(M_{\text{buffer}}\in\mathbb{R}^{m_{1}\times d}\) with size \(m_{1}\) (this buffer is empty at the beginning) and a momentum encoder \(E_{m}\).

**Reading.** We utilize MHA2 to extract information from the memory buffer, thereby implicitly providing cross-batch information for subsequent SMA. To achieve this, we first transform \(M_{\text{buffer}}\) from \(\mathbb{R}^{m_{1}\times d}\) to \(\mathbb{R}^{1\times m_{1}\times d}\) and \(C^{1}\) from \(\mathbb{R}^{bs\times d}\) to \(\mathbb{R}^{bs\times 1\times d}\). Then, we repeat \(M_{\text{buffer}}\)\(bs\) times in the first dimension to obtain \(M^{\prime}_{\text{buffer}}\in\mathbb{R}^{bs\times m_{1}\times d}\). Subsequently, we concatenate \(M^{\prime}_{\text{buffer}}\) and \(C^{1}\) along the second dimension, resulting in \(C^{1}_{\text{aug}}\in\mathbb{R}^{bs\times 1+m_{1}\times d}\)). RMA computes the output feature \(C^{2}\) by:

\[C^{2}=\text{MHA}(C^{1},C^{1}_{\text{aug}},C^{1}_{\text{aug}})\in\mathbb{R}^{ bs\times 1\times d} \tag{5}\]

After RMA, \(C^{2}\) incorporates information from real data beyond a single batch. In 5.3, we observed that using RMA alone seldom brings improvements. Moreover, while using SMA alone can lead to some improvements, adding RMA achieves pronounced improvements.

**Writing.** We insert new aggregated input features \([E_{m}(X),E_{\text{label}}(Y)]\in\mathbb{R}^{bs\times d}\) into the memory queue, and update the momentum encoder momentum update rule as in MoCo: \(E^{t+1}_{m}=\lambda E^{t}_{m}+(1-\lambda)E^{t+1}\), in which \(E^{t+1}\) is the SGD-updated encoder from \(E^{t}\).

### Synthetic Memory Augmentation (SMA)

In NPT, the ABD mechanism allows models to utilize non-parametric components flexibly. We can also leverage this mechanism to introduce learnable synthetic memory, which captures dataset-relevant information. These synthetic memories also serve as compact non-parametric components for the ABD mechanism, eliminating the need for NPT's reliance on large batch sizes.

Although we can directly initialize a series of learnable parameters as memory slots with dimensions \(d_{1}+d_{2}\), we also aim to incorporate label information into the inference process of SMA. Specifically, for each category \(i\in[n]\), we initialize \(m_{2}\) learnable memory slots. Each memory slot is in \(\mathbb{R}^{d_{1}}\) and has corresponding label \(y_{i}\). With the label embedder \(E_{\text{emb}}\), we map the label \(y_{i}\) to the label embedding \(e_{\text{syn},i}\in\mathbb{R}^{d_{2}}\). By concatenating all learnable memory slots and label embeddings, we obtain a learnable memory \(M_{\text{SMA}}\in\mathbb{R}^{(n*m_{2})\times(d_{1}+d_{2})}\). We reshape \(C^{2}\) and \(M_{\text{SMA}}\), and concatenate them to obtain \(C^{2}_{\text{aug}}\in\mathbb{R}^{1\times(bs+m_{2}*n)\times d}\). The SMA output features are obtained by

\[C^{3}=\text{MHA}(C^{2},C^{2}_{\text{aug}},C^{2}_{\text{aug}})\in\mathbb{R}^{ 1\times bs\times d}. \tag{6}\]

Finally, we use a linear projection head to compute the logits from \(C^{3}\). During the training phase, the learnable memory slots are updated by gradient, while the label embeddings are regenerated using the updated embedder \(E_{\text{emb}}\).

## 5 Experiments

In this section, we evaluate the performance of HMA and explore its characteristics. Section 5.1 demonstrates that HMA can be combined with various backbones and improve image and graph tasks. Following the success of semi-parametric methods on OOD tasks, in section 5.2, we examine HMA's out-of-distribution generalization capabilities under different distributional shifts. Finally, in Section 5.3, we demonstrate the characteristics of both real and synthetic memory slots. Ablation study shows while ABD itself can bring improvements, with the help of real memory and synthetic memory slots, HMA consistently achieves competitive performance against task-specific state-of-the-art methods.
**Experiment Setup.** Our experiments strictly follow the basic settings like batch size, learning rate schedule, number of epochs, optimizer, evaluation metric, and test model selection in previous works for a fair comparison. In CIFAR-10 task, we use the codebase3. Each method is trained 200 epochs, and the final test accuracies over three random seeds are reported. In tuning ViT [15] task, we follow [52] to tune the ViT-B/32 model, which was pretrained on Imagenet-21K [13]. In Colored MNIST task 2, we use code base from [3]. In obg molecular property prediction tasks 3, we use the official code base. In real-world graph ID tasks 4 and graph OOD tasks 65, we follow the official implementation of G-mixup [62] and CIGA [21]. More details can be found in the supplementary material.

Footnote 3: [https://github.com/kuangliu/pytorch-cifar](https://github.com/kuangliu/pytorch-cifar)

**Table Notations.** We have **bolded** the best results and added an underline to the second-best results in all the tables. To present the effects of different components in HMA and provide a comprehensive view, we included the results of ablations for each component in the subsequent tables: 1) RMA represents only real memory augmentation from 4.2. 2) ABD and 3) ABD+SYN represent the use of only attention between datapoints from 7.4.2 without and with learnable synthetic memory, respectively. 4) ABD+RMA represents using attention between datapoints and RMA.

### HMA is a General Augmentation Mechanism

We initially tested HMA on CIFAR-10. We adopted a ResNet-18 [23] encoder, similar to the one used in NPT [31]. The original NPT framework reported a performance of 93.7% on CIFAR-10, which is even lower than using ResNet alone (93.9%). HMA achieves an accuracy of 95.54\(\pm\)0.11%, slightly outperforming the backbone ResNet (which achieves 95.43\(\pm\)0.01%). We acknowledge that minor differences may be attributed to the fact that this setting has already been well-studied.

To evaluate the performance of HMA on larger models and its adaptability to pre-trained models, we conducted experiments following the setup of ViT [15] on the pets37 [46], flowers102 [43], CIFAR100 [32], and DTD [10] datasets. We performed a simple hyperparameter search for HMA using different values of \(m_{1}=1,2,4\times bs\) and \(m_{2}=4,8,16\). In addition to directly tuning the ViT backbone, we compared HMA with the SOTA methods [52] that introduced soft prompts (referred to as learnable memory in the paper). "prompt-p" in the table indicates adding soft prompts only at the input layer, similar to p-tuning [39]. "prompt-a" refers to independently adding soft prompts at each layer of ViT, similar to p-tuning v2 [38]. PARAM refers to the ablation where non-parametric memory is not introduced, and only the parameters of the HMA 2 module are used. We report the mean and standard deviation for 3 seeds. The results are summarized in Table 1.

From Table 1, we can observe that the complete HMA consistently achieves better results than the best prompt tuning. The inferior performance of prompt tuning may be attributed to the interference caused by inserting new randomly initialized parameters. In contrast, HMA operates solely in the feature space of the backbone's output, avoiding such interference and allowing for more effective utilization of the original parameters.

We also conducted experiments following [52] in the scenario where only the CLS token and head of the pretrained model are tuned, and the results are shown in Table 2. Since HMA and prompt tuning are orthogonal techniques, we also examined the effects of adding HMA on top of prompt tuning. Aloughth HMA, operating solely in the feature space, does not show as significant improvement as prompt tuning when most of the model parameters are fixed, it can be combined with prompt tuning to achieve more pronounced improvements.

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline Dataset & DTD & PETS37 & FLOWERS102 & CIFAR100 & AVG \\ \hline \hline ViT & 75.11\(\pm\)0.24 & 90.71\(\pm\)0.07 & 98.48\(\pm\)0.40 & 91.48\(\pm\)0.27 & 88.95 \\ prompt-p & 75.69\(\pm\)0.56 & 90.49\(\pm\)0.22 & 98.59\(\pm\)0.22 & 91.61\(\pm\)0.14 & 89.09 \\ prompt-a & 75.35\(\pm\)0.81 & 90.47\(\pm\)0.28 & 98.42\(\pm\)0.11 & 91.62\(\pm\)0.07 & 88.97 \\ \hline \hline param & 75.12\(\pm\)0.67 & 90.27\(\pm\)0.05 & 98.46\(\pm\)0.29 & 91.69\(\pm\)0.08 & 88.89 \\ rma & 75.67\(\pm\)0.16 & 89/96\(\pm\)0.31 & 98.61\(\pm\)0.05 & 91.63\(\pm\)0.31 & 88.97 \\ abd & 75.89\(\pm\)1.04 & 90.66\(\pm\)0.20 & 98.57\(\pm\)0.05 & 91.65\(\pm\)0.14 & 89.19 \\ abd+rma & 75.69\(\pm\)0.40 & 90.54\(\pm\)0.77 & 98.69\(\pm\)0.14 & 91.63\(\pm\)0.19 & 89.14 \\ abd+syn & 75.69\(\pm\)0.75 & 90.09\(\pm\)1.10 & 98.49\(\pm\)0.15 & 91.72\(\pm\)0.08 & 89.00 \\ HMA & **76.05\(\pm\)**0.39 & **91.08\(\pm\)**0.22 & **98.70\(\pm\)**0.16 & **91.93\(\pm\)**0.19 & **89.44** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Accuracy of full fine-tuning pretrained ViT on 4 image datasets. PROMPT-P and PROMPT-A indicate adding soft prompts only at the input layer and at each layer, respectively. PARAM is the parameter ablation. Other notations are introduced in 5.
To further validate the effectiveness of HMA with diverse backbones, we combine HMA with GCN [30], GIN [72] and their variants GCN-v, GIN-v[18, 36] backbones on three molecular property prediction tasks on the ogb benchmark [26]. The results are summarized in Table 3. HMA brings improvement among 10 of 12 reported AUROCs.

Furthermore, we follow G-mixup [21] and compare it with various augmentation methods, including DropEdge [51], DropNode [74], Subgraph [61], and Manifold-Mixup [62] on different datasets. All methods use the same GIN [72] backbone. The results are summarized in Table 4. We can see the three components of HMA have a similar effect to CIFAR-10, and full HMA consistently outperforms vanilla GIN backbone. By doing meaningful mixup with the help of elaborately designed graphon, G-mixup achieves better performance than simply mixing graphs in feature space [62]. Although HMA also only operates in feature space, it significantly outperforms Manifold-Mixup and shows competitive results to G-Mixup.

### HMA Performs Competitively on OOD Tasks

We test the OOD performance of HMA on the Colored MNIST dataset introduced in IRM [3]. This task is distinct from traditional MNIST as it introduces misleading correlations between labels and colors. Specifically, digits 0-4 and 5-9

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline Dataset & IMDB-B & IMDB-M & REDD-B & REDD-M5 & REDD-M12 & AVG \\ \hline vanilla & 71.55\(\pm\)3.53 & 48.83\(\pm\)2.75 & 92.59\(\pm\)0.86 & 55.19\(\pm\)1.02 & 50.23\(\pm\)0.83 & 63.68 \\ D-edge & 72.20\(\pm\)1.82 & 48.83\(\pm\)3.02 & 92.00\(\pm\)1.13 & 55.10\(\pm\)0.44 & 49.77\(\pm\)0.76 & 63.58 \\ D-node & 72.16\(\pm\)0.28 & 48.33\(\pm\)0.98 & 90.25\(\pm\)0.98 & 53.26\(\pm\)4.99 & 49.95\(\pm\)1.70 & 62.79 \\ S-graph & 68.50\(\pm\)0.86 & 47.25\(\pm\)3.78 & 90.33\(\pm\)0.87 & 54.60\(\pm\)3.15 & 49.67\(\pm\)0.90 & 62.07 \\ M-mixup & 70.83\(\pm\)1.04 & 49.88\(\pm\)1.34 & 90.75\(\pm\)1.78 & 54.95\(\pm\)0.86 & 49.81\(\pm\)0.80 & 63.24 \\ G-mixup & 71.94\(\pm\)3.00 & **50.46\(\pm\)**1.49 & **92.90\(\pm\)**0.87 & 55.49\(\pm\)0.53 & **50.50\(\pm\)**0.41 & 64.25 \\ \hline RMA & 71.17\(\pm\)4.75 & 49.89\(\pm\)2.50 & 91.92\(\pm\)1.70 & 55.23\(\pm\)0.58 & 49.32\(\pm\)1.45 & 63.51 \\ ABD & 70.67\(\pm\)3.25 & 47.44\(\pm\)1.84 & 92.42\(\pm\)0.14 & 54.87\(\pm\)0.45 & 49.08\(\pm\)0.86 & 62.90 \\ abd+syn & 71.67\(\pm\)3.25 & 48.97\(\pm\)2.97 & 91.08\(\pm\)1.23 & 55.90\(\pm\)0.95 & 49.29\(\pm\)1.00 & 63.38 \\ HMA & **72.62\(\pm\)**0.48 & 49.78\(\pm\)0.74 & 92.80\(\pm\)1.37 & **56.37\(\pm\)**1.65 & 50.27\(\pm\)0.60 & **64.37** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Accuracies of different augmentation methods on real-world graphs. We strictly follow the official implementation and report the baseline results from G-mixup [21].

\begin{table}
\begin{tabular}{l c c c} \hline \hline Dataset & Dtd & PETS37 \\ \hline \hline vit\(\rightarrow\)+hma & 74.01\(\pm\)0.29\(\rightarrow\)74.91\(\pm\)0.39 & 89.17\(\pm\)0.07\(\rightarrow\)89.72\(\pm\)0.22 \\ prompt-p\(\rightarrow\)+hma & 73.98\(\pm\)0.56\(\rightarrow\)73.95\(\pm\)0.73 & 89.22\(\pm\)0.22\(\rightarrow\)89.58\(\pm\)0.60 \\ prompt-a\(\rightarrow\)+hma & 74.75\(\pm\)0.86\(\rightarrow\)**75.44\(\pm\)**0.06 & 89.88\(\pm\)0.28\(\rightarrow\)**90.71\(\pm\)**0.27 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Accuracy of only tuning CLS token and head of ViT. +HMA indicates adding HMA to the original method. The results on the right side of \(\rightarrow\) represent the performance after adding HMA.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Dataset & bace & bbbp & HIV & AVG \\ \hline GCN & 79.15\(\pm\)1.44 & 68.87\(\pm\)1.51 & 76.06\(\pm\)0.97 & 74.69 \\ GCN+hma & 80.59\(\pm\)1.15 & 69.96\(\pm\)0.55 & 77.26\(\pm\)1.42 & 75.85(\(\uparrow\)1.16) \\ \hline GCN-v & 68.93\(\pm\)6.95 & 67.80\(\pm\)2.35 & 75.99\(\pm\)1.19 & 70.91 \\ GCN-v+hma & 72.53\(\pm\)1.36 & 68.26\(\pm\)0.85 & 76.51\(\pm\)1.70 & 72.73(\(\uparrow\)1.82) \\ \hline GIN & 72.97\(\pm\)4.00 & 68.17\(\pm\)1.48 & 75.58\(\pm\)1.40 & 72.24 \\ GIN+hma & 77.18\(\pm\)2.20 & 69.47\(\pm\)1.50 & 76.47\(\pm\)0.17 & 74.38(\(\uparrow\)2.14) \\ \hline GIN-v & 73.46 \(\pm\)5.24 & 69.71\(\pm\)1.92 & 77.07\(\pm\)1.40 & 73.41 \\ GIN-v+hma & 76.56\(\pm\)1.38 & 68.43\(\pm\)1.19 & 75.83\(\pm\)1.23 & 73.61(\(\uparrow\)0.20) \\ \hline \hline \end{tabular}
\end{table}
Table 3: AUROC and standard deviation of backbones and HMA on ogb benchmark. We have marked the relative improvements over the backbone with an upward arrow (\(\uparrow\)) for the average results.
are considered two separate classes. Two domains \(p_{1},p_{2}\) can be accessed during training: in the first domain, green digits have a probability of \(p_{1}\) to be in 0-4; in the second, the probability is \(p_{2}\). During testing, there are two sets: 1. reversed color: green digits have a 90% chance of being 0-4; 2. gray: blue and red images are turned into gray. The original image size (28,28) is downsized for computational convenience to (14,14). HMA is tested with MLP from IRM [3] and CNN from [19] under both downsized and original image sizes. We modify \(p_{1},p_{2}\) from \((0.2,0.1)\) to \((0,0.02)\) to more accurately demonstrate the model's performance with distribution shifts.

This task provides a proof of concept for the semi-parametric approach in addressing out-of-distribution (OOD) tasks. While there is a distribution shift, there are still certain data points within the training set that share the same distribution as the test data. The semi-parametric approach effectively reduces the distribution shift by leveraging these data points as a support set. The average accuracy of the two test domains is presented in Figure 2. With different backbones, image sizes, and distribution shifts, we can observe that HMA consistently improves OOD performance.

In the following graph tasks, we follow CIGA [8] to examine HMA on diverse distribution shifts. Our baselines include ERM [57], SOTA interpretable GNNs like ASAP Pooling[49], and DIR[66], and SOTA OOD objectives like IRM[3], v-Rex[33], and IB-IRM[1].

In Table 6, we examine HMA on graph size shift scenarios on datasets converted from TU benchmarks[41]. Distribution shifts are generated through size-specific dataset splits following [8][73]: the training set comprises graphs with sizes smaller than the 50th percentile, 10% of which are held-out for the validation set, and the test set consists of those with sizes larger than the 90th percentile. Matthews correlation coefficient [4] is adopted due to the class imbalance. In Table 5, we further test HMA on larger real-world datasets with more complicated node degree biases. We use the Graph-SST datasets following CIGA[8]. Graphs are assigned based on their average degree to generate distribution shifts: the training set includes graphs with average degrees smaller than the 50th percentile, the validation set includes those with average degrees larger than the 50th percentile and smaller than the 80th percentile, and the test set includes the left. Graphs in Twitter are assigned in an inversed order to investigate OOD generalization ability from large degree graphs to small ones.

ERM is a strong baseline for these real-world tasks compared with specifically designed algorithms [19]. Our HMA consistently improves ERM baselines without introducing complex OOD objectives or learning strategies.

### Heterogenous Memory Shows Reasonable Properties

We use T-SNE [56] to visualize the relationships between the synthetic memory features and real data features in \(C^{2}\). We excluded the label embeddings to focus on the synthetic memory slots since they would result in identical embeddings for instances with the same label and affect the T-SNE results.

\begin{table}
\begin{tabular}{l c c c} \hline \hline Dataset & SST5 & TWITTER & avg \\ \hline ERM & 43.72\(\pm\)0.59 & 63.54\(\pm\)2.10 & 53.63 \\ ASAP & **44.16\(\pm\)**1.36 & 60.68\(\pm\)2.10 & 52.42 \\ DIR & 41.12\(\pm\)1.96 & 59.85\(\pm\)2.98 & 50.49 \\ \hline IRM & 43.69\(\pm\)1.26 & 63.50\(\pm\)1.23 & 53.60 \\ V-REX & 43.28\(\pm\)0.52 & 63.21\(\pm\)1.57 & 53.25 \\ IB-IRM & 40.85\(\pm\)2.08 & 61.26\(\pm\)1.20 & 51.06 \\ CIGAV1 & 43.70\(\pm\)1.98 & 62.02\(\pm\)2.28 & 52.68 \\ CIGAV2 & 43.30\(\pm\)0.90 & 61.80\(\pm\)2.03 & 52.55 \\ \hline HMA & 44.03\(\pm\)1.01 & **65.12\(\pm\)**2.17 & **54.58** \\ \hline \hline \end{tabular}
\end{table}
Table 5: Acc of OOD algorithms on averaged degrees shift.

Figure 2: OOD results for colored MNIST (left: original size, right: downsized) experiments. We change the OOD sample ratio from \((0.2,0.1)\) to \((0,0.02)\)
We present the synthetic memory slots' relationships in Figure 3 (left). It can be observed that slots belonging to the same class cluster together, indicating that the synthetic memory slots can learn class-specific information with the help of label embeddings. In Figure 3 (right), we show the relationships between three classes of synthetic memory slots and their corresponding real data features. It can be observed that different classes of real data features form separate clusters, while three separate clusters in Figure 3 (left) merge into one. This suggests that the learned memory in the synthetic memory slots differs from the real data features. The improvement in 5.1 and 5.2 may come from the information beyond the instance level in the synthetic memory slots.

We also visualize the synthetic memory slots features learned in graph tasks in Figure 4. Although the synthetic memory slots of class 1 scatter over other classes' clusters Figure 4 left, there is an interesting observation that the features of class 1 mix with the other three classes in the feature space correspondingly, as shown in Figure 4 right. The correspondence indicates that the synthetic memory slots might capture class-level information, similar to prototype representations [53].

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline Dataset & NC11 & NC109 & proteins & DD & avg \\ \hline ERM & 0.15\(\pm\)0.03 & 0.16\(\pm\)0.07 & 0.19\(\pm\)0.07 & 0.29\(\pm\)0.1 & 0.198 \\ ASAP & 0.16\(\pm\)0.10 & 0.15\(\pm\)0.07 & 0.22\(\pm\)0.16 & 0.21\(\pm\)0.08 & 0.185 \\ DIR & 0.21\(\pm\)0.06 & 0.13\(\pm\)0.05 & 0.25\(\pm\)0.14 & 0.20\(\pm\)0.10 & 0.198 \\ \hline IRM & 0.17\(\pm\)0.02 & 0.14\(\pm\)0.01 & 0.21\(\pm\)0.09 & 0.22\(\pm\)0.08 & 0.185 \\ V-REX & 0.15\(\pm\)0.04 & 0.15\(\pm\)0.04 & 0.22\(\pm\)0.06 & 0.21\(\pm\)0.07 & 0.183 \\ IB-IRM & 0.12\(\pm\)0.04 & 0.15\(\pm\)0.06 & 0.21\(\pm\)0.06 & 0.15\(\pm\)0.13 & 0.158 \\ CIGAV1 & 0.25\(\pm\)0.08 & 0.26\(\pm\)0.09 & **0.35\(\pm\)**0.17 & 0.21\(\pm\)0.12 & **0.268** \\ CIGAV2 & **0.28\(\pm\)**0.10 & **0.27\(\pm\)**0.10 & 0.30\(\pm\)0.09 & 0.19\(\pm\)0.12 & 0.260 \\ \hline HMA & 0.19\(\pm\)0.05 & 0.18\(\pm\)0.04 & 0.34\(\pm\)0.07 & **0.32\(\pm\)**0.03 & 0.258 \\ \hline \hline \end{tabular}
\end{table}
Table 6: Matthews correlation coefficients of different OOD algorithms on graph size shifts for real-world graphs. We strictly follow CIGA [8] and reproduce the results of ERM, GIGAv1&v2.

Figure 4: T-SNE on REDD-M5 for synthetic memory features (left) and input features (right)

Figure 3: T-SNE on CIFAR-10 for synthetic memory features (left) and real data features with synthetic memory features (right). ’0-train’ and ’0-mem’ refer to real and memory data features, respectively. The preceding label corresponds to the class from CIFAR10.
**Ablation Study.** The in-distribution performances on REDD-M5 are depicted in Figure 5(a) for different memory configurations, while the out-of-distribution performances on Twitter and SST are illustrated in Figure 5(b) and (c) respectively. It can be seen that while a suitable increase in the synthetic memory size \(m_{2}\) can bring improvement, larger \(m_{2}\) will hmu's performance. One potential explanation could be that a large \(m_{2}\) increases the optimization difficulty of ABD. Furthermore, it can be seen that with the addition of real memory augmentation, an increase in \(m_{2}\) provides more improvement.

From Table 1 and Table 4, we can find that singly using ABD or SMA proves beneficial in certain scenarios but not significant. The possible reason for this observation is the sensitivity of semi-parametric approaches. By solely relying on SMA without any data filtering, the quality of the non-parametric elements remains unassured. However, upon incorporating RMA, the effectiveness of HMA surpasses that of individual deployments of RMA and SMA. This aligns with our conjecture that RMA explicitly offers more information to the synthetic memory learning process.

## 6 Discussions, Conclusions, Limitations and Future work

**Two-stage Augmentation.** While it seems more natural to perform both RMA and SMA in one layer, _i.e._ to compute MHA with features aggregated from the input features, buffer features and the learnable features, the overall complexity is quadratic due to the SMA operations. It is possible to alternatively use linear approximations to standard attention (_e.g._, Performer [9]), it is beyond our scope and we leave it to future work.

We propose a universal memory augmentation method that can be seamlessly integrated with various backbone architectures. SMA leverages parametric and non-parametric approaches in the complementary combination of Attention Between Datapoints and learnable memory, thereby enabling the capture of class-specific and task-aware information, which is beneficial for ID and OOD downstream tasks. Additionally, based on momentum memory queue and attention-based reading, RMA enhances SMA with cross-batch information. Our experimental results demonstrate that HMA can consistently improve the performance of backbones and performs competitively with task-specific best-performing designs.
