SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Semantic Memory Update: Inspired by the characteristics of human memory [31-33], we develop a novel semantic memory update mechanism that considers the data information from previous semantic memory  $\mathbf{M}^{(s)}$ , episodic memory  $\mathbf{M}^{(e)}$ , and new task  $D_{i}$ . Specifically, the semantic memory module is empty at the beginning of learning, and each unseen block (category)  $\mathbf{M}_{n}^{(s)}$  will be initialized with the mean and variance of the latent representations of data samples, which is collected from the current task  $D_{i}$  and belongs to  $n$ -th category. For those seen categories, their semantic memory can be updated with the data sample from the current task  $D_{i}$  using self-attention mechanism.
Following Zhen et al. [18], we construct a graph with respect to the memory  $\mathbf{M}_n^{(s)}$  for updating, which consists of latent representations of  $J_{n}$  data samples that belongs to  $n$ -th category and  $T$  data samples stored in episodic memory  $\mathbf{M}_n^{(e)}$ , denoted as  $H_{n} = \{\pmb{h}_{n,1},\pmb{h}_{n,2},\dots ,\pmb{h}_{n,J_{n}},\pmb{M}_{n,1}^{(e)},\dots ,\pmb{M}_{n,T}^{(e)}\}$ . The mean of the semantic memory  $\mathbf{M}_n^{(s)}$  can be updated as:
$$
\boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {\text {n e w}} \leftarrow \alpha \boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {\text {o l d}} + (1 - \alpha) \bar {\boldsymbol {\mu}} _ {\mathbf {M} _ {n} ^ {(s)}}, \bar {\boldsymbol {\mu}} _ {\mathbf {M} _ {n} ^ {(s)}} = \sum_ {j = 1} ^ {J _ {n} + T} \beta_ {n, j} H _ {n, j}, \beta_ {n, j} = \frac {g \left(\mathbf {M} _ {n} ^ {(s)} , H _ {n , j}\right)}{\sum_ {j = 1} ^ {J _ {n} + T} g \left(\mathbf {M} _ {n} ^ {(s)} , H _ {n , j}\right)}, \tag {16}
$$
where  $\alpha \in (0,1)$  is a hyperparameter;  $\beta_{n,j}$  indicates the attention weight between  $j$ -th node features and the source node  $\mathbf{M}_n^{(s)}$ ;  $g(\cdot)$  is the similarity function defined in Eq. (13). This operation allows the useful information to be retained in memory while erasing less relevant or trivial information.
After that, we employ the incremental update strategy of the Gaussian distribution [34] to update the variance of semantic memory  $\mathbf{M}_n^{(s)}$  as follows:
$$
\sigma_ {\mathbf {M} _ {n} ^ {(s)}} ^ {2} ^ {\text {n e w}} \leftarrow \alpha \left[ \sigma_ {\mathbf {M} _ {n} ^ {(s)}} ^ {2} ^ {\text {o l d}} + \left(\boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {\text {n e w}} - \boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {\text {o l d}}\right) ^ {2} \right] + (1 - \alpha) \left[ \sigma_ {H _ {n}} ^ {2} + \left(\boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {\text {n e w}} - \bar {H} _ {n}\right) ^ {2} \right], \tag {17}
$$
where  $\bar{H}_n$  and  $\sigma_{H_n}^2$  are the mean and variance of the set of node features  $H_{n}$ , respectively.
Episodic Memory Update: Inspired by the fact that it is difficult for a person to forget those things he/she often recalls (and vice versa), we design a frequency-based episodic memory update stragy. Typically, for each data sample  $\mathbf{M}_{n,t}^{(e)}$  stored in episodic memory  $\mathbf{M}_n^{(e)}$ , we use a frequency matrix, denoted as  $C_{\mathbf{M}_{n,t}^{(e)}}$ , to record how many times  $\mathbf{M}_{n,t}^{(e)}$  was recalled during training. Then, for updating  $\mathbf{M}_n^{(e)} = \{\mathbf{M}_{n,t}^{(e)}\}_{t=1}^T$ , if the current episodic memory block has not been filled yet, we will directly append the latent representation of a new data sample to the  $\mathbf{M}_n^{(e)}$  when a new task arrives; Otherwise, we will update  $\mathbf{M}_n^{(e)}$  by replacing the least used data sample in the block if its capability exceeds the limit  $T$ . We note that the frequency matrix  $C_{\mathbf{M}_{n,t}^{(e)}}$  will be reset as the zero matrix after each update of episodic memory. The detailed memory updating process algorithm can be found in Appendix A.
Datasets: The experiments are conducted on five widely used benchmark datasets of various sizes, including binarized Omniglot [2], MNIST [35], DOUBLE-MNIST [36], CelebA [37] and FS-CIFAR100 [38]. The split of training/testing set (also known as background-evaluation split) follows Lake et al. [2], in which scenario all the samples at test time are from new classes. We resize all the images of binary datasets to 28x28, FS-Cifar100 to 32x32, and CelebA to 64x64.
Baselines: 1) For VAE-based models, we utilize NS-related models, specifically focusing on two NS variants: Convolutional Neural Statistician (CNS), where the latent space is shaped with convolutions at a given resolution; Set-Context-Hierarchical-Aggregation-Variational-Autoencoder (SCHAVAE) [24], which employs an additional hierarchy over the context latent variable  $c$ . The CNS and SCHAVAE can be naturally extended under the framework of VSM-VAE, leading to VSM-CNS and VSM-SCHA, respectively. Meanwhile, we create several variants by fusing either semantic memory or episodic memory with CNS or SCHAVAE, as appropriate, to further demonstrate the efficacy of the proposed variational structured memory; 2) For Diffusion-based models, we consider both
Table 1: Few-shot generative evaluation of VAE-based models on various datasets with the set size 5, including Omniglot, Double-Mnist, FS-Cifar100, and Mnist (trained on Omniglot and tested on Mnist.). For all datasets, we set the episodic memory size of all methods to 5 and the semantic memory size to be the same as the number of categories appearing in training data.
Omniglot-ns	Double-mnist	Mnist	FS-Cifar100
NELBO	NLL	NELBO	NLL	NELBO	NLL	NELBO(bpd)	NLL(bpd)
VAE	101.5	95.9	74.3	69.2	125.1	105.1	-	-
NS	96.6	92.4	67.3	60.2	118.7	101.2	-	-
CNS	92.9	89.7	66.5	54.3	114.1	90.4	9.979	9.803
CNS + Semantic	84.4	72.4	56.7	50.2	107.3	88.5	9.952	9.745
CNS + Episodic	83.9	71.8	54.3	51.1	97.8	80.3	9.946	9.741
VSM-CNS(our)	81.1	68.2	47.8	40.6	92.9	77.0	9.872	9.711
SCHA-VAE	89.4	85.8	68.1	53.3	114.7	88.5	10.364	10.338
SCHA + Semantic	81.5	66.8	53.7	48.9	98.5	77.7	10.317	10.215
SCHA + Episodic	82.5	63.4	53.0	47.2	95.3	79.1	10.292	10.176
VSM-SCHA(our)	74.0	56.8	42.4	36.5	89.9	68.2	10.219	10.041
Table 2: Few-shot generative evaluation of diffusion-based models on various datasets with the set size 5, including Fs-Cifar100 and Celeba. (FID: Frechet score; sFID: spatial FID; P: precision; R: recall.)
FS-Cifar100	Celeba
FID ↓	sFID ↓	P ↑	R ↑	FID ↓	sFID ↓	P ↑	R ↑
DDPM	62.84	28.91	0.58	0.40	14.02	27.46	0.71	0.39
sDDPM	45.50	29.87	0.54	0.46	12.87	26.44	0.72	0.37
vDDPM	62.58	27.50	0.58	0.41	13.01	26.24	0.73	0.39
vFSDM	63.73	28.85	0.58	0.38	13.52	27.75	0.71	0.38
VSM-vFSDM(our)	44.11	23.11	0.59	0.46	12.27	25.79	0.72	0.40
FSDM	35.07	20.95	0.62	0.53	10.11	24.58	0.73	0.39
VSM-FSDM(our)	33.21	18.32	0.68	0.61	9.25	23.21	0.77	0.43
unconditional and conditional diffusion models as baselines, where we use the DDPM [39] as the unconditional baseline and other three conditional diffusion models: VDDPM, where the context  $c$  is a latent variable and the dataset  $D$  is generated conditioned on  $c$ ; VFSDM [24], another variational diffusion model where employees a different way to extract and aggregate set-information using ViT [40]: stack all the samples on the channel dimension; sDDPM adapting ideas in Sinha et al. [41] without contrastive learning, where a ViT encoder splits the image in patches and processes them jointly.
Model Settings: The details about model settings can be found in Appendix B.
5.2 Quantitative experiment
VAE-based models: To evaluate the generative properties of VAE-based models, we use the evidence lower bound (ELBO) to approximate the log marginal likelihood with 1000 importance samples (NLL). Our qualitative experiments are performed on Omniglot, DOUBLE-MNIST, MNIST, and FS-Cifar100. The outcomes of the experiment are displayed in Table 5.2. It is clear that VSM-CNS/VSM-SCHA can achieve significant improvements over its baseline models, proving that VSM could be a promising mechanism for few-shot generation. Additionally, by incorporating either semantic or episodic memory into CNS/SCHA-VAE, the resulting models can also achieve better performance, proving both memory types can improve the model performance of few-shot generation.
Diffusion-based models: As exhibited in Table 5, we employ the metrics below to assess the generative capacity of diffusion-based models, including: FID [42] for sample quality, sFID [43] to capture spatial relationships, and Precision and Recall [44] for measuring variety and mode coverage. From the results, VSM-Diffusion beats the other unconditional and conditional baselines on both two datasets, demonstrating the effectiveness of VSM on few-shot generation.
(a) Query Set
(b) Generative Samples
(c) Query Set
(d) Generative Samples
Figure 3: Visualization of few-shot generative samples using a VSM-Diffusion with the set size 5. The query set and generative samples on the left side are from FS-Cifar100, while those on the right side are from Celeba.
(a) Query Set
(b) Recalled Semantic Memory
(c) Recalled Episodic Memory
Figure 4: (a) Query Set; (b) Recalled blocks (categories) of semantic memory, each of which is represented by the images of corresponding category; (c) Recalled samples of episodic memory.
(a) Omniglot
Figure 5: Few-shot generative evaluation of VAE-based models on various datasets with different episodic memory sizes.
(b) Doublemnist
(c) Mnist
5.3 Qualitative experiment
Few-shot image generation: The generation results of VSM-Diffusion are exhibited in Figure 3. From the results, we can find that our memory-augmented models can generate class-specific samples through successfully extracting content information and realistic classes that are both complicated and varied. Limited by pages, additional few-shot generative experimental results on Omniglot, CelebA, and FS-Cifar100 can be found in Appendix C.2.
Table 3: Ablation study on the update strategy of variational structured memory (VSM).
Omniglot-ns	Double-Mnist	Mnist
NELBO	NLL	NELBO	NLL	NELBO	NLL
W/o Episodic memory	75.4	58.2	45.3	40.1	91.1	69.5
W/o Frequency-based	74.8	58.2	42.9	37.7	91.7	70.6
Combined strategy	74.0	56.7	42.4	36.5	89.9	68.2
Memory visualization: To further verify the effectiveness of the proposed VSM, we conduct experiments of memory visualization here. For instance, given a new set, as shown in Figure 4(a), the images of recalled semantic memory corresponding to the class are shown in Figure 4(b), and the images of recalled episode memory are shown in Figure 4(c). The visualization results demonstrate that VSM can recall related image classes and samples to the query set in an effective way.
Effect of memory size: To investigate the effect of memory size in episodic memory, we use the VSM-SCHA model as a typical example. We train a model on Omniglot data, and evaluate it with the Omniglot, DOUBLE-MNIST, and MNIST datasets for unknown classes. The size of the episodic memory is raised from 0 to 10. The results in Figure 5 show that the model performance improves with the increase of episodic memory size, further demonstrating the importance of episodic memory in few-shot generation tasks. We note that the developed variational structured memory will reduce to variational semantic memory [18] when the episodic memory size is adjusted to 0.
Effect of memory update process: We conduct two variants of VSM-SCHA to investigate the benefits of the proposed structured memory update strategy. For the first variant, its semantic memory update strategy does not consider episodic memory, denoted as W/o Episodic memory; For the second variant, its episodic memory update strategy does not employ a frequency-based strategy but uses a first-in, first-out queue-based memory update method, denoted as W/o Frequency-based. The experiment results in Table 3 can verify the effectiveness of the proposed memory update strategy.
In this study, we develop a novel variational structured memory module that can concurrently retain generic semantic structure about a category and specific information about a sample, the former being achieved by semantic memory and the latter being caught by episodic memory. Meanwhile, we introduce a bionic memory updating strategy for the conversion between episodic and semantic memories stored in VSM, where episodic memory can provide context information that will later be transferred and stored in semantic memory. Through incorporating VSM with few-shot latent variable models, we develop a novel hierarchical memory-augmented neural statistician (HMNS), whose generative process can be specified with a VAE or diffusion model, leading to VSM-VAE and VSM-Diffusion. Extensive experiments demonstrate that our developed VSM-VAE and VSM-Diffusion can achieve promising performance on few-shot generation tasks.
This project was supported in part by the National Natural Science Foundation of China under Grant U21B2006; in part by Shaanxi Youth Innovation Team Project; in part by the Fundamental Research Funds for the Central Universities QTZX23037 and QTZX22160; in part by the 111 Project under Grant B18039; in part by the Fundamental Research Funds for the Central Universities; in part by the Innovation Fund of Xidian University under Grant YJSJ23016; Additionally, this project was supported in part by the National Research Foundation, Singapore under its Industry Alignment Fund – Pre-positioning (IAF-PP) Funding Initiative.
[1] Brenden Lake, Ruslan Salakhutdinov, Jason Gross, and Joshua Tenenbaum. One shot learning of simple visual concepts. In Proceedings of the annual meeting of the cognitive science society, volume 33, 2011.
[2] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332-1338, 2015.
[3] Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient-based meta-learning as hierarchical bayes. arXiv preprint arXiv:1801.08930, 2018.
[4] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017.
[5] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126–1135. PMLR, 2017.
[42] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[43] Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W Battaglia. Generating images with sparse representations. arXiv preprint arXiv:2103.03841, 2021.
[44] Tuomas Kynkänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019.
A Memory Recall and Update Algorithm
The procedure of memory recall and update algorithm is summarized in Algorithm. 1.
Algorithm 1 Semantic Memory recall and update algorithm
Input: Given task  $D_{i}$ ; Semantic memory  $\mathbf{M}^{(s)}$ , and episodic memory  $\mathbf{M}^{(e)}$ ; Recall task-specific semantic memory according to Eq. 14 and episodic memory according to Eq. 15; if The semantic memory for the given task is null: then Computing the mean and variance for the semantic memory with given task  $D_{i}$ ; else Update the semantic memory according to Eq. 16 and Eq. 17; end if if The episodic memory for the given task is null: then Append episodic memory with given task  $D_{i}$ ; else Update the episodic memory according to frequency matrix  $C$ ; end if Infer variational posteriors for the latent variables of different layers  $\{\pmb{\theta}^{(l)}\}_{1 = 1}^{L}$ ;
Model Design: In most memory-based models, the images  $D_{i} = \{x_{ij}\}_{j=1}^{J_{i}}$  are fed into an encoder to produce image features  $\{h_{ij}\}_{j=1}^{J_{i}} \in \mathbb{R}^{D \times 1}$ , which are then represented as vectors. These vector representations are utilized for both recall and update operations in the model. Nonetheless, due to the high dimensionality of images, relying solely on vectors can lead to a loss of valuable information. By drawing inspiration from the approach proposed in the Vision Transformer, we can overcome these challenges and obtain more accurate memory information. We begin by obtaining image features, denoted by  $\{h_{ij}\}_{j=1}^{J_{i}} \in \mathbb{R}^{D \times H \times W}$ , then calculate the average of the set dimension resulting in  $h_{i} = \frac{1}{J_{i}} \sum_{j=1}^{J_{i}} h_{ij}$ . Afterwards, we reshape the size of the image features to  $\{h_{i}\} \in \mathbb{R}^{N \times D}$  where  $D$  denotes the patch size. When updating  $M_{n}^{(s)} \in \mathbb{R}^{N \times D}$ , we employ the resized feature  $h_{i}^{m} \in \mathbb{R}^{(J_{i} \times N) \times D}$  for graph attention. As for  $M_{n}^{(e)} \in \mathbb{R}^{L \times N \times D}$ , where  $L$  denotes the total number of episodic memory, we directly replace the least recently used patch in  $M_{n}^{(e)}$  with  $h_{i}^{m}$ . By using per-patch recall and update, we obtain information from different semantic levels. This makes the Semantic Memory more suitable for task generalization, while the Episodic Memory is more capable of capturing sample-specific details.
VSM-SCHA In SCHA-VAE and VSM-SCHA, we use a shared routing path between c and z:  $q_{\phi}(c_i^{(l)}|D_i,c_i^{(l + 1)},z_i^{(l + 1)},m_i^{(l)})$  where  $z_{l + 1}$  and  $c_{l + 1}$  are samples from layer  $l$ . We additionally remove any form of normalization (batch/layer norm) in the hierarchical formulation
Memory size Across all datasets, we set the size of each semantic memory block (category) in the VSM-based model to 1, and each episodic memory block to 5. For models that exclusively depend on episodic memory (CNS + Episodic, SCHA + Episodic, and Episodic-diffusion), we establish a memory size of 10000, with the episodic memory functioning as a first-in, first-out queue.
B Details of Model Design
C.1 Quantitative Comparison of Generative Samples
In this section, we apply VSM to FSDM, a conditional diffusion model where  $c$  is not a variational variable. To align with FSDM, we only use the following update rule for the semantic memory:
Figure 6: The process of memory recall mechanism in VSM-SCHA.
Figure 7: The process of memory recall mechanism in VSM-FSDM.
$$
\boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {n e w} \leftarrow \alpha \boldsymbol {\mu} _ {\mathbf {M} _ {n} ^ {(s)}} ^ {o l d} + (1 - \alpha) \bar {\boldsymbol {\mu}} _ {\mathbf {M} _ {n} ^ {(s)}}, \bar {\boldsymbol {\mu}} _ {\mathbf {M} _ {n} ^ {(s)}} = \sum_ {j = 1} ^ {J _ {n} + T} \beta_ {n, j} H _ {n, j}, \beta_ {n, j} = \frac {g \left(\mathbf {M} _ {n} ^ {(s)} , H _ {n , j}\right)}{\sum_ {j = 1} ^ {J _ {n} + T} g \left(\mathbf {M} _ {n} ^ {(s)} , H _ {n , j}\right)} \tag {18}
$$
Table 4 presents the results of VSM-FSDM on Fs-Cifar100 and celeba datasets. It can be observed that VSM performs well even without using variational latent variables, demonstrating its strong performance and scalability.
C.2 Additional Visualization of Generative Samples
More visualization results of generative samples can he found in Figures 8, 9, 10 and 11.
Table 4: Few-shot generative evaluation of diffusion-based models on various datasets with the set size 5, including Fs-Cifar100 and Celeba. (FID: Frechet score; sFID: spatial FID; P: precision; R: recall.)
FS-Cifar100	Celeba
FID ↓	sFID ↓	P ↑	R ↑	FID ↓	sFID ↓	P ↑	R ↑
DDPM	62.84	28.91	0.58	0.40	14.02	27.46	0.71	0.39
FSDM	35.07	20.95	0.62	0.53	10.11	24.58	0.73	0.39
VSM-FSDM	33.21	18.32	0.68	0.61	9.25	23.21	0.77	0.43
Table 5: Ablation study on diffusion-based models
FS-Cifar100	Celeba
vFSDM	FID ↓63.73	sFID ↓28.85	P ↑0.58	R ↑0.38	FID ↓13.52	sFID ↓27.75	P ↑0.71	R ↑0.38
vFSDM + semantic	60.21	27.63	0.58	0.42	12.98	26.01	0.73	0.39
vFSDM + episodic	52.13	25.25	0.59	0.44	12.36	25.97	0.72	0.41
VSM-diffusion	44.11	23.11	0.59	0.46	12.27	25.79	0.72	0.40
(a) Query Set
Figure 8: Visualization of few-shot generative samples using a VSM-Diffusion with the set size 5 on Celeba dataset. (Conditioning set and samples from in-distribution classes).
D Border Impact
The impact of our research can be highlighted in two key aspects:
For Few-shot Learning Task: Our structured memory has demonstrated a high level of effectiveness in storing and managing task-related information, enabling efficient transfer of prior knowledge from previous tasks to new ones. By combining semantic and episodic information storage, our model enhances its ability to capture both task-level information and fine-grained image-level details. Diverging from traditional few-shot models, our approach leverages the stored and recalled memory to extract more valuable insights while minimizing the influence of irrelevant information on new tasks.
For Deep Generative Models: Our modularized memory offers a seamless integration into a wide range of generative models, such as diffusion models and VAEs, while also easily merging with latent variables. This flexibility allows for efficient utilization of the memory component within the generative framework. In addition, our structured memory can be divided into multiple semantic levels, providing the ability to learn at various refined semantic layers. This hierarchical organization of memory facilitates capturing complex dependencies and relationships across different levels of
(a) Query Set
Figure 9: Visualization of few-shot generative samples using a VSM-Diffusion with the set size 5 on Celeba dataset. (Conditioning set and samples from out-distribution classes)
abstraction. By incorporating multiple semantic layers, generative models can potentially achieve enhanced performance and generate more meaningful and coherent outcomes.
E Limitation
The proposed model variant, VSM-Diffusion, is based on the denoise diffusion model, which requires multiple steps to generate samples. The framework can be extended to other diffusion models, such as implicit denoise diffusion models, to accelerate the generative process and improve its efficiency.
(a) Query Set
Figure 10: Visualization of few-shot generative samples using a VSM-Diffusion with the set size 5 on FS-Cifar100 dataset. (Conditioning set and samples from in-distribution classes).
(a) Query Set
Figure 11: Visualization of few-shot generative samples using a VSM-Diffusion with the set size 5 on FS-Cifar100 dataset. (Conditioning set and samples from out-distribution classes).