SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We dub this OOD score tailored to EBML Energy Sum, whose full derivation is deferred to Appendix A.2. This energy sum enjoys not only the theoretical advantage, i.e., being provably proportional to the negative log-likelihood of a task, but also simple computation benefits. During meta-testing, evaluating the score of Eqn. (9) for each task requires only a single forward pass of the support set samples through the two energy functions.
More remarkably, the energy sum is intuitively appealing in the sense that it characterizes (1) how far a task is from the overall ID meta-training task distribution via the latent prior energy score  $E_{\lambda}$  and (2) how difficult it is to predict the observed support set conditioned on  $\phi_{i}$  via the task-specific data energy score  $E_{\omega}$ . First, the terms in the last line of Eqn. (8) for learning the latent prior EBM altogether correspond to maximizing the likelihood  $\mathbb{E}_{\mathcal{T}_i \sim p(\mathcal{T})} \mathbb{E}_{q_{\psi}(\phi_i | \mathcal{T}_i^s)}[\log p_{\lambda}(\phi_i)]$ , which enforces the latent prior energy score  $E_{\lambda}$  to capture the overall ID meta-training distribution. As illustrated in
Figure 2 (right), the further away a task is from the overall ID meta-training distribution measured in Euclidean distance, the larger the energy score  $E_{\lambda}$  is as expected. Second, conditioned on even the ID latent prior  $\phi_{i}$ , those tasks with support samples as scattered as possible are especially difficult to predict. These tasks are considered to be OOD, as evidenced in higher values of  $E_{\omega}$  in Figure 2 (left).
4.3 EBML for OOD Generalization
The Energy Sum proposed in Section 4.2 develops OOD awareness of a meta-testing task, based on which we differentiate our meta-testing procedures for effective meta-generalization.
Meta-testing for ID tasks Given the support set  $\mathcal{T}^s$  of a meta-testing task, prediction of the label for its query  $\mathbf{x}_j^q$  amounts to maximizing our approximated log-likelihood (see Eqn. (5)) of the task, i.e.,
$$
y _ {j} ^ {q} = \underset {y} {\arg \min } \mathbb {E} _ {\phi \sim q _ {\psi} (\phi \mid \mathcal {T} ^ {s})} \left[ E _ {\omega} \left(\mathbf {x} _ {j} ^ {q}, y, \phi\right) + E _ {\lambda} (\phi) \right]. \tag {10}
$$
Provided that the task has already been identified within the ID region, the second energy  $E_{\lambda}(\phi)$  is negligibly small. Consequently, we reduce the above optimization problem to consider only the first term  $E_{\omega}(\mathbf{x}_j^q,y,\phi)$ , and solve it via gradient descent. We provide the pseudo codes in Appendix E.
Meta-testing for OOD tasks For an OOD task, its meta-learned prior  $\phi \sim q_{\psi}(\phi|\mathcal{T}^s)$  is located out of the ID meta-training task distribution and likely loses its effectiveness. We seek a solution that adapts this inadequate meta-learned prior back to the ID region, so as to make the most of the ID latent priors with guaranteed meta-generalization. This shares the idea with classifier editing in [45], where the editing parameters are trained to map an OOD image to an ID one for improving generalization. Therefore, we introduce task-specific parameters  $\zeta$  which are optimized via the following,
$$
\underset {\zeta} {\arg \min } \mathbb {E} _ {\boldsymbol {\phi} \sim q _ {\psi \cup \zeta} (\boldsymbol {\phi} | \mathcal {T} ^ {s})} \left[ \sum_ {j = 1} ^ {N ^ {s}} E _ {\omega} \left(\mathbf {x} _ {j} ^ {s}, y _ {j} ^ {s}, \boldsymbol {\phi}\right) + \max  \left(E _ {\lambda} (\boldsymbol {\phi}) - m, 0\right) \right], \tag {11}
$$
where  $m$  is a hyper-parameter. We find that setting  $m$  as the empirical average of the latent prior energy over all ID training tasks works well in practice, i.e.,  $m = \mathbb{E}_{p_{ID}}[\mathbb{E}_{\phi_i\sim q_\psi (\phi_i|\mathcal{T}_i^s)}[E_\lambda (\phi_i)]]$ .
As a result of optimizing the second term in Eqn. (11), the task-specific parameters  $\zeta$  enable  $q_{\psi \cup \zeta}(\phi|\mathcal{T}^s)$  to accommodate for OOD tasks by mapping the meta-learned prior back to ID meta-training tasks; while optimizing the first term preserves the data-level predictive ability of the model. We highlight that the task energy minimization approximates the minimization of a KL divergence between the task-specific posterior and the meta-learned prior, thereby inducing a metaregularization effect during adaptation. See Appendix A.3 for details. Eventually, we use the adapted task-specific parameters for final prediction on query samples as in Eqn. (10). Pseudo code for the EBML adaptation and inference algorithms described above can be found in Appendix E.
In Figure 3, we visualize the adaptation process when optimizing Eqn. (11) for OOD few-shot classification tasks in Meta-dataset [49]. As the prior
energy of these OOD tasks decreases, their  $\phi_{i}$  gradually shift towards to the ID region as desired. Within this region, minimizing the first term in Eqn. (11) continuously improves generalization. In contrast, given only a few support samples, existing SOTA methods that solely rely on feed-forward inference [1] and gradient-based optimization [28] for OOD task adaptation without a prior are both prone to overfitting. We provide more empirical evidence on this in Appendix C. On the other hand, meta-learning a BNN, which imposes a prior distribution on the parameter space during adaptation may be computationally cumbersome and often lead to sub-optimal performance in comparison to their non-Bayesian counterparts.
Figure 3: Illustration of the OOD task adaptation process on OOD domains of the metadata set [49] where each dot in (a) represents an OOD task in latent space  $\phi$ . Minimizing Eqn. (11) leads to (a) the latent  $\phi$  of the OOD task moving to the ID region (contour plot), (b) the Euclidean distance between class prototypes enlarging, and consequently (c), the classification accuracy on query samples increasing.
In the experiments, we test EBML on both few-shot regression and image classification tasks in search for answers to the following key questions: RQ1: Whether the improved expressiveness of EBML over traditional Bayesian meta-learning methods can lead to a more accurate model of the meta-training ID task distribution, hence a more reliable OOD task detector. RQ2: Whether Energy Sum can be an effective score for detection of OOD meta-testing tasks. RQ3: Whether EBML instantiated with SOTA algorithms can exploit the meta-learned EBM prior in OOD task adaptation to achieve better prediction performance on OOD tasks.
We now discuss two instantiations of the EBML framework with SOTA meta-learning algorithms for regression and classification. We illustrate our approach in Figure 4 below and defer a more detailed description for our models to Appendix B.
Figure 4: Overview of the EBML framework. The task latent variable  $\phi_{i}$  is inferred from the support set  $\mathcal{T}_i^s$  following the implementation of the base algorithm. The data and task energy scores are evaluated by the data and prior EBMs  $E_{\omega_1}$  and  $E_{\lambda}$ , respectively; while the query labels are predicted by the classifier  $p_{\omega_2}$  of the base algorithm.
Regression. Take CNPs [8] as an example base model. CNPs implements  $q_{\psi}(\phi_i|\mathcal{T}_i^s)$  as a neural network encoder that outputs a function embedding in finite vector form, i.e.,  $\phi_i \in \mathbb{R}^D$ , from a given support set,  $\mathcal{T}_i^s$ . That said, we let the prior EBM to model the empirical distribution over such finite-dimension function embedding, i.e.,  $E_{\lambda}(\phi_i): \mathbb{R}^D \to \mathbb{R}$ .
Classification Many cross-domain few-shot classification algorithms [28, 42, 1] rely on a metric-based classifier for prediction, which assigns query sample to the class with nearest prototype to the query representation based on some distance measure. In these cases, it is natural to specify the task-specific latent  $\phi_{i}$  as the set of class prototypes in each ID training task. Since  $\phi_{i}$  is a set of variables, we build the prior EBM model as a permutation-invariant neural network function. Suitable choices include DeepSets [56] and set transformer [26].
To align with the state-of-the-art prediction performance, we follow the practice in [50, 37] to train another decoder  $\omega_{2}$  with the loss function (e.g., cross entropy) in the base meta-learning model, which serves as a surrogate for  $E_{\omega}(x_j^q,y,\phi)$  in Eqn. (10) and Eqn. (11). We use this decoder for prediction.
Baseline Models For regression, we compare against: 1) MAML [6] which is a deterministic meta-learning method, and 2) Bayesian meta-learning methods that use Gaussians for prediction or prior, including ABML [41], MetaFun [53], CNPs [8] and F-PACOH-GP [43]. For classification, we consider Simple-CNAPs [1] and TSA [28], which respectively resort to amortized variational inference and gradient-based optimization for estimating the task-specific parameters from the support set. Both are SOTA cross-domain few-shot classification approach on the Meta-dataset [49] benchmark. For more experimental details, hyper-parameter configurations, and additional experimental results, please refer to Appendix B and C.
5.2 Datasets and Evaluation Metrics
Sinusoids Few-shot Regression We consider 1D sinusoids regression tasks in the form  $y(x) = A\sin (B(x + C))$ . For ID meta-training, we consider frequency  $B = 1$ , while sample amplitude  $A$  and phase  $C$  uniformly from a set of equally-spaced points  $\{1, 1.1, 1.2, \dots, 4\}$  and  $\{0, 0.1, 0.2, \dots, 0.5\pi\}$ , respectively. Each training task consists of 2 to 5 support and 10 query points with  $x$  uniformly sampled from  $\mathcal{X} \in [-5.0, 5.0]$ . During testing, we evaluate the models on 500 ID and OOD tasks each with 512 equal-distant query points in  $\mathcal{X}$ . For ID testing, we expand the range of the tasks by uniformly sampling  $A \in [1, 4]$  and  $C \in [0, 0.5\pi]$ . For OOD tasks, we randomly change either the phase distribution to  $C \in [0.6\pi, 0.75\pi]$ , amplitude to  $A \in [0.1, 0.8] \cup [4.2, 5.0]$  or frequency to  $B \in [1.1, 1.25]$ . Details for the multi-sinusoids regression experiment can be found in C.1. We use MSE and negative log-likelihood on query samples to evaluate the regression performance.
Drug Activity Prediction Few-shot Regression In each task, we aim to predict the drug-target binding affinity of query molecular compounds given 10 to 50 labelled examples from the same domain defined by molecular size. We use the lbap-general-ic50-size ID/OOD task split in the DrugOOD [21] benchmark, which divides the molecules into  $222/145/23$  domains by molecular size for ID Train / ID Test / OOD Test, respectively. The regression performance is evaluated by the square of Pearson coefficient  $(\mathbf{R}^2)$  between predictions and the ground-truth values. We report the mean and median  $R^2$  on 500 tasks sampled from ID and OOD testing domains.
Meta-dataset [49] 5-way 1-shot Classification This experiment considers image classification problems on Meta-dataset [49]. Each task contains up to 10 query images per class from the same domain. Following the current state-of-the-art practice [28, 1], we use Aircraft, dtd, cub, vgg-flower, fungi, quickdraw and omniglot as the ID datasets for meta-training and meta-testing, while traffic, mscoco, cifar10, cifar100 and mnist are treated as OOD datasets for meta-testing only.
ODD Task Detection Evaluation We compare the OOD task detection performance of Energy Sum against several model-agnostic OOD detection baselines. Concretely, for classification, we compare against max-softmax score [16], ODIN [30], MAH [27], and max-logits score [15]; for regression, we consider Averaged Bayesian prediction uncertainty in standard deviation (Std) on support samples, and Averaged Support samples Negative Log-Likelihood (SNLL) under model's task-specific predictive probability, i.e.,  $-\mathbb{E}_{\phi_i \sim q_{\psi}(\phi_i | \mathcal{T}_i^s)}[\mathbb{E}_j[\log p_{\omega}(y_{ij}^s | x_{ij}^s, \phi_i)]]$  for baselines and  $\mathbb{E}_{\phi_i \sim q_{\psi}(\phi_i | \mathcal{T}_i^s)}[\mathbb{E}_j[E_{\omega}(y_{ij}^s, x_{ij}^s, \phi_i)]]$  for EBML. Following common practice [17, 16], we report AUROC, AUPR and FPR95 for OOD detection performance. Details for these metrics can be found in Appendix B.1.
5.3 OOD Detection Results
Energy sum performs best in OOD task detection. Table 1 and 8. The proposed energy sum further improves our SNLL-only results in all three OOD detection metrics - with  $15.2\%$  and  $11.8\%$  significant reduction in FPR95, outperforming the best baseline methods by  $20.0\%$  and  $39.1\%$ , in single and multi-sinusoids situations respectively. In Table 2 for OOD classification task detection, Energy Sum consistently results in superior OOD detection performance, outperforming the best baselines by large margins of  $36.84\%$  and  $20.19\%$  in FPR95 for Simple-CNAPs and TSA, respectively.
Table 1: OOD task detection performance on single-sine and DrugOOD [21] few-shot regression tasks.
OOD Scores	Models	Sinusoids	DrugOOD
AUROC↑	AUPR↑	FPR95↓	AUROC↑	AUPR↑	FPR95↓
Std	ABML [41]	50.14	54.80	97.20	57.82	50.31	74.80
F-PACOH-GP [43]	49.52	51.30	94.20	81.74	71.99	32.00
CNPs [8]	22.72	35.34	99.60	93.56	89.58	13.00
Metafun [53]	76.57	80.33	82.40	85.68	80.55	58.18
SNLL	ABML [41]	82.48	81.31	61.00	80.99	79.12	47.60
F-PACOH-GP [43]	91.78	93.23	52.40	37.73	45.01	85.21
CNPs [8]	95.63	96.46	34.22	17.25	34.07	91.40
Metafun [53]	96.25	97.11	32.00	83.54	85.54	65.17
EBML-CNPs (Ours)	96.46	97.41	29.40	99.71	99.71	2.20
Energy Sum	EBML-CNPs (Ours)	97.74	98.31	14.20	99.79	99.78	1.40
Modelling the joint distribution improves OOD detection under Domain-shift. In Table 1 DrugOOD regression tasks, using either our SNLL or Energy Sum as OOD scores can achieve better detection performance than baselines. In particular, our method outperforms the best OOD detection results obtained using Gaussian SNLL and Std by  $43.84\%$  and  $11.6\%$  in FPR95, respectively.
Qualitative Illustration. In Figure 5, we visualize the predictive distribution  $p(y_{ij}|\mathbf{x}_{ij},\boldsymbol {\phi}_i)$  learned using an EBM decoder and a Gaussian decoder on a sampled ID multi-sinusoids task. The EBM clearly shows two prediction modes at all non-overlapping positions, whereas the Gaussian decoder is unable to model the multi-modality, resulting in a blurry prediction.
Figure 5: Predictive distribution of Middle and data EBM vs Right a Gaussian for an ID task.
Computational Complexity Analysis. We conduct a computational complexity analysis for EBML by comparing its wall-clock training time and convergence to baselines in Figure 6 below. EBML-CNPs eventually achieves better OOD detection performance than baseline CNPs meanwhile matching its regression performance at all training epochs. In Table 15 Appendix C.4, we show EBML-CNPs is computationally cheaper and faster than traditional Bayesian methods, namely, F-PACOHGP [43] which requires matrix inversion for inference with Gaussian processes prior, and ABML [41] which imposes a Gaussian prior over the entire parameter space of the model.
Figure 6: Left: Wall-clock convergence in seconds, and Right: performance vs number of training epochs, for EBML-CNPs vs CNPs in single-sinosoid few-shot regression tasks. The plots show the regression (MSE  $\downarrow$ ) and combined OOD tasks detection (1-AUROC)+(1-AUPR)+FPR95  $\downarrow$  performance on single sine few-shot regression tasks during training. Curves are moving averages with window size 3. EBML-CNPs achieves better final performance than CNPs.
Energy sum achieves better OOD detection results with EMB prior than Gaussian. In Table 3 and 9, we investigate the contribution of the prior EBM in improving the modelling of meta-training task distribution. We train CNPs and ABML using diagonal Gaussian distribution as the prior in ELBO, and compute OOD scores as (a) SNLL, and (b) the sum between SNLL and the NLL of task-specific latent evaluated under the learned Gaussian prior (indicated by  $+Gauss\ Prior$ ). The results show that energy sum using an EBM prior outperforms all ablated models. The OOD detection performance of our model benefits from adding the prior EBM energy to the data EBM energy (SNLL), resulting in the most reduction in FPR95 on both single and multi-sinusoids tasks (15.2% and 11.8%, respectively). This suggests the improved expressive of EBM over simple distributions can indeed lead to learning a more accurate model of the meta-training ID task distribution.
Energy sum achieves better OOD detection results when learning the joint distribution In Table 16, we compare EBML-joint, which is exactly our proposed training procedure in the paper,
Table 2: OOD task detection performance on Meta-dataset 5-way 1-shot classification tasks.
OOD Scores	Simple-CNAPs [1]	TSA [28]
AUROC↑	AUPR↑	FPR95↓	AUROC↑	AUPR↑	FPR95↓
max-softmax [16]	85.50	85.54	65.43	89.25	87.14	46.02
max-logits [15]	50.00	70.83	95.00	50.14	44.64	95.28
ODIN [30]	90.49	89.42	43.57	92.02	90.18	37.36
MAH [27]	71.18	69.76	90.52	94.54	93.95	23.83
Domain Classifier	83.10	73.18	53.17	n/a	n/a	n/a
EBML Energy Sum	97.01	94.92	6.74	99.10	98.48	3.64
Table 3: Ablation study on Energy Sum for OOD detection on single-sinusoids.
Models	OOD Scores	Sinusoids
AUROC↑	AUPR↑	FPR95↓
ABML [41]	SNLL	82.48	81.31	61.00
+Gauss Prior	86.95	86.64	52.20
CNPs [8]	SNLL	94.81	96.34	38.40
+Gauss Prior	94.61	96.10	34.40
EBML-CNPs	SNLL	96.46	97.41	29.40
+EBM Prior	97.74	98.31	14.20
Table 4: Few-shot regression performance on single-sinusoids and DrugOOD [21] tasks.
Models	Sinusoids	DrugOOD
ID MSE ↓	ID Mean R2↑	ID Median R2↑	OOD Mean R2↑	OOD Median R2↑
F-PACOH-GP [43]	0.068±0.016	0.492	0.454	0.055	0.027
Metafun[53]	0.009±0.002	0.537	0.541	0.054	0.027
CNPs [8]	0.009±0.002	0.540	0.549	0.066	0.046
ABML [41]	0.127±0.013	0.452	0.443	0.051	0.029
MAML [6]	0.119±0.013	0.462	0.475	0.055	0.024
EBML-CNPs	0.009±0.002	0.533	0.553	0.071	0.043
and EBML-conditional, which follows the same training with EBML-joint but models  $p(\mathbf{Y} \mid \mathbf{X})$  instead of  $p(\mathbf{X}, \mathbf{Y})$ . With all other factors being the same, EBML-joint significantly outperform EBML-conditional in OOD detection on DrugOOD regression tasks with domain shift in  $\mathbf{X}$ . This supports our motivation for using the joint distribution instead of the conditional distribution for training a potentially better OOD detector. Detail of this ablation study can be found in Appendix D.
5.4 OOD Generalization Results
EBML achieves SOTA regression performance. In Table 4, for single-sinusoids, EBML is able to match the MSE of the best-performing baseline methods; while on multi-sinusoids in Table 7, EBML obtains the lowest ID NLL, specifically 0.58 lower than the best baseline, thanks to our energy-based decoder which is sufficiently expressive for modelling the multimodality at each input.
Task adaption using Eqn. (11) improves few-shot classification performance. In Table 5, we report the average classification accuracy computed over 600 test tasks per ID and OOD domains. In meta-testing, we obtain classification results for EBML-TSA by running gradient descent on the objective in Eqn. (11) to optimize the task-specific modules in TSA from scratch. With this addition of prior energy in the OOD adaption objective, EBML-TSA further improves TSA results in 5/7 ID domains and all 5 OOD domains. Additional OOD classification results in Table 11 Appendix C further confirm the superiority of our proposed OOD task adaptation strategy in Eqn. (11) over prior baselines.
Table 5: Classification performance on 5-way 1-shot tasks for both ID and OOD domains in Meta-dataset.
Datasets	TSA [28]	EBML-TSA (Ours)
Omniglot	98.63±0.26	98.67±0.26
Textures	51.93±0.87	52.35±0.88
Aircraft	78.91±0.86	78.47±0.86
Birds	75.02±0.90	75.52±0.90
VGG Flower	80.37±0.80	80.30±0.83
Fungi	70.89±0.93	72.29±0.94
Quickdraw	79.02±0.84	80.27±0.85
MSCOCO	52.28±0.94	53.03±0.97
Traffic Sign	57.40±0.94	58.85±1.01
CIFAR10	49.16±0.82	50.04±0.89
CIFAR100	62.25±1.01	62.77±1.05
MNIST	74.72±0.83	76.08±0.88
Avg ID	76.40	76.84
Avg OOD	59.16	60.15
Avg All	69.22	69.89
This paper proposes a new energy-based meta-learning (EBML) framework for the first time, which directly characterizes any arbitrary meta-training task distribution using two data and prior energy functions. EBML is compatible with many existing SOTA meta-learning algorithms and allows both detection and adaption of OOD tasks. The sum of the two learned energy functions gives an unnormalized probability distribution proportional to the underlying task likelihood, deployable as OOD scores. The experiment results show the superiority of Energy Sum over traditional methods in detecting both OOD regression and classification tasks, and the possibility of achieving improved OOD adaptation performance with EBML through minimizing the task energy. One limitation of EBML is that our current OOD task adaptation strategy does not consider the effect of negative transfer, as some OOD tasks may benefit from adapting from scratch without ID energy prior regularization. Thus, in future works, we are interested in designing task-specific adaptation strategies for EBML that can selectively adapt OOD tasks for better performance.
[48] Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
[49] Eleni Triantafillou, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Utku Evci, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, and Hugo Larochelle. Meta-dataset: A dataset of datasets for learning to learn from few examples. In International Conference on Learning Representations, 2020.
[50] Haotao Wang, Aston Zhang, Yi Zhu, Shuai Zheng, Mu Li, Alex J Smola, and Zhangyang Wang. Partial and asymmetric contrastive learning for out-of-distribution detection in long-tailed recognition. In International Conference on Machine Learning, pages 23446-23458, 2022.
[51] Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML'11, page 681-688, Madison, WI, USA, 2011. Omnipress.
[52] Jeffrey Ryan Willette, Hae Beom Lee, Juho Lee, and Sung Ju Hwang. Meta learning low rank covariance factors for energy based deterministic uncertainty. In International Conference on Learning Representations, 2022.
[53] Jin Xu, Jean-Francois Ton, Hyunjik Kim, Adam R Kosiorek, and Yee Whye Teh. Metafun: Meta-learning with iterative functional updates. In International Conference on Machine Learning (ICML), 2020.
[54] Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian Model-Agnostic Meta-Learning. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[55] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRL, pages 1094–1100, 2020.
[56] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep Sets. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[57] Yang Zhao, Jianwen Xie, and Ping Li. Learning Energy-Based Generative Models via Coarse-to-Fine Expanding and Sampling. February 2022.