SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhibin Duan,\* Zhiyi Lv,\* Chaojie Wang, Bo Chen,
National Laboratory of Radar Signal Processing
Xidian University, Xi'an, Shaanxi 710071, China
xd_zhibin@163.com, xd_silly@163.com, bchen@mail.xidian.edu.cn
Aimed at adapting a generative model to a novel generation task with only a few given data samples, the capability of few-shot generation is crucial for many real-world applications with limited data, e.g., artistic domains. Instead of training from scratch, recent works tend to leverage the prior knowledge stored in previous datasets, which is quite similar to the memory mechanism of human intelligence, but few of these works directly imitate the memory-recall mechanism that humans make good use of in accomplishing creative tasks, e.g., painting and writing. Inspired by the memory mechanism of human brain, in this work, we carefully design a variational structured memory module (VSM), which can simultaneously store both episodic and semantic memories to assist existing generative models efficiently recall these memories during sample generation. Meanwhile, we introduce a bionic memory updating strategy for the conversion between episodic and semantic memories, which can also model the uncertainty during conversion. Then, we combine the developed VSM with various generative models under the Bayesian framework, and evaluate these memory-augmented generative models with few-shot generation tasks, demonstrating the effectiveness of our methods.
A remarkable capability of human intelligence is its ability to quickly grasp the concepts of new objects that it has never encountered before [1]. Essentially, this rapid adaptation is achieved by utilizing the past memories stored in the human brain, which can greatly improve the efficiency of later learning [2, 3]. Unlike human brain, existing machine learning methods often require to be trained on a large amount of data when facing new tasks, motivating the research interests on few-shot learning that aim at efficiently solving these unseen tasks with only a few given data samples [4-6]. While there have been promising processes for few-shot adaptation on classification tasks, less work has been done on few-shot generation [7-10], which is mainly due to the challenging nature of learning the generative process with only a few samples in an unsupervised manner [11, 12].
For existing deep neural networks, to improve their abilities of few-shot adaptation, one of the most effective methods is to introduce an augmented memory module, which can store the memory of past experiences to adapt the model to a new task quickly [6, 13-16]. However, the most widely used memory mechanism in existing works is short-term [17] with limited capacity, which focuses on the
collection of personal experience related to a particular place or time. To remove the limitation of short-term memory, Zhen et al. [18] introduce the concept of semantic memory, which allows the storage of general conceptual information into meta-learning to acquire long-term knowledge for few-shot learning. Although these memory modules have separately achieved attractive progress in few-shot learning [1, 2, 19, 20], there is little work to explore a memory mechanism to organically combine the semantic and episodic memories, and further apply it on few-shot generation tasks.
Inspired by the memory mechanism of the human brain, where both episodic and semantic memory work together to help humans understand the world around us and also accomplish creative tasks, in this work, we first develop a variational structured memory module (VSM), which can simultaneously store both episodic and semantic memories to improve the few-shot generation capability of the generative models. Then, we design two important processes in VSM: 1) structured memory recall, which retrieves the relevant information of previous tasks from VSM to be applied on current task; 2) structured memory update, which continuously collects data information from training tasks and gradually consolidates the structured memories stored in VSM. Finally, we combine VSM with a branch of few-shot latent variable models, specifically neural statistician (NS) [7], and develop a series of memory-augmented latent variable models for few-shot generation.
- Inspired by the memory mechanisms of the human brain, we develop a novel variational structured memory module (VSM) that can concurrently store both episodic and semantic memories, effectively aiding generative models in memory recall during sample generation.
- Meanwhile, we introduce a bionic memory updating strategy for the conversion between episodic and semantic memories stored in VSM, where episodic memory can provide detailed episodes that will be converted later and stored in semantic memory.
- Through incorporating VSM with NS, we develop a novel hierarchical memory-augmented neural statistician (HMNS), whose generative process can be specified with VAE or diffusion model, leading to VSM-VAE and VSM-Diffusion, respectively.
- Extensive experiments demonstrate that our developed VSM-VAE and VSM-Diffusion can achieve promising generative performance on few-shot generation tasks.
Few-Shot Latent Variable Models: For few-shot generation, a series of latent variable models based on NS [7] have been developed [9-12], which can be further improved by adopting a non-parametric formulation for the context [21], introducing a hierarchical learnable aggregation for the input set [11], increasing expressivity for the conditional prior [10], or exploring supervision [22]. Specifically, NS [7] introduces a context constant variable to specify the summary statistics of a dataset, which provides a shared prior for these samples within the same dataset, contributing to handling new datasets at test time; GMN [9] defines a variational model to learn a per-sample context-aware latent variable. Distinct from previous works, our method introduces a memory module to collect historical datasest information to inference the context of a new dataset.
Memory-Augmented Neural Networks: It has been shown that augmenting existing neural networks with an external memory module can improve their capability of few-shot learning [23]. Specifically, these memory-augmented neural networks mainly collect the information stored in the support set of the current task [6, 15, 16], which will be wiped form episode to episode [17] instead of maintaining long-term information that has been shown to be crucial for effectively learning new unseen tasks. For instance, variational semantic memory [18] is proposed to accumulate and store the semantic information from previous tasks for the inference of the prototypes of new tasks. Distinct from previous methods only constructing semantic memory for the prototype network, in this paper, we develop a structured memory module to store the historical information at both the episodic and semantic sides, which can largely enhance the model's capability of few-shot generation.
Figure 1: The generative process of (a) neural statistician (NS), (b) hierarchical neural statistician (HNS), (c) memory-augmented variational autoencoder (VSM-VAE), and (d) memory-augmented diffusion model (VSM-Diffusion), where  $\boldsymbol{x}_{i,j}$  denotes the observed data sample,  $z_{i,j}^{(l)}$  is the latent variable at the  $l$ -th hidden layer, and  $\boldsymbol{c}_i^{(l)}$  and  $\boldsymbol{m}_i^{(l)}$  are the context variable and memory variable of dataset  $D_i$ , respectively. More details can be found in Section 3.2.
3 Memory-Augmented Neural Statistician
Few-shot Generation: Few-shot generation [22] borrows a similar setting of few-shot learning [4] but for a more challenging task to generate objects given only a few samples of that object at test time. Specifically, given a set of training datasets  $D_{old} = \{D_i\}_{i=1}^N$ , where each dataset  $D_i = \{\pmb{x}_{ij}\}_{j=1}^{J_i}$  consists of  $J_i$  data samples, the goal of few-shot generation is to train a generative model  $p_\theta(\pmb{x})$  that is trained on given datasets  $D_{old}$  but can be quickly adapted to a new unseen task given only a few data samples collected from a new dataset  $D_{\mathrm{new}}$ .
3.2 Memory-Augmented Neural Statistician
Through incorporating memory mechanism with few-shot latent variable models, specifically NS [7] in Figure 1(a), in this paper, we develop a novel memory-augmented NS (MNS) for few-shot generation.
Generative Process of MNS: The design of MNS aims at introducing a memory module to store and update the data information at various semantic levels, which can be recalled when facing a new generative task. As shown in Figure 1(c), the generative process of hierarchical MNS is formulated as:
$$
p _ {\theta} \left(D _ {o l d}\right) = \iint \iint \prod_ {i = 1} ^ {N} \prod_ {j = 1} ^ {J _ {i}} p _ {\theta} \left(\boldsymbol {x} _ {i, j} \mid \boldsymbol {z} _ {i, j}, \boldsymbol {c} _ {i}, \boldsymbol {m} _ {i}\right) \tag {1}
$$
$$
\prod_ {i = 1} ^ {N} \prod_ {j = 1} ^ {J _ {i}} p _ {\theta} (\boldsymbol {z} _ {i, j} \mid \boldsymbol {c} _ {i}) \prod_ {i = 1} ^ {N} p _ {\theta} (\boldsymbol {c} _ {i} \mid \boldsymbol {m} _ {i}) \prod_ {i = 1} ^ {N} p (\boldsymbol {m} _ {i}) d \boldsymbol {z} _ {i, j} d \boldsymbol {c} _ {i} d \boldsymbol {m} _ {i}.
$$
To clarify the meanings of  $\pmb{z}_{i,j}$ ,  $\pmb{c}_i$ , and  $\pmb{m}_i$ :  $\pmb{z}_{i,j}$  denotes the per-sample latent variables of the  $j$ th sample in dataset  $D_i$ ,  $\pmb{c}_i$  denotes the task-specific context variable of dataset  $D_i$ , whose detailed definition can be found in [7], and  $\pmb{m}_i$  denotes the task-specific memory variable, which stores the data information of previous tasks that will be adapted to the new task.
3.3 Variants of Hierarchical Memory-Augmented Neural Statistician
Through specifying the generative process of HMNS and replacing its memory module with the variational structured memory (VSM) as described in Section 4, we can obtain a series of model
variants, such as VSM-HVAE and VSM-Diffusion, both of which can achieve promising generative performance on few-shot generation task.
3.3.1 VSM-VAE
Generative Process of VSM-VAE: To capture the hierarchical data information at various semantic levels, we extend the latent variables in Eq. (1) to a hierarchical version, specifically  $\boldsymbol{z}_{i,j} = \{z_{i,j}^{(l)}\}_{l=1}^{L}$ ,  $\boldsymbol{c}_{i,j} = \{c_{i,j}^{(l)}\}_{l=1}^{L}$ , and  $\boldsymbol{m}_{i,j} = \{m_{i,j}^{(l)}\}_{l=1}^{L}$ , which will be included in a variational autoencoder (VAE) with  $L$  hidden layers, leading to a VSM-VAE shown in Figure 1(c). Under the generative framework of MNS, using the same symbol notation in Eq. (1), the per-sample generative process of  $\boldsymbol{x}_{i,j}$  in the dataset  $D_i$  can be formulated as
$$
p _ {\theta} \left(\boldsymbol {x} _ {i, j}\right) = \int \int \int p _ {\theta} \left(\boldsymbol {x} _ {i, j} \mid \boldsymbol {z} _ {i, j}\right) p _ {\theta} \left(\boldsymbol {z} _ {i, j} \mid \boldsymbol {c} _ {i}\right) p _ {\theta} \left(\boldsymbol {c} _ {i} \mid \boldsymbol {m} _ {i}\right) p (\boldsymbol {m} _ {i}) d \boldsymbol {z} _ {i, j} d \boldsymbol {c} _ {i} d \boldsymbol {m} _ {i}, \tag {2}
$$
where  $p_{\theta}(\boldsymbol{z}_{i,j} \mid \boldsymbol{c}_i)$ ,  $p_{\theta}(\boldsymbol{c}_i \mid \boldsymbol{m}_i)$ , and  $p(\boldsymbol{m}_i)$  are formulated as the product of corresponding hierarchical latent variables as follows
$$
p _ {\theta} \left(\boldsymbol {z} _ {i, j} \mid \boldsymbol {c} _ {i}\right) = \prod_ {l = 1} ^ {L} p _ {\theta} \left(z _ {i, j} ^ {(l)} \mid z _ {i, j} ^ {(l + 1)}, c _ {i} ^ {(l)}\right), \quad p _ {\theta} \left(\boldsymbol {c} _ {i} \mid \boldsymbol {m} _ {i}\right) = \prod_ {l = 1} ^ {L} p _ {\theta} \left(c _ {i} ^ {(l)} \mid c _ {i} ^ {(l + 1)}, m _ {i} ^ {(l)}\right), \tag {3}
$$
$$
p _ {\theta} \left(\boldsymbol {m} _ {i}\right) = \prod_ {l = 1} ^ {L - 1} p _ {\theta} \left(m _ {i} ^ {(l)} \mid m _ {i} ^ {(l + 1)}\right) p \left(m _ {i} ^ {(L)}\right), \quad p \left(m _ {i} ^ {(L)}\right) = \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {4}
$$
Training Objective of VSM-VAE: For VSM-VAE, the training objective is to maximize the evidence lower bound (ELBO) of a set of training datasets  $\{D_i\}_{i=1}^N$ , which can be expressed as
$$
\mathcal {L} _ {\mathrm {E L B O}} = \sum_ {i = 1} ^ {N} \mathbb {E} _ {Q} \left[ \log p _ {\theta} \left(D _ {i} \mid \boldsymbol {c} _ {i}\right) \right] - \mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} \left(\boldsymbol {c} _ {i} \mid D _ {i} , \boldsymbol {m} _ {i}\right)}{p _ {\theta} \left(\boldsymbol {c} _ {i} \mid \boldsymbol {m} _ {i}\right)} \right] - \mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} \left(\boldsymbol {m} _ {i} \mid D _ {i} , \mathbf {M}\right)}{p _ {\theta} \left(\boldsymbol {m} _ {i}\right)} \right]. \tag {5}
$$
The first item in Eq. (5) can be further written as
$$
\mathbb {E} _ {Q} \left[ \log p _ {\theta} (D _ {i} \mid -) \right] = \sum_ {j = 1} ^ {J _ {i}} \mathbb {E} _ {Q} \left[ \log p _ {\theta} (\pmb {x} _ {i, j} | z _ {i, j} ^ {(1)}) \right] - \sum_ {l = 1} ^ {L} \mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} (z _ {i , j} ^ {(l)} \mid \pmb {x} _ {i , j} , z _ {i , j} ^ {(l + 1)} , c _ {i} ^ {(l)})}{p _ {\theta} (z _ {i , j} ^ {(l)} \mid c _ {i} ^ {(l)} , z _ {i , j} ^ {(l + 1)})} \right]
$$
$$
\mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} (\boldsymbol {c} _ {i} | -)}{p _ {\theta} (\boldsymbol {c} _ {i} | \boldsymbol {m} _ {i})} \right] = \sum_ {l = 1} ^ {L} \mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} \left(c _ {i} ^ {(l)} \mid D _ {i} , c _ {i} ^ {(l + 1)} , m _ {i} ^ {(l)}\right)}{p _ {\theta} \left(c _ {i} ^ {(l)} \mid c _ {i} ^ {(l + 1)} , m _ {i} ^ {(l)}\right)} \right], \tag {6}
$$
$$
\mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} (\pmb {m} _ {i} \mid -)}{p _ {\theta} (\pmb {m} _ {i})} \right] = \sum_ {l = 1} ^ {L} \mathbb {E} _ {Q} \left[ \log \frac {q _ {\phi} (m _ {i} ^ {(l)} \mid D _ {i} , \mathbf {M})}{p _ {\theta} (m _ {i} ^ {(l)} \mid m _ {i} ^ {(l + 1)})} \right],
$$
and  $Q$  denotes the variational distribution defined by the inference network of VSM-VAE
$$
Q = \prod_ {l = 1} ^ {L} q _ {\phi} \left(z _ {i, j} ^ {(l)} \mid \boldsymbol {x} _ {i, j}, z _ {i, j} ^ {(l + 1)}, c _ {i} ^ {(l)}\right) q _ {\phi} \left(c _ {i} ^ {(l)} \mid D _ {i}, m _ {i} ^ {(l)}\right) q _ {\phi} \left(m _ {i} ^ {(l)} \mid D _ {i}, \mathbf {M}\right), \tag {7}
$$
where  $\mathbf{M}$  indicates the memory module to store the semantic information collected from previous tasks, and will be introduced in Section 4.
3.3.2 VSM-Diffusion
Generative Model of VSM-Diffusion: Further, through specifying the generative process of HMNS with the diffusion model, we can obtain the VSM-Diffusion, as shown in Figure 1(d). Distinct from the hierarchy of latent variables in VSM-VAE, we share the same context and memory variables across all hidden layers of a diffusion model, specifically  $\boldsymbol{c}_i = c_i^{(l)}$  and  $\boldsymbol{m}_i = m_i^{(l)}$  for  $l \in \{1, \dots, L\}$ . Thus, under the generative formulate of MNS, using the same symbol notation in Eq. (1), the per-sample generative process of  $x_{i,j}$  in the dataset  $D_i$  can be formulated as
$$
p _ {\theta} (\pmb {x} _ {i, j}) = \int \int \int p _ {\theta} (\pmb {x} _ {i, j} | \{\pmb {x} _ {i, j} ^ {(l)} \} _ {l = 1} ^ {L}, \pmb {c} _ {i}) p _ {\theta} (\pmb {c} _ {i} | \pmb {m} _ {i}) p (\pmb {m} _ {i}) d \{\pmb {x} _ {i, j} ^ {(l)} \} _ {l = 1} ^ {L} d \pmb {c} _ {i} d \pmb {m} _ {i},
$$
where the backward diffusion process is consistent with that of conditional diffusion model [24, 25], formulated as
$$
p _ {\theta} \left(\boldsymbol {x} _ {i, j} \mid \left\{\boldsymbol {x} _ {i, j} ^ {(l)} \right\} _ {l = 1} ^ {L}, \boldsymbol {c} _ {i}\right) = p _ {\theta} \left(\boldsymbol {x} _ {i, j} \mid \boldsymbol {x} _ {i, j} ^ {(1)}, \boldsymbol {c} _ {i}\right) \prod_ {l = 1} ^ {L - 1} p _ {\theta} \left(\boldsymbol {x} _ {i, j} ^ {(l)} \mid \boldsymbol {x} _ {i, j} ^ {(l + 1)}, \boldsymbol {c} _ {i}\right) p _ {\theta} \left(\boldsymbol {x} _ {i, j} ^ {(L)}\right). \tag {8}
$$
Training objective of VSM-Diffusion: For VSM-Diffusion, the training objective is derived from maximizing the evidence lower bound (ELBO) of a set of training datasets  $\{D_i\}_{i=1}^N$ , which is consistent with Eq. (5). Similar to regular DDPMs, the first term can be decomposed into a summation of  $L$  terms, each of which can be computed independently, and can be estimated as
$$
\mathbb {E} _ {Q} \left[ \log p _ {\theta} (D _ {i} \mid \boldsymbol {c} _ {i}, \boldsymbol {m} _ {i}) \right] = \sum_ {j = 1} ^ {J _ {i}} \mathbb {E} _ {Q} \left[ - \log \frac {p _ {\theta} (\boldsymbol {x} _ {i , j} , \{x _ {i , j} ^ {(l)} \} _ {l = 1} ^ {L} | -)}{q _ {\phi} (\{x _ {i , j} ^ {(l)} \} _ {l = 1} ^ {L} \mid \boldsymbol {x} _ {i , j})} \right] = \sum_ {j = 1} ^ {J _ {i}} L _ {j, 0} + \sum_ {t = 2} ^ {L} L _ {j, l - 1} + L _ {j, L},
$$
where  $L_{L}$  is a fixed constant,  $L_{0}$  is the expectation of likelihood term and the intermediate term  $L_{l-1}$  can be formulated as
$$
L _ {j, l - 1} = \mathbb {E} _ {q _ {\phi} \left(x _ {i, j} ^ {(l)} \mid \boldsymbol {x} _ {i, j}\right)} \left[ \log \frac {q _ {\phi} \left(x _ {i , j} ^ {(l - 1)} \mid x _ {i , j} ^ {(l)} , \boldsymbol {x} _ {i , j}\right)}{p _ {\theta} \left(x _ {i , j} ^ {(l - 1)} \mid x _ {i , j} ^ {(l)} , \boldsymbol {c} _ {i}\right)} \right]. \tag {9}
$$
4 Variational Structured Memory
In this section, we will introduce the storage architecture of VSM in Section 4.1, equipped with the mechanism of memory recall and update in Sections 4.2 and 4.3, respectively. The basic intuition of VSM is inspired by the memory mechanism in the human brain: 1) there are two advanced forms of memory in the human brain: semantic memory allows the storage of general conceptual information and episodic memory allows the collection of detailed episodes; 2) episodic memory can provide detailed episodes that will be converted into conceptual information and stored in semantic memory. The structured memories stored in VSM will be recalled to infer the posterior of task-specific memory, formulated as:
$$
q _ {\phi} \left(\boldsymbol {m} _ {i} \mid D _ {i}, \mathbf {M}\right) \tag {10}
$$
which will be used in the inference process of either VSM-VAE or VSM-Diffusion.
4.1 Structured Memory Storage
As shown in Figure 7, there are  $N$  blocks (categories) in the VSM module, denoted as  $\mathbf{M} = \{\mathbf{M}_n\}_{n=1}^N$ , and each memory block  $\mathbf{M}_n$  consists of both semantic memory  $\mathbf{M}_n^{(s)}$  and episodic memory  $\mathbf{M}_n^{(e)}$ .
Semantic Memory Storage: The semantic memory in the human brain can provide conceptual information for quickly learning concepts of new object categories by seeing only a few data samples. Following Zhen et al. [18], for each block (category) of semantic memory in VSM, we try to store the conceptual information by averaging the latent representations of data samples belonging to this category, contributing to lightweight storage and fast lookup. To model the uncertainty, we model the  $n$ -th block (category) of semantic memory with a Gaussian distribution, formulated as:
$$
\mathbf {M} ^ {(s)} = \left\{\mathbf {M} _ {n} ^ {(s)} \right\} _ {n = 1} ^ {N}, \quad \mathbf {M} _ {n} ^ {(s)} = \mathcal {N} \left(\boldsymbol {\mu} _ {n}, \operatorname {d i a g} \left(\boldsymbol {\sigma} _ {n}\right)\right), \tag {11}
$$
where  $\pmb{\mu}_n, \pmb{\sigma}_n \in \mathbb{R}^d$  denotes the  $d$ -dimensional mean and variance vectors, respectively.
Episodic Memory Storage: Distinct from semantic memory allowing the storage of general conceptual information, episodic memory tends to focus on storing more detailed and vivid experiences or episodes. The design of episodic memory is inspired by previous works [26-30], and the main idea is to employ a query memory module to collect a subset of representative data samples for each block (category) in VSM. Specifically, the  $n$ -th block (category) of episodic memory is a set of  $T$  embedding vectors of data samples:
$$
\mathbf {M} ^ {(e)} = \left\{\mathbf {M} _ {n} ^ {(e)} \right\} _ {n = 1} ^ {N}, \quad \mathbf {M} _ {n} ^ {(e)} = \left\{\mathbf {M} _ {n, t} ^ {(e)} \right\} _ {t = 1} ^ {T}, \tag {12}
$$
where  $T$  is the storage limit of each block (category) of episodic memory, and the dimension of  $\mathbf{M}_{n,t}^{(e)}\in \mathbb{R}^d$  is the same as the latent representations of data samples.
Figure 2: The workflow of memory recall mechanism in variational structured memory (VSM) module  $\mathbf{M}$ , which stores both semantic memory  $\mathbf{M}^{(s)}$  and episodic memory  $\mathbf{M}^{(e)}$ .
4.2 Structured Memory Recall
Given a query dataset  $D_{i} = \{x_{i,j}\}_{j = 1}^{J_{i}}$ , which can be treated as a new dataset  $D_{new}$  for previous tasks, the mechanism of memory recall aims at effectively using the memory module to quickly adapt a model trained on previous tasks to a new unseen task. The workflow of structured memory recall in VSM has been depicted in Figure 7, which includes both the recalls of semantic memory  $\mathbf{M}^{(s)}$  and episodic memory  $\mathbf{M}^{(e)}$ , resulting in two kinds of task-specific memories, denoted as  $m_{i}^{(s)}$  and  $m_{i}^{(e)}$  respectively.
Semantic Memory Recall: For a new few-shot generation task on the dataset  $D_{i}$  ( $D_{new}$ ), the memory recall in VSM will first select the top-K relevant blocks (categories) of semantic memory as the candidates, denoted as  $\{\mathbf{M}_k^{(s)}\}_{k=1}^K$ , and then use an attention mechanism to aggregate these candidates into a task-specific semantic memory  $m_{i}^{(s)}$ , where the attention weight of  $n$ -th can be calculated as
$$
g \left(\mathbf {M} _ {n} ^ {(s)}, \boldsymbol {h} _ {i}\right) = \exp \left(\boldsymbol {\mu} _ {n} - \boldsymbol {h} _ {i}\right) ^ {2} / 2 \sigma_ {n} ^ {2}, \tag {13}
$$
where  $g(\cdot)$  can be treated as a similarity function, and  $h_i$  indicates the semantic representation of dataset  $D_i$  and can be obtained by averaging the latent representations of data samples in  $D_i$ .
Then, with the top-K relevant blocks of semantic memory  $\{\mathbf{M}_k^{(s)}\}_{k = 1}^K$  in hand, we can further aggregate these candidates with their attention weights, formulated as:
$$
\boldsymbol {m} _ {i} ^ {(s)} = \sum_ {k = 1} ^ {K} \lambda_ {i, k} \mathbf {M} _ {k} ^ {(s)}, \lambda_ {i, k} = \frac {g \left(\mathbf {M} _ {k} ^ {(s)} , \boldsymbol {h} _ {i}\right)}{\sum_ {k = 1} ^ {K} g \left(\mathbf {M} _ {k} ^ {(s)} , \boldsymbol {h} _ {i}\right)}, \mathbf {M} _ {k} ^ {(s)} \sim \mathcal {N} (\boldsymbol {\mu} _ {k}, \operatorname {d i a g} (\boldsymbol {\sigma} _ {k})), \tag {14}
$$
where  $\pmb{m}_i^{(s)}\in \mathbb{R}^d$  has gathered the relevant information from the semantic memory  $\mathbf{M}^{(s)}$  for the generation task on dataset  $D_{i}$
Episodic Memory Recall: After selecting the top-K relevant blocks of semantic memory  $\{\mathbf{M}_k^{(s)}\}_{k = 1}^K$ , we can further dive into recalling episodic memory, which is more specific and detailed. In practice, each block of episodic memory contains  $T$  embedding vectors of data samples  $\mathbf{M}_k^{(e)} = \{\mathbf{M}_{k,t}^{(e)}\}_{t = 1}^T$ , and we can finally obtain  $K*T$  embedding vectors of data samples. After that, following a similar process of recalling episodic memory, we can further aggregate these obtained embedding vectors with their attention weights as
$$
\boldsymbol {m} _ {i} ^ {(e)} = \sum_ {k = 1} ^ {K} \sum_ {t = 1} ^ {T} \boldsymbol {\lambda} _ {i, k, t} \mathbf {M} _ {k, t} ^ {(e)}, \quad \boldsymbol {\lambda} _ {i, k, t} = \frac {g \left(\mathbf {M} _ {k , t} ^ {(e)} , \boldsymbol {h} _ {i}\right)}{\sum_ {k = 1} ^ {K} \sum_ {t = 1} ^ {T} g \left(\mathbf {M} _ {k , t} ^ {(e)} , \boldsymbol {h} _ {i}\right)}, \tag {15}
$$
where  $g(\cdot)$  and  $h_i$  have been defined in Eq. (13), and the obtained  $\mathbf{m}_i^{(e)} \in \mathbb{R}^d$  has gathered the relevant information from the episodic memory  $\mathbf{M}^{(e)}$  for the generation task on dataset  $D_i$ .
4.3 Structured Memory Update and Consolidation