SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Saurav Jha
saurav.jha@unsw.edu.au
Dong Gong*
dong.gong@unsw.edu.au
Lina Yao
CSIRO's Data61, UNSW Sydney
lina.yao@data61.csiro.au
Continual learning (CL) aims to train deep neural networks efficiently on streaming data while limiting the forgetting caused by new tasks. However, learning transferable knowledge with less interference between tasks is difficult, and real-world deployment of CL models is limited by their inability to measure predictive uncertainties. To address these issues, we propose handling CL tasks with neural processes (NPs), a class of meta-learners that encode different tasks into probabilistic distributions over functions all while providing reliable uncertainty estimates. Specifically, we propose an NP-based CL approach (NPCL) with task-specific modules arranged in a hierarchical latent variable model. We tailor regularizers on the learned latent distributions to alleviate forgetting. The uncertainty estimation capabilities of the NPCL can also be used to handle the task head/module inference challenge in CL. Our experiments show that the NPCL outperforms previous CL approaches. We validate the effectiveness of uncertainty estimation in the NPCL for identifying novel data and evaluating instance-level model confidence. Code is available at https://github.com/srvCodes/NPCL.
Continual learning (CL) aims to help deep neural networks (DNNs) learn from a stream of nonstationary tasks by retaining the previously acquired knowledge [54, 35]. To achieve this, CL agents target alleviating the catastrophic forgetting issue with restricted computational and memory costs [42]. This requires balancing the plasticity for new knowledge with the stability for old [37].
To handle forgetting in CL, experience replay (ER) methods [30, 6] are one effective way to train DNNs on a memory buffer with a subset of the past tasks' experiences. Other than the ER methods, many regularization-based approaches have been proposed to penalize the forgetting on the DNNs' parametric [30] or representation spaces [5, 4]. However, these may still suffer from interference due to the regularization on the entire parameter space [7]. To address this, parameter isolation methods [53, 33] define task-specific training components but are usually confined to task incremental CL setups requiring task ID during testing [47]. It is thus challenging for CL agents to maintain transferable and shareable knowledge. Lastly, a hurdle to the real-world deployment of CL agents is their inability to measure predictive uncertainties, which impacts the potential utilization of CL across various practical applications, particularly those with critical safety considerations [31].
$^{*}$ D. Gong is the corresponding author.
To tackle the above issues, we propose to explore CL models using neural processes (NPs) [13, 14], a class of meta-learners that model tasks as data-generating functions from a stochastic process. NPs learn a prior over functions by marginalizing over a set of data points, or context, thus enabling rapid adaptation to new observations through inference on functions. Additionally, their probabilistic nature endows them with reliable uncertainty quantification capabilities [14, 27, 24, 25]. Our motivations to explore NPs for CL are thus two-fold. First, NPs exploit Bayes' theorem, which naturally enables CL through sequential posterior construction. Namely, NPs perform inference over the function space by learning context-based priors, which are updated to posteriors upon observing (additional) targets. Second, NPs meta-learn input correlations through a set of latent variables, which could be a key to meta-learn knowledge transfer across multiple correlated tasks. However, NPs face challenges in directly addressing CL tasks, given that (a) the reliance on a single global latent leads to suboptimal modeling of complex CL tasks where multiple correlated tasks could occur simultaneously, (b) NPs cannot directly handle the forgetting of past task correlations arising from the non-static data stream.
To address the above desiderata, we propose Neural Processes for Continual Learning (NPCL), a hierarchical latent variable model with a global latent variable to capture inter-task correlation and task-specific latent variables for finer knowledge. Fig. 1 shows the NPCL exploiting functional correlation among current and past task training samples of ER. The drift of global and past task-specific distributions away from their original forms is the major cause of forgetting in the NPCL. We thus propose to regularize the latent variables to be similar to their old forms and show the merits of regularization over typical parameter-based regularization. We then leverage the uncertainty encoded by the NPCL for the aforesaid CL challenge of task head inference. To this end, we propose using entropy as an uncertainty quantification metric (UQM). The NPCL outperforms previous probabilistic CL models
and delivers better or comparable results than state-of-the-art deterministic CL methods, which usually have an edge over their probabilistic counterparts in terms of accuracy. Moreover, our ablations show the enhanced efficacy offered by the NPCL on continual learning settings requiring model calibration and few-shot replay. To study the further usages of the NPCL's uncertainty estimation, we show its out-of-the-box readiness for novel data detection and instance-level confidence evaluation [17]. Lastly, we list the key limitations of the NPCL as an attempt to lay further solid directions for uncertainty-aware continual learning.
Figure 1: Neural Processes for Continual Learning (NPCL): each training step involves minimizing the distance  $D$  between the context-based prior and the target-based posterior, alongside regularizing the task-specific and global distributions towards their old forms.
Continual Learning (CL). Existing CL methods address catastrophic forgetting through three major approaches: (a) Regularization-based methods penalize changes in a model's important weights for previous tasks, such as Elastic Weight Consolidation (EWC) [30], Synaptic Intelligence (SI) [54], etc. (b) Parameter Isolation-based methods partition the network's parameters to specialize on individual tasks, e.g., Douillard et al. [9] learning task-specific tokens for Transformers. (c) Replay-based methods use an episodic memory to preserve a fraction of the past tasks' experience for preventing forgetting while learning on new tasks; e.g., experience replay (ER) [6] storing past inputs, dark experience replay (DER) [4] storing past logits, and Yan et al. [52] using a loss-aware memory for ER. Our method uses (a) via regularization of distributions, (b) via task-specific latent heads, and (c) via replay of past task inputs and distributions.
Neural Processes (NPs). NPs were introduced to meta-learn a distribution of a family of functions modeling the data-generating process through their deterministic [13] and/or latent summaries [14]. Attentive NPs (ANPs) [27] replaced the averaging operation in NPs with a dot-product attention [48] to enhance their expressivity. NPs/ANPs rely on a global latent that limits their ability to
model observations from multiple functions. Some works address this through local latent variables that model fine-grained correlation among a subset of the observations [50]. Recently, multi-task processes (MTPs) [26] have been studied to model multiple tasks with NPs, owing to the hierarchy of task-specific latent variables conditioned on a global latent. However, existing MTPs cannot directly handle CL problem because (a) MTPs are not designed to learn on sequential tasks and thus do not handle forgetting; (b) MTPs target the multi-task learning problem where the label for an input spans the exhaustive output space of available tasks unlike the CL setup where each input may belong to one specific and unknown task, out of multiple seen tasks.
Besides, the added complexity of variational inference has limited NP applications to mostly proof-of-concept focused regression tasks [23]. The potential of NPs for large-scale classification tasks thus remains largely under-explored except in some recent works. Wang et al. [49], for instance, leverage the predictive uncertainties of NPs to decide on pseudo labels for unlabeled data in semi-supervised classification. In our work, we use NPs to handle CL with classification tasks, reflecting the benefits of principled Bayesian learning, uncertainty estimation, and easily integrated existing ER.
3 Preliminaries: Neural Processes
Given the data  $\{(x^t,y^t)\} = (X^t,Y^t)\sim \mathcal{D}^t$  of a task  $t$ , the goal is to learn the mapping  $F_{*}^{t}:X^{t}\to Y^{t}$  reflecting the data-generating process. NPs [13, 14] meta-learn the distribution over the mapping functions from the given tasks. This is equivalent to meta-learning the distribution over the predictions  $p(y_i^t |x_i^t,\mathcal{C}^t)$  for the target output  $y_{i}^{t}$  belonging to a target data set  $\mathcal{T}^t$ , given the corresponding target input  $x_{i}^{t}$  and a context set  $\mathcal{C}^t$  [14, 23]. To reflect the meta-learning behavior of NPs [14], the training samples are split into the context set  $|\mathcal{C}^t| = m$  and a target set  $|\mathcal{T}^t| = m + n$  containing context set  $\mathcal{C}$  and additional samples. NPs learn the Gaussian priors and posteriors using a neural network  $F_{[\phi ;\theta ]}^{t}\approx F_{*}^{t}$  for the predictive distribution, where  $\phi$  and  $\theta$  parameterize an encoder  $q$  and a decoder  $p$  respectively. This involves deriving a global variable  $z^{G}$  to estimate the prior  $p(z^{G}|\mathcal{C}^{t};\phi)$ , and then maximizing the marginal likelihood  $p(Y_{\mathcal{T}}^{t}|\mathcal{C}^{t},X_{\mathcal{T}}^{t};\theta)$ :
$$
p \left(Y _ {\mathcal {T}} ^ {t} \mid X _ {\mathcal {T}} ^ {t}, \mathcal {C} ^ {t}\right) = \int p \left(Y _ {\mathcal {T}} ^ {t} \mid X _ {\mathcal {T}} ^ {t}, z ^ {G}\right) p \left(z ^ {G} \mid \mathcal {C} ^ {t}\right) d z ^ {G}, \tag {1}
$$
where  $p(Y_{\mathcal{T}}^{t}|X_{\mathcal{T}}^{t},z^{G}) = \prod_{i = 1}^{m + n}p(y_{i}^{t}|x_{i}^{t},z^{G})$  is the generative likelihood. In CL with streaming tasks, maintaining the memorization of the task prior  $p(z^{G}|\mathcal{C}^{t};\phi)$  can help NPs avoid forgetting the  $t$ -th task. Our aim behind enabling NP for CL is to seek a trade-off to preserve such task priors while sharing the parameters among tasks.
4 Continual Learning with Neural Processes
CL considers learning from a series of different tasks arriving sequentially, i.e.,  $\mathcal{D}^t\mid 0\leq t\leq T - 1$  Here,  $\mathcal{D}^t$  can belong to classification tasks with different classes in class incremental CL [8]. Let  $l$  be the cross-entropy (CE) loss for classification, the CL objective for the task  $t$  involves minimizing:
$$
\mathcal {L} _ {C E} ^ {t} = \mathbb {E} _ {(x, y) \sim \mathcal {D} ^ {t}} l \left(F _ {[ \phi ; \theta ]} (x), y\right) \tag {2}
$$
on all  $[0, t]$  tasks seen sequentially. Achieving Eq. (2) is challenging in real-world CL scenarios, where the previous datasets can be unavailable due to constraints on privacy, storage, etc. Learning on the sequential data with varying distributions causes catastrophic forgetting. To alleviate the issue, experience replay (ER) is used in CL to store and periodically revisit some past experiences, e.g., samples  $(x^{t}, y^{t})$  of task  $t$ , in a small episodic memory  $\mathcal{M}$  for replay in the future [8, 6]. In this work, we develop our method with the classical reservoir sampling-based ER [6] for a task boundary-agnostic updating of  $\mathcal{M}$ .
CL methods with ER solely still suffer from severe forgetting issues [6, 52, 8]; and jointly optimizing parameters on  $\mathcal{D}^t$  and  $\mathcal{M}$  has several drawbacks [35, 4]. Considering that a deterministic mapping  $F^t$  limits capturing the randomness behind the real-world data in a stream, to utilize the meta-learning ability of NPs, we next propose extending models with Eq. (2) and Eq. (1) to arrive at our NPCL model. It allocates small subsets of parameters to learn robust per-task and global priors and uses stochastic factors to meet data-driven challenges such as deducing the right parameters for inference.
Figure 2: Overview of the NPCL architecture: the decoding mechanism differs during training and inference. Red, Cyan, and Orange denote three different tasks.
4.1 Neural Processes for Continual Learning
Given the task in a stream, we model the CL task based on NPs formulated in Eq. (1). In ER framework with a small memory buffer, where the context and target could be from tasks indexed by  $t$ , Eq. (1) can be extended to derive the joint posterior for NPs [14] as:
$$
p \left(Y _ {\mathcal {T}} ^ {0: t} \mid \mathcal {C} ^ {0: t}, X _ {\mathcal {T}} ^ {0: t}\right) = \int p \left(Y _ {\mathcal {T}} ^ {0: t} \mid X _ {\mathcal {T}} ^ {0: t}, z ^ {G}\right) p \left(z ^ {G} \mid \mathcal {C} ^ {0: t}\right) d z ^ {G}, \tag {3}
$$
where  $z^{G}$  models the joint distribution  $F_{*}^{0:t}$  of CL tasks and is an enabler of the knowledge transfer [35] between these (see App. A.3 for ELBO). Eq. (3) poses two challenges. First, a labeled context  $\mathcal{C}$  is needed for inferring predictions as all NPs, which is unprepared in CL setups by default. To overcome this, we use the memory  $\mathcal{M}$  offered by the ER-based setups as context during inference. Second, jointly modeling  $F_{*}^{0:t}$  ignores the dynamics of per-task stochasticities and is still prone to the bottlenecks of Eq. (2). We address the issue by introducing hierarchical modeling and redefining Eq. (3) in the following.
4.2 NPs with Hierarchical Task-specific Priors for CL
To learn informative task priors while tackling the forgetting issue in CL, we propose a hierarchical modeling of the NP model. We preserve the global latent  $z^{G}$  to induce the direct knowledge transfer and add the task-specific upon the global variable to enhance the capturing of task-specific knowledge in CL. We thus extend Eq. (3) with task-specific latent variables  $z^{t} = (z^{0},\dots,z^{t})$ . As a result, our posterior is a two-step hierarchical latent variable model (Fig. 2) where the global and the per-task latent variables model the inter and intra-task correlations, respectively:
$$
p \left(Y _ {\mathcal {T}} ^ {0: t} \mid X _ {\mathcal {T}} ^ {0: t}, \mathcal {C} ^ {0: t}\right) = \iint \left[ \prod_ {t = 0} ^ {T - 1} p \left(Y _ {\mathcal {T}} ^ {t} \mid X _ {\mathcal {T}} ^ {t}, z ^ {t}\right) p \left(z ^ {t} \mid z ^ {G}, \mathcal {C} ^ {t}\right) \right] p \left(z ^ {G} \mid \mathcal {C} ^ {0: t}\right) d z ^ {0: t} d z ^ {G}, \tag {4}
$$
where the entire context  $\mathcal{C}^{0:t}$  is first encoded into  $z^G$  and then conditioned on  $z^G$ , the task-specific context  $C^t \coloneqq (C^0,..,C^t)$  are encoded into their respective latent variables. We refer to Eq. (4) as NP for CL (NPCL). The hierarchical modeling enables NPCL to learn the shareable knowledge via  $z^G$  and the task-specific knowledge via  $z^t$  in the meta-learning fashion of NPs. Task identity is used in training to specify the task-specific latent variables. Unlike MTP [26] making predictions of all tasks for all inputs, NPCL needs to specify the corresponding output space for each test sample. We further discuss the relationship between NPCL and NP-based meta-learning in App. B.
4.3 The NPCL Architecture
As standard NPs [14, 27], the training samples are split into a context  $\mathcal{C}$  and a target set  $\mathcal{T}$  containing  $\mathcal{C}$  and additional samples. Given the inputs  $x_{i}$  from  $\mathcal{C}$  or  $\mathcal{T}$ , we first pass these to a feature extractor
$f$ . With a slight abuse of notation, we denote the features as  $x_{i}: x_{i} \in \mathbb{R}^{|f|}$  and let  $|f|$  denote the dimension.  $x_{i}$  concatenated with the one-hot encoded labels, i.e.,  $[x_{i}; y_{i}]$ , is fed to the NPCL encoder with a deterministic and a latent path, and then to the decoder (Fig. 2). All the NPCL layers use multi-layer perceptrons (MLPs) projections, i.e.,  $\mathrm{MLP}(x): \mathbb{R}^{|f|} \to \mathbb{R}^{|o|}$ , where  $|o|$  is the output feature dimension as a hyperparameter. We denote a normal distribution with a mean  $\mu$  and a variance  $\sigma^2$  by  $\mathcal{N}(\mu, \sigma^2)$ ; the global and the task-specific distributions are  $\mathcal{N}(\mu_G, \sigma_G^2)$  and  $\mathcal{N}(\mu_t, \sigma_t^2)$ .
Latent Encoder. The latent path comprises of the projection  $\Phi_i^{\mathrm{lat}} = \mathrm{MLP}([x_i; y_i])$  followed by two attention operations [48]. First, per-task projections form the keys, values and queries to taskwise self-attention layers  $SA_{\mathrm{lat}}^t$  that produce order-invariant encodings  $s_i^t$  over the samples of task  $t$ . Second, all encodings  $\{s_i^{0:t}\}_{i=1}^{n+m}$  serve as the keys, values and queries to cross-attention layers  $CA_{\mathrm{lat}}^{0:t}$  that enrich their order-invariance from intra-task  $s^t$  to inter-task  $s^G$ .  $s^t$  and  $s^G$  are used to derive the  $N$  and  $M$  Monte Carlo samples of the global  $z^G$  and the task-specific latent variables  $z^t$ , respectively (see App. C for more details) using the reparameterization trick [29]. We set  $M = 1$  to enhance the inter-task stochasticity in posterior while retaining superior computational efficiency (see App. E). For each input, we thus get  $N \times (t + 1)$  latent outputs.
Deterministic Encoder. The deterministic path is similar to that of the ANP [27] and outputs an order-invariant representation  $r_*$  for target  $x_*$  (see App. C).
Decoder. Based on the task information, the decoder adopts separate mechanisms during training and inference. At train time, we use the available task labels to filter the  $N$  true latent variables  $\{z_i^t\}_{i = 1}^N$ , combine them with  $r_*$  and  $x_{*}$ , and decode the logits  $h_*$ . We discuss the decoding operation in the testing phase without task ID in Sec. 4.5.
4.4 Learning Objectives for the NPCL
The learning of the NPCL involves variational inference alongside additional regularizations.
Evidence Lower Bound (ELBO). The intractability of Eq. (4) leads us to the following ELBO:
$$
\begin{array}{l} \log p _ {\theta} \left(Y _ {\mathcal {T}} ^ {0: t} \mid X _ {\mathcal {T}} ^ {0: t}, \mathcal {C}\right) \geq \mathbb {E} _ {q _ {\phi} (z | \mathcal {T})} \left[ \sum_ {t = 0} ^ {T - 1} \mathbb {E} _ {q _ {\phi} \left(z ^ {t} \mid z ^ {G}, \mathcal {C} ^ {t}\right)} \left[ \log p _ {\theta} \left(Y _ {\mathcal {T}} ^ {t} \mid X _ {\mathcal {T}} ^ {t}, z ^ {t}\right) \right] \right. \tag {5} \\ \left. \right.\left. - D ^ {t} \left(q _ {\phi} \left(z ^ {t} \mid z ^ {G}, \mathcal {T} ^ {t}\right) \| q _ {\phi} \left(z ^ {t} \mid z ^ {G}, \mathcal {C} ^ {t}\right)\right)\right] - D ^ {G} \left(q _ {\phi} \left(z ^ {G} \mid \mathcal {T}\right) \| q _ {\phi} \left(z ^ {G} \mid \mathcal {C}\right)\right), \\ \end{array}
$$
where  $p_{\theta}(Y_T^t |X_T^t,z^t)$  is approximated by the CE loss.  $D^{t}$  and  $D^{G}$  denote the KL divergence (KLD) between the approximate posterior and prior for the task-specific and global distributions, respectively. We derive the ELBO in App. A.1. We next propose two techniques to counter forgetting in the NPCL. Henceforth, we use  $D$  to denote the Jenshen-Shannon (JS) divergence [11] between two distributions.
Global Regularization (GR). The training data of a CL task  $t$  is dominated by the t-th task samples. For the NPCL, this drifts the global distribution  $\mathcal{N}(\mu_G^t,\sigma_G^t)$  of past tasks towards the new task (Fig. 1). We thus regularize their global distribution using the one learned at step  $t - 1$ :
$$
\mathcal {L} _ {\mathrm {G R}} = D \left(\mathcal {N} \left(\mu_ {G}, \sigma_ {G} ^ {2}\right) _ {t}, \mathcal {N} \left(\mu_ {G}, \sigma_ {G} ^ {2}\right) _ {t - 1}\right) \tag {6}
$$
Task-specific Regularization (TR). While GR helps preserve the joint distribution of the past tasks, the hierarchy in the NPCL leaves their task-specific distributions to be still prone to forgetting (Fig. 3(a)). This can further amplify the posterior collapse [46] for past task-specific latent variables during CL training (Fig. 3(b)). To alleviate these, we regularize the learning of previous task distributions as:
$$
\mathcal {L} _ {\mathrm {T R}} ^ {t} = D \left(\mathcal {N} \left(\mu_ {t}, \sigma_ {t} ^ {2}\right) _ {t}, \mathcal {N} \left(\mu_ {t}, \sigma_ {t} ^ {2}\right) _ {j}\right), \tag {7}
$$
where  $j$  is the step at which the task  $t$  arrived. Given the reliance of Eq. (6) and Eq. (7) on past distributions, we maintain a separate buffer, which we refer to as the distribution memory  $\mathcal{M}_{\mathcal{N}}$ , to store the global  $\mathcal{N}(\mu_G,\sigma_G^2)$  and the task-specific distributions  $\mathcal{N}(\mu_{0:t - 1},\sigma_{0:t - 1}^2)$ .  $\mathcal{M}_{\mathcal{N}}$  is updated after each incremental training step, where we run an additional pass over the training data of task  $t$  alongside replaying  $\mathcal{M}$  to record the batchwise averaged global and task-specific means and variances.
Integrated objective. Using  $\alpha$ ,  $\beta$ ,  $\gamma$ , and  $\delta$  to denote the loss weights, our total loss can be written as:
$$
\mathcal {L} = \frac {1}{| \mathcal {D} ^ {t} | + | \mathcal {M} |} \sum_ {\left(x ^ {t}, y ^ {t}\right) \in \mathcal {D} ^ {t} \cup \mathcal {M}} \left(\mathcal {L} _ {\mathrm {C E}} + \alpha D ^ {t} + \beta D ^ {G}\right) + \frac {1}{| \mathcal {M} |} \sum_ {\left(x ^ {t}, y ^ {t}\right) \in \mathcal {M}} \gamma \mathcal {L} _ {\mathrm {G R}} + \delta \mathcal {L} _ {\mathrm {T R}} ^ {t}, \tag {8}
$$
(a) Drift of past task distributions
(b)  $\log K' L(q|p) \, \mathrm{w/o TR}$
Figure 3: Analyses on the need for distribution regularization: (a) shows the increasing distances between current distributions of past tasks and their original distributions (learned while the tasks were introduced). (b) and (c) show the effect of global (GR) and task regularization (TR) on the activation of the global and task-specific latent units. Low KLD corresponds to an inactive unit.
(c)  $\log KL(q|p)$  w/TR
where  $\mathrm{CE}$ ,  $D^t$ , and  $D^G$  act on the current task data  $\mathcal{D}^t$  and on the buffer  $\mathcal{M}$  while GR and TR act only on  $\mathcal{M}$ . By setting  $0 < \{\alpha, \beta, \gamma, \delta\} < 1$ , we resort to using the (respective) cold posteriors [55].
4.5 Inference with Uncertainty Awareness
NPCL's inference uses  $f$  to obtain the features  $x_{*}$  for the target test images. Although the task identification information is used to train the task-specific module, task identification of test samples is usually unavailable in general real CL tasks (except the restricted task-incremental setting [8]). Given  $x_{*}$  from the encoder, this leaves us with  $\{z_i^{0:t}\}_{i=1}^N$  possible modules and the corresponding latent variables to use and infer for obtaining the prediction. A naive solution is to average over  $N*(t+1)$  logits. But as the number of tasks grows, the noise from incorrect task priors would dominate the posterior. We thus propose using entropy as an uncertainty quantification metric (UQM) to filter the logits of the true task head  $\psi^t$ :
$$
h _ {*} = \underset {j \in [ 1, t ]} {\arg \min } U \left(h _ {\psi j}\right), \quad U \left(h _ {\psi}\right) = - \sum_ {i \in N} \delta (i) \log (\delta (i)), \tag {9}
$$
where  $\delta$  is the softmax function and  $U$  is the total Shannon entropy [45] over the  $N$  logits per head. As we use true head  $\psi_{\phi}^{t}$  during training,  $\phi \in \psi^t$  produces low entropy for within distribution data. In light of Eq. (9), the NPCL can be seen as a special case of the mixture-of-expert (MoE) modeling [36, 51], where we leverage uncertainty to select the top-1 expert during inference.