SCORE NEURAL OPERATOR: A GENERATIVE MODEL FOR LEARNING AND GENERALIZING ACROSS MULTIPLE PROBABILITY DISTRIBUTIONS

## Abstract

Abstract Most existing generative models are limited to learning a single probability distribution from the training data and cannot generalize to novel distributions for unseen data.
An architecture that can generate samples from both trained datasets and unseen probability distributions would mark a significant breakthrough.
Recently, score-based generative models have gained considerable attention for their comprehensive mode coverage and high-quality image synthesis, as they effectively learn an operator that maps a probability distribution to its corresponding score function.
In this work, we introduce the Score Neural Operator , which learns the mapping from multiple probability distributions to their score functions within a unified framework.
We employ latent space techniques to facilitate the training of score matching, which tends to over-fit in the original image pixel space, thereby enhancing sample generation quality.
Our trained Score Neural Operator demonstrates the ability to predict score functions of probability measures beyond the training space and exhibits strong generalization performance in both 2-dimensional Gaussian Mixture Models and 1024-dimensional MNIST double-digit datasets.
Importantly, our approach offers significant potential for few-shot learning applications, where a single image from a new distribution can be leveraged to generate multiple distinct images from that distribution.

## 1 Introduction

Generative modeling has become a cornerstone of modern artificial intelligence and machine learning, with applications ranging from image synthesis to natural language processing Brown (); Karras et al. (). These models can be broadly categorized into two main approaches: likelihood-based methods, such as autoregressive models and normalizing flows Rezende & Mohamed (); Kingma & Dhariwal (); Papamakarios et al. (), and implicit generative models, exemplified by Generative Adversarial Networks (GANs) Goodfellow et al. (); Arjovsky et al. (). While these approaches have shown remarkable success, they often face limitations in terms of model architecture constraints or reliance on surrogate objectives for likelihood approximation Theis et al. (); Salimans et al. (); Papamakarios et al. ().

In recent years, score-based generative models have emerged as a powerful alternative, gaining considerable attention for their ability to produce high-quality samples and provide comprehensive mode coverage. These models operate by approximating the score function (the gradient of the log-density) of the target distribution through minimizing the Fisher divergence Song & Ermon (); Hyvärinen & Dayan (). Once trained, new samples can be generated using Langevin dynamics Song et al. (). The appeal of score-based methods lies in their ability to avoid the limitations of both likelihood-based and implicit generative models while maintaining high-quality output Song & Ermon (); Song et al. (); Ho et al. ().

However, most existing generative models, including traditional score-based approaches, are designed to learn and sample from a single probability distribution. This limitation becomes particularly apparent in scenarios where we need to generate samples from multiple related but distinct distributions, or when faced with novel, unseen distributions. The ability to learn and generalize across multiple distributions would mark a significant advancement in the field, enabling more flexible and adaptable generative models.

To address this capability gap, we introduce the Score Neural Operator, a novel framework that learns to map multiple probability distributions to their corresponding score functions within a unified model. Our approach leverages recent advancements in operator learning Azizzadenesheli et al. (); Lu et al. (); Seidman et al. () to capture the underlying relationships between different distributions and their score functions. By doing so, we enable the model to not only generate high-quality samples from trained distributions but also to generalize to unseen distributions without requiring retraining. The main contributions of our work can be summarized as follows:

- •
We introduce the Score Neural Operator, a novel architecture that learns to map multiple probability distributions to their score functions, enabling generalization to unseen distributions.
- •
We employ latent space techniques to facilitate score matching in high-dimensional spaces, improving both the quality of generated samples and the efficiency of the training process.
- •
We demonstrate the effectiveness of our approach on both low-dimensional Gaussian Mixture Models and high-dimensional image data (MNIST double-digit), showing strong generalization performance to unseen distributions.
- •
We showcase the potential of our method for few-shot learning applications, where a single image from a new distribution can be used to generate multiple distinct samples from that distribution.

Our work represents a significant step towards more flexible and generalizable generative models, with potential applications in transfer learning, few-shot generation, and adaptive AI systems.

## 2 Related Work

#### Score-based Generative Models.

Score-based generative models have gained prominence in recent years due to their ability to produce high-quality samples without the need for adversarial training or explicit density estimation. The foundational work by Song and Ermon Song & Ermon () introduced the concept of score matching with Langevin dynamics for sample generation. This was further developed by Song et al. Song et al. (), who proposed a continuous-time formulation of the score-based generative process using stochastic differential equations (SDEs).
Recent advancements in score-based models have focused on improving sample quality and generation speed. Notably, Vahdat et al. Vahdat et al. () introduced a method to perform score matching in a lower-dimensional latent space using a variational autoencoder, significantly improving both the quality of generated samples and the efficiency of the training process. Rombach et al. Rombach et al. () further extended this idea by incorporating cross-attention mechanisms, enabling conditioning on various types of embeddings.
While these developments have greatly enhanced the capabilities of score-based models, they remain primarily focused on learning and sampling from a single distribution. Our work aims to address this limitation by learning a mapping across multiple distributions.

#### Operator Learning.

Operator learning has emerged as a powerful paradigm for capturing mappings between function spaces, with applications in various scientific and engineering domains Kovachki et al. (). Notable approaches in this field include the Deep Operator Network (DeepONet) Lu et al. (), which can approximate nonlinear operators based on the universal approximation theorem, and the Fourier Neural Operator Li et al. (), which leverages spectral methods for efficient operator learning.
Of particular relevance to our work is the NOMAD (Nonlinear Manifold Decoders for Operator Learning) framework introduced by Seidman et al. Seidman et al. (). NOMAD’s ability to handle high-dimensional outputs and provide continuous representations makes it well-suited for our task of learning score functions across multiple distributions.
While operator learning has found success in various scientific computing applications, its potential in generative modeling has been relatively unexplored. Our work bridges this gap by applying operator learning techniques to the domain of score-based generative models.

#### Few-shot Learning in Generative Models.

Few-shot learning, the ability to learn from a small number of examples, has been a long-standing challenge in machine learning. In the context of generative models, few-shot learning is particularly challenging due to the need to capture complex data distributions from limited samples.
Recent works have explored few-shot learning in generative models, such as the approach by Rezende et al. () using meta-learning for one-shot generalization in generative models. However, these methods often struggle with generating diverse samples from limited data, frequently resulting in overfitting to the few available examples. Our Score Neural Operator aims to address this challenge by learning a generalizable mapping across distributions, enabling the generation of diverse samples even when presented with a single example from a new distribution. This capability sets our approach apart from existing few-shot generative methods and opens up new possibilities for adaptive and flexible generative modeling.

## 3 Background

### 3.1 Score-based generative model

Recall that for score-based generative model the diffusion process is governed by a stochastic differential equation (SDE):

$$ $\displaystyle\text{d}{\bf x}={\bf f}({\bf x},t)\text{d}t+{\bf G}({\bf x},t) \text{d}{\bf w},$ (1) $$

where ${\bf f}(\cdot,t):\mathbb{R}^{d}\to\mathbb{R}^{d}$ and ${\bf G}(\cdot,t):\mathbb{R}^{d}\to\mathbb{R}^{d\times d}$. At $t=0$, $\mathbf{x}(0)$ is sampled from the original data distribution $p_{0}(\mathbf{x})$. Those input samples are disturbed by noise through a Stochastic Differential Equation equation 1, and evolve into a Gaussian prior with fixed mean and covariance at $t=1$. Subsequently, a sample is drawn from this prior and denoised via the reverse diffusion process, which is structured as follows to reconstruct the original data distribution:

$$ $\displaystyle\text{d}{\bf x}=\left[{\bf f}({\bf x},t)-\nabla\cdot[{\bf G}({\bf x },t){\bf G}({\bf x},t)^{\text{T}}]-{\bf G}({\bf x},t){\bf G}({\bf x},t)^{\text {T}}\nabla_{\bf x}\log p_{t}({\bf x})\right]\text{d}t+{\bf G}({\bf x},t)\text{ d}{\bf\bar{w}}.$ (2) $$

All the needed information to generate new samples is the score function $\nabla_{\bf x}\log p_{t}({\bf x})$, which is approximated by a parameterized neural network $\mathbf{s}_{\theta}(\mathbf{x}(t),t)$. In Song et.al Song et al. (), the loss function is given by,

$$ (3) $$

where $\lambda(t)$ is the weighting function usually chosen to be inversely proportional to $\mathbb{E}[\|\nabla_{\mathbf{x}}\log p_{0t}(\mathbf{x}(t)\mid\mathbf{x}(0))\|_
{2}^{2}],$ $p_{0t}$ is the transitional kernel that transforms the distribution $p_{0}(\mathbf{x})$ at time 0 0 into the distribution $p_{t}(\mathbf{x})$ at time $t$ and has a closed form by choosing specific ${\bf f}$ and ${\bf G}$.

Equation 3 targets the score function of a single distribution. We extend this training objective to multiple probability distributions by utilizing an operator learning architecture, which is adept at learning mappings between functional spaces. To achieve this, we require a method to embed a probability distribution into a vector such that it captures sufficient statistical information. A straightforward approach is to compute the characteristic function of the probability measure at predetermined points. However, the size of the embedding increases exponentially with the dimension, making this method computationally prohibitive for high-dimensional data, such as images. The Vector Quantized-Variational Autoencoder (VQ-VAE) Van Den Oord et al. () employs an encoder to compress input samples into a discrete latent space, or codebook. However, different samples from the same distribution are mapped to distinct latent vectors in the codebook. Ideally, we desire that different samples from the same distribution yield the same embedding, as this would enhance the stability of the training process. Thus, an efficient and effective method for embedding probability distributions is essential. In the next section, we introduce two scalable approaches for embedding probability distributions as data dimensionality increases.

### 3.2 Kernel Mean Embeddings

Consider a set $\mathcal{X}$, and a kernel $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ which is positive definite and symmetric. Define the vector space of functions

$$ $\mathcal{H}_{0}:=\left\{\sum_{i=1}^{n}\alpha_{i}k(\cdot,x_{i})\;\Big{|}\;n\in \mathbb{N},\;\alpha_{i}\in\mathbb{R},\;x_{i}\in\mathcal{X}\right\}.$ $$

We can give this space the inner product

$$ $\Big{\langle}\sum_{i=1}^{n}\alpha_{i}k(\cdot,x_{i}),\sum_{i=1}^{m}\beta_{i}k( \cdot,y_{i})\Big{\rangle}=\sum_{i=1}^{n}\sum_{j=1}^{m}\alpha_{i}\beta_{j}k(x_{ i},y_{j}).$ (4) $$

The completion of $\mathcal{H}_{0}$ with respect to the norm induced by the inner product equation 4 is the Reproducing Kernel Hilbert Space (RKHS) associated with the kernel $k$, denoted $\mathcal{H}_{k}$. It is a Hilbert space of functions from $\mathcal{X}\to\mathbb{R}$ with an inner product from equation 4.

An important property of $\mathcal{H}_{k}$ is the *reproducing property* which states that for any $f\in\mathcal{H}_{k}$ and any $x\in\mathcal{X}$,

$$ $f(x)=\langle f,k(\cdot,x)\rangle,$ (5) $$

that is, pointwise evaluation is a linear functional on $\mathcal{H}_{k}$ (this is sometimes taken as part of the definition of an RKHS). For more information on RKHS’s see Manton et al. ().

Let $\mathcal{P}(\mathcal{X})$ be the space of probability measures on a set $\mathcal{X}$. A kernel $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ (under some reasonable assumptions) induces a mapping from $\mathcal{P}(\mathcal{X})\to\mathcal{H}_{k}$ defined by

$$ $\mathbb{P}\longmapsto\mu_{\mathbb{P}}:=\int_{\mathcal{X}}k(\cdot,x)d\mathbb{P} (x),$ (6) $$

where we call $\mu_{\mathbb{P}}$ the *kernel mean embedding* of $\mathbb{P}$. Note that $\mu_{\mathbb{P}}\in\mathcal{H}_{k}$ and is therefore a function $\mu_{\mathbb{P}}:\mathcal{X}\to\mathbb{R}$. In the case that $\mathbb{P}$ is an empirical distribution, that is

$$ $\hat{\mathbb{P}}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}},$ (7) $$

we have that

$$ $\mu_{\hat{\mathbb{P}}}=\frac{1}{n}\sum_{i=1}^{n}k(\cdot,x_{i}).$ (8) $$

The definition of the kernel mean embedding along with the reproducing property in $\mathcal{H}_{k}$ allows one to show that we may represent an inner product of kernel mean embeddings in $\mathcal{H}_{k}$ as an expectation of the kernel over the product measure of the associated distributions. Specifically, for any probability measures $\mathbb{P},\mathbb{Q}\in\mathcal{P}(\mathcal{X})$,

$$ $\langle\mu_{\mathbb{P}},\mu_{\mathbb{Q}}\rangle_{\mathcal{H}_{k}}=\int_{ \mathcal{X}}\int_{\mathcal{X}}k(x,y)\textnormal{d}\mathbb{P}(x)\textnormal{d} \mathbb{Q}(y).$ (9) $$

In the case that we consider two empirical distributions

$$ $\hat{\mathbb{P}}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}},\quad\hat{\mathbb{Q}} =\frac{1}{m}\sum_{i=1}^{m}\delta_{y_{i}},$ $$

we can compute the inner product of their embedding in $\mathcal{H}_{k}$ as

$$ $\langle\mu_{\hat{\mathbb{P}}},\mu_{\hat{\mathbb{Q}}}\rangle^{\phantom{|}}_{ \mathcal{H}_{k}}=\frac{1}{n}\frac{1}{m}\sum_{i=1}^{n}\sum_{j=1}^{m}k(x_{i},y_{ j}).$ (10) $$

For more details about kernel mean embeddings see Muandet et al. (; ).

We then conduct Principal Component Analysis on $\mu_{\mathbb{P}_{i}}$ in RKHS to reduce dimensions. For any probability measure $\mathbb{P}_{t},$ the $k$-th element of $\mathbf{u}$ is given by,

$$ $\displaystyle u_{k}=\sum_{j=1}^{N}\alpha_{j}^{k}\left[\left<\mu_{\mathbb{P}_{t }},\mu_{\mathbb{P}_{j}}\right>+\frac{1}{N^{2}}\sum_{k=1}^{N}\sum_{m=1}^{N} \left<\mu_{\mathbb{P}_{k}},\mu_{\mathbb{P}_{m}}\right>-\frac{1}{N}\sum_{k=1}^{ N}\left<\mu_{\mathbb{P}_{k}},\mu_{\mathbb{P}_{j}}\right>-\frac{1}{N}\sum_{m=1} ^{N}\left<\mu_{\mathbb{P}_{m}},\mu_{\mathbb{P}_{t}}\right>\right],$ (11) $$

where the expression for $\alpha_{j}^{k}$ is given in the Appendix.

### 3.3 Prototype Embedding

Let $\mathcal{N}$ be a neural network. Then the prototype embedding of probability measure $\nu$ can be given as,

$$ $\displaystyle\mathbf{u}=\mathbb{E}_{\mathbf{x}\sim\nu}\mathcal{N}(\mathbf{x}).$ (12) $$

A suitable choice for $\mathcal{N}$ is the encoder part of a trained variational autoencoder. This encoder compresses all training probability distributions into a latent space, from which the original data can be fully reconstructed using latent samples. By taking the ensemble average of these latent samples, the embedding for each probability distribution is extracted. These probability embeddings are then fed into the score neural operator, which identifies the score function of the current probability distribution.

## 4 Model Formulation

### 4.1 Score neural operator

We utilize NOMAD as the operator learning architecture to learn the score functions of probability distributions, primarily because it provides a continuous
representation of the output function, which is essential for managing high-dimensional datasets such as images. In NOMAD, the encoder processes the vector representation $\mathbf{u}^{\nu}$ of each probability measure $\nu$, utilizing either equation 11 or equation 12. The decoder, on the other hand, processes the input $\mathbf{y}=[\mathbf{x},t]$.
The output from NOMAD is the score function of $\nu$, evaluated at $\mathbf{y}$. We generalize the objective function, as shown in equation 3 to accommodate multiple probability measures in the following form:

$$ (13) $$

Minimizing equation 13 will derive the optimal $\mathbf{s}_{\theta}^{\text{opt}}(\mathbf{u}^{\nu},\mathbf{x}(t),t)$ for any $\nu\sim\mathcal{P}(\mathcal{X})$. With the trained $\mathbf{s}_{\theta}^{\text{opt}}$, we can integrate the probability flow ODE equation 14 derived by Song et al. () to generate samples for any $\nu\sim\mathcal{P}(\mathcal{X})$,

$$ $\displaystyle d\mathbf{x}=\left[\mathbf{f}(\mathbf{x},t)-\frac{1}{2}\nabla \cdot[\mathbf{G}(\mathbf{x},t)\mathbf{G}(\mathbf{x},t)^{\text{T}}]-\frac{1}{2} \mathbf{G}(\mathbf{x},t)\mathbf{G}(\mathbf{x},t)^{\text{T}}\mathbf{s}_{\theta} ^{\text{opt}}(\mathbf{u}^{\nu},\mathbf{x}(t),t)\right]dt.$ (14) $$

### 4.2 Latent space score neural operator

(Vahdat et al., ), (Rombach et al., ) enhance the performance and speed of score-based generative models by conducting score matching in low-dimensional latent spaces. They use a variational autoencoder to compress the original data distribution into a low-dimensional latent space and perform the diffusion process there. Score matching in the latent space is significantly easier and faster than in the original image pixel space. Another advantage of utilizing the latent space is the wide mode coverage, which leads to a more expressive generated distribution. Inspired by their approach, we employ a VAE to map probability measures from the original pixel space $\mathcal{P}(\mathcal{X})$ to a latent space $\mathcal{Z}(\mathcal{Y})$ with Kullback-Leibler (KL) regularization. Subsequently, we train the score neural operator in this latent space.

Specifically, let $\mu\sim\mathcal{P}(\mathcal{X})$ be a probability measure. Recall that for a single distribution, the variational lower bound on negative data log-likelihood is given by,

$$ $\displaystyle\mathcal{L}_{\varphi,\phi}(\mathbf{x})=\mathbb{E}_{q_{\phi}( \mathbf{z}|\mathbf{x})}\left[-\log p_{\varphi}(\mathbf{x}|\mathbf{z})\right]+ \beta\mathbf{KL}\left(q_{\varphi}(\mathbf{z}|\mathbf{x})||\mathcal{N}(\mathbf{ 0},\mathbf{I})\right).$ (15) $$

Equation 15 can be generalized to multiple probability distributions with the following optimization target,

$$ $\displaystyle\min_{\varphi,\phi}\mathcal{L}^{\text{VAE}}_{\varphi,\phi}(\beta) =\mathbb{E}_{\mu\sim\mathcal{P}(\mathcal{X})}\mathbb{E}_{\mathbf{x}\sim\mu} \mathbb{E}_{q_{\varphi}(\mathbf{z}|\mathbf{x})}\left[-\log p_{\varphi}(\mathbf {x}|\mathbf{z})\right]+\beta\mathbf{KL}\left(q_{\varphi}(\mathbf{z}|\mathbf{x} )||\mathcal{N}(\mathbf{0},\mathbf{I})\right).$ (16) $$

We can conduct score matching in the latent space through,

$$ $\displaystyle\min_{\theta}\mathcal{L}^{\text{SGM}}_{\theta}=$ $\displaystyle\mathbb{E}_{\nu\sim\mathcal{Z}(\mathcal{Y})}\mathbb{E}_{\mathbf{z }(0)\sim\nu}\mathbb{E}_{t\sim\mathcal{U}(0,1)}\lambda(t)$ $\displaystyle\mathbb{E}_{\mathbf{z}(t)\sim p_{0t}(\mathbf{z}(t)\mid\mathbf{z}( 0))}\left.\Big{[}\left\|\mathbf{s}_{\theta}(\mathbf{u}^{\nu},\mathbf{z}(t),t)- \nabla_{\mathbf{z}(t)}\log p_{0t}(\mathbf{z}(t)\mid\mathbf{z}(0))\right\|_{2}^ {2}\right].$ (17) $$

The dependency on the distribution $\mu$ is replaced by its vectorized embedding $\mathbf{u}$ during computation, and we do not differentiate between them in the expressions. In equation 16 $\beta$ is a hyperparameter that balances the reconstruction error and the regularization error. Our training objective can be summarized as follows:

$$ $\displaystyle\min_{\phi,\varphi,\theta}\mathcal{L}_{\phi,\varphi}^{\text{VAE}} (\beta)+\gamma\mathcal{L}_{\theta}^{\text{SGM}},$ (18) $$

where $\gamma$ is a hyperparameter that balances the VAE loss and score-matching loss. When sampling from a test probability measure, we first compute its embedding $\mathbf{u}$ using equation 11 or equation 12. Next, we sample from a standard normal distribution and denoise back to the latent distribution using the predicted score function $\mathbf{s}_{\theta}^{\text{opt}}(\mathbf{u},\mathbf{x}(t),t)$ and probability flow ODE equation 14. Finally, the generated samples are transformed back to the original pixel space through the decoder.

## 5 Experiments

### 5.1 Evaluation metric of generated samples

We discuss several metrics to assess the quality of samples generated by a generative model.
Fréchet inception distance (FID) is a metric commonly used to evaluate the quality of generated images. To evaluate the quality of generated samples in low-dimensional space, we use the MMD (Maximum Mean Discrepancy) as the metric.
The maximum mean discrepancy between two distribution $\mathbb{P},\mathbb{Q}$ is an integral probability metric, given by Muandet et al. (),

$$ $\displaystyle\text{MMD}(\mathbb{P},\mathbb{Q},\mathcal{H})=\sup_{\|f\|\leq 1} \left\{\int f(\mathbf{x})\mathbb{P}(\mathbf{x})-\int f(\mathbf{y})\mathbb{Q}( \mathbf{y})\right\}=\|\mu_{\mathbb{P}}-\mu_{\mathbb{Q}}\|_{\mathcal{H}}.$ (19) $$

Equation 19 is non-negative and lower bounded by 0 when $\mathbb{P}\,{\buildrel d\over{=}}\,\mathbb{Q}$. It can be estimated using empirical MMD,

$$ (20) $$

For grayscale images, such as those from the MNIST dataset, we employ a ResNet-18 architecture He et al. () to train a digit classifier. We then utilize this trained classifier to assess the classification accuracy of digits generated from a test probability measure.

### 5.2 2D Gaussian Mixture Models (GMMs)

For our 2D toy example, we utilize a family of Gaussian Mixture Models. Each model comprises four components, with centers positioned at the center of cells in a $6\times 6$ lattice. All components follow a uniform distribution within the range of the cell they lie in. We define two panels within the lattice: the left panel consists of the first three columns, and the right panel includes the last three columns. For the training distributions, we randomly select two components from the same row in the left panel and two from the same column in the right panel, a configuration we denote as “left row right col.” Similarly, we select two components from the same column in the left panel and two from the same row in the right panel, referred to as “left col right row.” The testing distributions, however, are composed of combinations labeled as “left row right row.”
Given the low-dimensional nature of our data, we employ the MMD metric, as introduced in Section 5.1 to assess the quality of the generated samples. Each probability distribution is mapped to a RKHS using KME, and the probability embedding is evaluated using equation 11. The computation of inner products in RKHS is notably efficient in this low-dimensional setting. We generate 2000 training examples to train the score neural operator. The model is then used to predict the score functions for testing distributions, which feature patterns not present during the training phase. The results demonstrated in figure 1 shows that Score Neural Operator can generate high-quality samples from unseen datasets, comparable to those produced by individual score-based generative models.

Figure: Figure 1: We examine four 2D Gaussian Mixture Model (GMM) distributions: the first two are derived from the training sets, and the last two are from the testing sets. To generate samples for these distributions, we utilized four score-based generative models (SGMs) and a Score Neural Operator (SNO). The first row displays raw samples from all four distributions. The second row presents samples generated using the SNO, while the third row shows samples generated by the four distinct SGMs. The maximum mean discrepancy (MMD) between samples generated by the SGMs and the original dataset are 0.0070, 0.0053, 0.0079, and 0.0056 (from left to right). For the SNO, the MMD values are 0.0054, 0.0064, 0.0081, and 0.0148, respectively. These results indicate that the SNO produces samples of comparable quality to those generated by individual SGM models.
Refer to caption: extracted/5927651/figures/2DGMM.png

### 5.3 32 × \times × 32 Double-Digit MINST

We developed an MNIST double-digit dataset, ranging from 00 to 99, by concatenating pairs of single digits from the original MNIST dataset. This process generated 100 unique probability distributions, from which we selected 70 for training the Score Neural Operator and 30 for testing. Performing score-matching directly in a 1024-dimensional pixel space proved time-consuming and susceptible to overfitting. To address this, we implemented the latent space technique outlined in section 4.2, which involves training a Variational Autoencoder (VAE) and conducting low-dimensional score-matching in an end-to-end manner using equation 18. We executed multiple experiments under varied settings to assess the classification accuracy on both training and testing datasets. The results are summarized in Table 1. Our experiments revealed that end-to-end training faces convergence challenges when KMEs are computed in a dynamically changing latent space. However, these challenges are mitigated when prototype embeddings are utilized. By computing KMEs in the original pixel space and employing fixed probability embeddings u for end-to-end training, we achieved rapid convergence. We also observed that the stability of the training process is significantly impacted when the Variational Autoencoder and score-matching are trained separately. This approach resulted in considerable variability in training and testing classification accuracies, which were highly dependent on the random seed used. Figure 5 and Figure 5 show the generated training and testing digits after the training has fully converged. The results indicate that the Score Neural Operator can generate high-quality images for both the training and testing sets, demonstrating its ability to accurately predict the score functions of unseen datasets.

To demonstrate our model’s ability to accurately predict score functions for distributions not included in the training set, we employed a conditional score model as a baseline. This model also incorporates a probability embedding as input to identify specific probability distributions. Unlike the embeddings in our Score Neural Operator, which are interconnected across different distributions, the embeddings in the conditional score model resemble isolated one-hot encodings, lacking inherent generalization capabilities. To verify this, we trained the conditional score model on 70 training digits over 10,000 epochs, followed by finetuning on 30 testing digits. Both training and testing accuracies were continuously monitored (see Figure 2). The results indicate that while the score functions for new distributions can be efficiently trained using well-pretrained weights, the model does not exhibit generalization proficiency in predicting score functions for unseen datasets.

**Table 1: Comparative Performance of Score Neural Operator Configurations: We compare the performance of our model in latent and pixel spaces, using various embedding methods (Prototype, KME, and Conditional) and sample sizes.**
| Exps | Space | # Samples | Embedding | Train Acc. | Test Acc. |
| --- | --- | --- | --- | --- | --- |
| 1 | Latent | 2000 | Prototype | 89.5% | 84.2% |
| 2 | Latent | 2000 | KME | 88.0% | 80.0% |
| 3 | Latent | 2000 | Conditional | 87.2% | 0.9% |
| 4 | Pixel | 2000 | KME | 94.8% | 61.1 % |
| 5 | Pixel | 2000 | Prototype | 95.2% | 60.1 % |

Figure: Figure 2: Classification accuracy for training and testing sets over the course of training using a conditional score-based generative model (left and middle) and a Score Neural Operator (right). For the conditional model, training initially involves 10,000 epochs on 70 training probability distributions, each containing 2,000 samples, followed by fine-tuning for another 10,000 epochs on 30 testing distributions. The left sub-figure displays results with 2,000 samples per testing distribution, while the middle sub-figure presents results with only 1 sample per testing distribution. The right sub-figure shows the results of the Score Neural Operator using the prototype embedding method, trained on 70 training digits with 2,000 samples each. The results indicate that the conditional score-based generative model can only perform well on either the training set or the testing set, but not both simultaneously. In contrast, the Score Neural Operator demonstrates robust performance on both sets, despite being trained solely on the training set.
Refer to caption: extracted/5927651/figures/condition_score.png

Figure: Figure 3: t-SNE visualization of the probability embedding $\mathbf{u}$ with Prototype embedding (left) and KME embedding (middle) for 70 training distributions (red) and 30 testing distributions (blue). The right sub-figure displays the t-SNE visualization of latent samples from all 100 distributions. The results indicate that all 100 distributions are well-separated in the latent space, with distributions sharing many similar features positioned close to each other.
Refer to caption: extracted/5927651/figures/tsne_all.png

Figure: Figure 4: The Score Neural Operator is trained on 70 training distributions, each containing 2,000 samples. Subsequently, the model generates an image for each of the 70 training distributions. The results demonstrate that the model is capable of producing high-quality images across multiple training datasets.
Refer to caption: extracted/5927651/figures/train70.png

### 5.4 Applications to few shot learning

Utilizing a score-based generative model to learn the distribution of a single image typically results in a delta function centered on that image, with the generated sample being essentially the original image plus some noise. However, when our trained score neural operator is applied to an unseen image, it is capable of generating a variety of distinct images that adhere to the same underlying data distribution. As illustrated in figure 6, this approach can generate distinct images from only a single input image from a new distribution, achieving approximately 74% classification accuracy. This demonstrates a significant potential for enhancing few-shot learning capabilities. In contrast, a conditional score-based generative model cannot generalize to unseen test distributions once trained. Even with fine-tuning using a single sample, it can only learn the delta function and reproduce the test image (see the right part of Figure 2).

Figure: Figure 6: The Score Neural Operator is trained on 70 training distributions, each containing 2,000 samples. It subsequently generates images for 30 testing distributions using a probability embedding $\mathbf{u}$ computed from only a few samples, denoted as K. We use ‘ACC’ to represent the averaged classification accuracy of 1,000 generated images for each of the 30 testing distributions. The ACC values are 0.7398, 0.8186, 0.8434, and 0.8505 for K=1, 10, 100, and 2000, respectively. The test digit ‘26’ is included to illustrate the quality of the generated images. The results indicate that the model can produce diverse samples from unseen probability distributions using just one testing sample, highlighting its potential for few-shot learning applications.
Refer to caption: extracted/5927651/figures/fewshot.png

## 6 Discussion

#### Summary.

Our study introduces the Score Neural Operator, a novel generative modeling framework that learns to map multiple probability distributions to their corresponding score functions. This approach enables generalization to unseen distributions without retraining, marking a significant advancement in the field. We have demonstrated its effectiveness on both low-dimensional Gaussian Mixture Models and high-dimensional MNIST double-digit data, showcasing strong generalization performance. Our implementation of latent space techniques for score matching has proven effective in handling high-dimensional data, improving both training efficiency and sample quality. Furthermore, we’ve shown the potential of our method for few-shot learning applications, where a single example from a new distribution can generate diverse, high-quality samples.

#### Limitations.

Despite these promising results, the Score Neural Operator has limitations. Its performance may degrade when faced with distributions significantly different from those seen during training. The computational cost of training on a large number of diverse distributions could become prohibitive for very large-scale applications. Additionally, while our method shows improved generalization compared to traditional approaches, it still requires a substantial amount of training data to learn the underlying operator mapping. These limitations highlight the need for further research to enhance the model’s scalability and generalization capabilities.

#### Future Work.

Future work could explore several promising directions. Investigating the theoretical foundations of the Score Neural Operator’s generalization capabilities could provide insights into its performance and guide future improvements. Integrating this approach with other advanced generative modeling techniques, such as transformer architectures or neural ordinary differential equations, could further enhance its capabilities. Extending the model to handle conditional generation tasks would greatly increase its practical utility. Additionally, exploring more efficient training algorithms or adaptive learning strategies could address the scalability concerns identified in our current implementation.

#### Potential Societal Impact.

The potential societal impact of this work is significant, with applications ranging from improved data augmentation techniques in healthcare imaging to more adaptable AI systems in rapidly changing environments. However, as with any advanced generative model, there is potential for misuse, such as in the creation of deepfakes or other synthetic media for malicious purposes. It is crucial that future development and deployment of these technologies be accompanied by robust ethical guidelines and safeguards.

### Reproducibility Statement

The code used to carry out all experiments is provided as a supplementary material for this submission. If the paper is accepted, we plan on making our entire code-base publicly available on GitHub.

## Appendix A Appendix

### A.1 Implementation details

#### NOMAD.

Our NOMAD architecture consists of three MLPs: a branch layer, a trunk layer, and an output layer. Each of these MLPs is 7 layers deep, with 500 neurons per layer and GELU activation functions. The NOMAD is employed for running 2D-GMMs experiments.

#### VAE.

Given the simplicity of the MNIST double-digit dataset, we utilized a 3-layer deep MLP with 512 neurons per layer and ReLU activation functions for both the encoder and decoder components.

#### ScoreNet.

For conducting score-matching in the original pixel space, we used a network comprising 3 down-sampling blocks and 3 up-sampling blocks. Each block includes a convolutional layer and a max-pooling layer, with LogSigmoid activation functions. For the score network in the latent space, we employed 2 down-sampling blocks and 2 up-sampling blocks, each consisting of a 2-layer deep MLP with LogSigmoid activation functions.

### A.2 Experiment Setup

#### Hyperparameters.

For the 2D-GMMs experiments, we used the Variance Exploding Stochastic Differential Equation (VESDE) and an ODE sampler with the hyperparameter $\sigma=25.0$. For the MNIST double-digit dataset, we employed the Variance Preserving Stochastic Differential Equation (VPSDE) and used the Euler-Maruyama method to solve the reverse diffusion process, with $\beta_{\text{max}}=20$ and $\beta_{\text{min}}=0.1$. Additionally, we incorporated Fourier Feature Embedding into the NOMAD architecture for the 2D-GMMs experiments, where the random Gaussian matrix was sampled from $\mathcal{N}(0,10^{2})$. The parameter $\gamma$ is set to 1 in equation 18, and the parameter $\beta$ is set to 2048 in equation 16.

#### Environment Setup.

Our experiments were conducted on an NVIDIA GeForce RTX 3070 GPU. The software environment included JAX version 0.4.33, JAXLIB version 0.4.33, PyTorch version 2.0.1+cu117, CUDA version 11.5, CUDNN version 8.9.2, and Python 3.10.

### A.3 KMEs coefficients

Let $\hat{C}$ be the operator on the RKHS. First we scale each KME by the empirical mean of KMEs over $\textnormal{N_Train}=5000$ training examples $\bar{\mu}_{\mathbb{P}_{i}}=\mu_{\mathbb{P}_{i}}-\bar{\mu}_{\mathbb{P}},$ where $\bar{\mu}_{\mathbb{P}}=\frac{1}{N}\sum_{j=1}^{N}\mu_{\mathbb{P}_{j}}$. We consider $\mathbf{v}$ from a finite subspace spanned by a linear combination of $\bar{\mu}_{\mathbb{P}_{i}},\mathbf{v}=\sum_{i=1}^{N}\alpha_{i}\bar{\mu}_{
\mathbb{P}_{i}}.$ Then,

$$ $\displaystyle\hat{C}\mathbf{v}$ $\displaystyle=$ $\displaystyle\frac{1}{N}\sum_{i=1}^{N}(\bar{\mu}_{\mathbb{P}_{i}}\otimes\bar{ \mu}_{\mathbb{P}_{i}})\mathbf{v}=\frac{1}{N}\sum_{i=1}^{N}<\mathbf{v},\bar{\mu }_{\mathbb{P}_{i}}>\bar{\mu}_{\mathbb{P}_{i}}=\frac{1}{N}\sum_{i=1}^{N}\sum_{k =1}^{N}\alpha_{k}<\bar{\mu}_{\mathbb{P}_{k}},\bar{\mu}_{\mathbb{P}_{i}}>\bar{ \mu}_{\mathbb{P}_{i}}$ (21) $\displaystyle=$ $\displaystyle\lambda\mathbf{v}=\lambda\sum_{i=1}^{N}\alpha_{i}\bar{\mu}_{ \mathbb{P}_{i}},$ $$

from which we get

$$ $\displaystyle\sum_{k=1}^{N}\alpha_{k}<\bar{\mu}_{\mathbb{P}_{k}},\bar{\mu}_{ \mathbb{P}_{i}}>-N\lambda\alpha_{i}$ $\displaystyle=$ $\displaystyle 0,\quad i=1,2,\cdots,N$ (22) $\displaystyle(\mathbf{M}-N\lambda\mathbf{I})\bm{\alpha}$ $\displaystyle=$ $\displaystyle\mathbf{0},$ (23) $$

where $\mathbf{M}_{ij}=<\bar{\mu}_{\mathbb{P}_{i}},\bar{\mu}_{\mathbb{P}_{j}}>$. Note that

$$ $\displaystyle<\bar{\mu}_{\mathbb{P}_{i}},\bar{\mu}_{\mathbb{P}_{j}}>$ $\displaystyle=\left<\mu_{\mathbb{P}_{i}}-\frac{1}{N}\sum_{k=1}^{N}\mu_{\mathbb {P}_{k}},\mu_{\mathbb{P}_{j}}-\frac{1}{N}\sum_{m=1}^{N}\mu_{\mathbb{P}_{m}}\right>$ $\displaystyle=\left<\mu_{\mathbb{P}_{i}},\mu_{\mathbb{P}_{j}}\right>+\frac{1}{ N^{2}}\sum_{k=1}^{N}\sum_{m=1}^{N}\left<\mu_{\mathbb{P}_{k}},\mu_{\mathbb{P}_{ m}}\right>-\frac{1}{N}\sum_{k=1}^{N}\left<\mu_{\mathbb{P}_{k}},\mu_{\mathbb{P} _{j}}\right>-\frac{1}{N}\sum_{m=1}^{N}\left<\mu_{\mathbb{P}_{m}},\mu_{\mathbb{ P}_{i}}\right>.$ (24) $$

It’s clear that $\bm{\alpha}$ is eigenvector of $\mathbf{M}-N\lambda\mathbf{I}$ with eigenvalue $N\lambda$. Hence, we can solve the first $N_{x}$ eigenvectors $\bm{\alpha}^{1},\cdots,\bm{\alpha}^{N_{x}}$ and use them to compute the first $N_{x}$ eigenvectors $\mathbf{v}$ for $\hat{C}$ the spans of which covers enough information of $\{\bar{\mu}_{\mathbb{P}_{i}}\}$. Then for $\mathbb{P}_{i}$, the input for the encoder of NOMAD can be evaluated as $[<\mathbf{v}_{1},\bar{\mu}_{\mathbb{P}_{i}}>,\cdots,<\mathbf{v}_{N_{x}},\bar{
\mu}_{\mathbb{P}_{i}}>]$. Note that $<\mathbf{v}_{k},\bar{\mu}_{\mathbb{P}_{i}}>=\sum_{j=1}^{N}\alpha_{j}^{k}<\bar{
\mu}_{\mathbb{P}_{j}},\bar{\mu}_{\mathbb{P}_{i}}>$. For a testing distribution $\mathbb{P}_{t},$
we need to compute

$$ $\displaystyle<\mathbf{v}_{k},\bar{\mu}_{\mathbb{P}_{t}}>=\sum_{j=1}^{N}\alpha_ {j}^{k}<\bar{\mu}_{\mathbb{P}_{j}},\bar{\mu}_{\mathbb{P}_{t}}>=\sum_{j=1}^{N} \alpha_{j}^{k}\left<\mu_{\mathbb{P}_{t}}-\frac{1}{N}\sum_{k=1}^{N}\mu_{\mathbb {P}_{k}},\mu_{\mathbb{P}_{j}}-\frac{1}{N}\sum_{m=1}^{N}\mu_{\mathbb{P}_{m}}\right>$ $\displaystyle=\sum_{j=1}^{N}\alpha_{j}^{k}\left[\left<\mu_{\mathbb{P}_{t}},\mu _{\mathbb{P}_{j}}\right>+\frac{1}{N^{2}}\sum_{k=1}^{N}\sum_{m=1}^{N}\left<\mu_ {\mathbb{P}_{k}},\mu_{\mathbb{P}_{m}}\right>-\frac{1}{N}\sum_{k=1}^{N}\left< \mu_{\mathbb{P}_{k}},\mu_{\mathbb{P}_{j}}\right>-\frac{1}{N}\sum_{m=1}^{N} \left<\mu_{\mathbb{P}_{m}},\mu_{\mathbb{P}_{t}}\right>\right].$ (25) $$

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: It is very possible that I am missing some crucial details (maybe due to the readability issues). If that is the case, I will be more than happy to raise the rating. Otherwise, it seems to me that the paper is sort of reinventing conditional models, and does so very poorly written. With that said, I do find a contribution in the combination of representing continuous distributions and using them as a condition for score-based models. Currently that is not the story told by the authors. If it were, then I'm not sure, but perhaps It could be a reasonable contribution, subject to pin pointing exactly what existed before and what techniques that make it possible are novel.
soundness: 3
strengths: 1. The paper ambitiously addresses the challenge of generalizing score-based generative models across multiple distributions, relevant for few-shot learning and adaptive AI. Deriving a new model is refreshing and should be appreciated.

2. It introduces a mathematical framework using NOMAD for continuous distribution mapping, combining operator learning with generative modeling. This is a good idea to my taste, and a non-trivial contribution. 

3. The use of KME and PCA for embedding distributions is quite elegant.

4. The authors demonstrate their method for latent space models as well. This is useful and important.
summary: The authors propose an improved score based model, capable of generalizing to unseen continuous distributions that are given as an input at test time. The distributions have a special representation and the model uses their continuous representation to calculate the score function. The method is demonstrated for several examples such as 2d toy examples, MNIST double digits. They also demonstrate using a single example or few of them to construct the distribution representation and get reasonable results.
weaknesses: ## 1) Readability
Usually one doesn't mark the first weakness on the list to be readability, but in this caseI find it crucial. I was having a very hard time to read this paper and follow the derivations.
1. First the structure is hard to follow. I would expect starting from high level and then get in to low level derivations. I don't mean the intro as the high level, but for example in section 4, I would expect to see clearly what is the input, output and goal.
2. I think some simple examples are also needed for inputs and expected outputs.
3. A figure describing the method is also missing. 
4. In the background I think a brief explanation about NOMAD is needed.
5. Taking Figure 1 as a test case: It took me unreasonable amount of time to understand it. The related text in sec 5.2 is confusing and almost cryptic. It is not clear what is Gaussian there, no motivation for this very specific choice of data. No visual explanation in the figure itself. The color coding is not clear to me. The caption says "the first two are derived from the training sets, and the last two are from the testing sets." does that mean the left two and right two? If so why not putting some lines making a table with titles? Also if that is the case, I don't understand why the rightmost column has examples that are "left col right row" which is a training characteristic, not test.

## 2) Comparing to conditional generative models
1. Simply defining the goal, I did not see any definition that makes the premise of this work actually different from a conditional generative model. Conditional generative models in many cases get conditions that they were never trained on. They produce samples from a conditional distribution that I think can be any one of the examples in the experiments. In general I think that explaining why the proposed method is different/better than conditional models should be the most important claim to justify in this paper.
2. The authors do use a baseline they refer to as a conditional model. To my understanding it is something else. A model that is fine-tuned on the unseen distributions. First, this is not what a conditional model is- a conditional model gets a condition along with its input. Moreover, comparing to fine-tuned models implemented by the authors is always problematic. We don't have any details about how many iterations, what was the procedure, what weights trained etc. Without pointing any finger to any.one, usually authors don't have a huge motivation to make a baseline work better and fine-tuning is very tricky.
3. Mathematically, while I appreciate the interesting derivation, as far as I see, the objective boils down to a standard conditional setup. eq.13 is eq.3 with an additional conditioned term, sampled from a distribution of conditions. $s_\theta(u^v,z(t),t)$. I think any condition can be framed as a representation of a distribution. If the novelty is in the specific representation, then the premise should have been around it and indicating differences from the representation used in NOMAD. If not, then the contribution of the paper is combining NOMAD with conditional score models.
4. l503: "conditional score-based generative model cannot generalize to unseen test distributions once trained." I don't agree with that statement. While it might be true for discrete class-conditioning, it is not true in most other cases. Conditioning can be done for example on images in various ways and of course can be applied with an unseen image. Text-to-image models are another example for conditioning that generalizes to not necessarily seen conditions. 

## 3) Experiments
1. In general experiments are done on very low-scale examples and very few. While I wouldn't expect generating HD videos with a novel method, I do think that more evidence for its capabilities is needed.
2. The toy example experiment is very specific and complicated. Why not choosing simple distributions? Like spirals or just Gaussians, as used in many generative models works? Seeing such examples raises the concern whether they were carefully chosen.
3. I'm not sure about the few-shot learning part. I think the fact that it is possible to use a single example actually very strongly shows that this is equivalent to a conditional model. If it really adhered accurately to this distribution it should, in fact, make a Dirac Delta over the given example and generate only that. Thinking about it as a condition also makes this application less impressive, as single images as conditions are common.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 1
questions: 1. What advantages does SNO offer over existing baselines in practical applications? 
- For few-shot generation tasks, using models like StyleGAN with fine-tuning can yield high-quality results with less complexity. SNO does not clearly demonstrate an advantage over these existing models in terms of efficiency or output quality for few-shot tasks. 
2. Is SNO truly scalable and beneficial for large datasets?
- If SNO aims to generalize across complex distributions, especially in large-scale benchmarks, it remains unclear whether its assumptions hold or if the model would perform meaningfully better than traditional score-based models on large, complex datasets. Without evidence of SNO’s scalability and robustness to diverse data, the practical benefits seem limited.
3. How does it differ from Latent Diffusion Model?
- It seems that it is lot similar to LDM. LDM also moves the image into latent space, then conducts diffusion process.
soundness: 2
strengths: A notable strength of this approach is its similarity to Latent Diffusion, enabling the simultaneous training of the VAE and SNO. This joint training allows for a meaningful generation schema even with limited data, making it effective in scenarios where data is sparse. The use of score matching in the latent space further enhances the quality of generated samples, leveraging the efficiency of learning in a compressed representation.
summary: This paper proposes the Score Neural Operator (SNO) algorithm, addressing the limitations of existing generative models that often rely on overly simplified distributions. SNO first uses a Variational Autoencoder (VAE) to map data into a latent space and subsequently applies score matching to shift the distribution toward a Gaussian. The algorithm demonstrates effective performance on toy examples, such as MNIST and Gaussian Mixture Models (GMMs).”
weaknesses: Discrepancy Between Claims and Experimental Evidence:
- While the paper claims that SNO can improve performance across various distributions, the chosen datasets—GMM and MNIST—are actually toy dataset. These limited datasets are insufficient to convincingly demonstrate the model’s effectiveness for handling complex and high-dimensional distributions. If SNO’s strength lies in handling diverse and complex data, additional experiments with richer and more challenging datasets are necessary to substantiate this claim.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: Please use \citet or \citep appropriately to avoid sentences like: "Song and Ermon Song & Ermon (2019) introduced .... "


What are the exact inputs and outputs to the score neural operator? Why does the neural operator generalize to the new score function?
soundness: 2
strengths: The problem of learning a shared generative model across different data distributions is a relevant and important problem with significant interest from the community. 

The idea of using neural operators for this task seems interesting.
summary: This paper introduces the score neural operator, an architecture that learns to map multiple probability distributions to their score functions, enabling generalization to unseen distributions. The ideas are developed in the latent space and are examined on toy and small datasets such as Gaussian Mixture Models or 1024-dimensional MNIST.
weaknesses: - The experiments are conducted in toy and small datasets such as Gaussian Mixture Models or 1024-dimensional MNIST. From my experience working with these datasets, it is hard to expect that findings in these datasets naturally generalize to real-world datasets. 

- The idea of learning one generative model across different data distributions was originally discussed in "Distribution Augmentation for Generative Modeling" by Jun et al. ICML 2020. This paper shows that one can learn one generative across different distributions by providing a simple input conditioning that specifies the target distribution. This idea has been explored in diffusion modeling literature as well (for example see "Elucidating the Design Space of Diffusion-Based Generative Models" by Karras et al. NeurIP 2022). This submission does not discuss the idea of specifying the data distribution via simple embedding. It also does not discuss why a simple approach like passing a data distribution embedding would enable learning a diffusion model across different distributions.  

- The paper is hard to follow. Section 3.2 suddenly introduces many different concepts without explaining how they fit into the bigger picture. 

- Equation 18 trains the VAE and SGM prior jointly. However, the KL is evaluated with respect to a standard Normal prior in Eq. 16 while the prior is set to be a diffusion model.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see Weaknesses
soundness: 3
strengths: The proposed formulation is novel and has the scope to be applicable to multiple scenarios if successful. Although my expertise is not in the space of neural operators, the theoretical formulation presented is well established and correct, to the best of my understanding.
summary: The authors propose a Score neural operator, that learns to model the score in a distribution-agnostic way, thereby potentially allowing for the score of unseen distributions to be estimated with very few training samples.
weaknesses: The paper has a few weaknesses that could be addressed:

-  Writing: In general, the writing could use some improvement. For example, citations are poorly done, with incorrect usage of \citep and \citet (L34-42,L59,L46-50, etc.) especially in places where the authors themselves write a paper’s author names, followed by \citet. (E.g., L86,89,91). These could be fixed. I see similar issues with equations (L129).

-  Math: Sometimes, the notations are defined when they first appear, and this could cause confusion. For example, the definition of $\mathbf{u}$ in L210 is not accompanied by a definition of what it represents. Similarly, the use of $\mathbf{G}$ in Eqn. 14, or $\nu$ in Eqn 12. In most cases, this is described eventually, say half a page later, but this does not yield a good reading experience and could be fixed. 

- Experimental validation: Though the proposed idea and methodology seems promising, the experiments, in my opinion, don’t live up to the claim. Unless I am misunderstanding, It would be best to either make the claims more tamer, or have experiments the validate the existing claim that the SNO would be able to work few-shot with any data distribution. The experiments on 2-digit MNIST with a hold out set do not seem convincing. Even if they chose to work with lower resolution images, it would be more appealing to see actual generalization across datasets. For example, how about using SVHN (greyscale) and MNIST. Both being digit datasets, the task would be manageably hard, but possible also able to justify that few-shot alignment to a brand new dataset is indeed possible.

---

confidence: 4
contribution: 2
initial_rating: 1
presentation: 1
questions: Line 47: I would recommend adding "and by reversing the stochastic differential equation" or something like this. Very few people using Langevin anymore.

Line 142: Right now there is background and new upcoming parts in the background section. Its fine, but it would be nice to more clearly separate prior from new work. Adding let say a small heading like a \paragraph or subsection before new information would be nice to clearly separate new from old.

Line 212: μ_P goes to μˆ_Pi with no explanation (what is Pi?), then to P_t (what is t?), then to the k-th element of u (not μ, idk what this is). The equation looks like it contains MMD, but it comes out of nowhere. It refers to Appendix. But not where in the Appendix. I found it in A.3, but this really doesn't help me understand better. The u is also written differently then in equation 12. Its badly presented and confusing.

Line 218: N is a very weird choice for neural network notation. I suggest changing it since N is also for N(0,1). Now v is the measure, before it was u_P and u_Q. I have no idea why it changed again.

Line 235: Would be great to define u^v in line 220... It has never been defined, but I can guess that it is just N(x).

Line 232: NOMAD is only mentioned in the Related Work, which is an optional section. NOMAD needs to be in the background.

Line 249: Remove the parenthesis of the citations.

Line 255: Would be nice to say that these prior work use VAE with KL and following this, you do it too because otherwise a reader unfamiliar with StableDiffusion might be wondering where did the KL regularization came from and if you invented it.

Line 277: I would like to remind you that u^v has never been defined, we can only guess what it is. And now equation 14 and 17 have u^v, but the text mention u! What is happening here! I finally understood that u^v is equation 12, its just a typo, its missing the v. Same thing for Equation 11, which incorrectly defined, since it doesn't even have a v in it. This finally make sense. This makes me understand equation 17. 

Line 298: Doesn't there exists FID networks for MNIST in the literature? Just wondering. But the approach of training your own is fine.
soundness: 2
strengths: Simple way to get zero-shot generation on new distributions. Has potential for interesting applications (not tackled in this paper, which only deal with toy problems).
summary: The authors tackle the problem of generating sampling from unseen distributions using score-based approaches at inference. They propose to use the encoder of the VAE to embed specific samples. Then they represent the embedding of specific dataset/data-distribution as the average of the sample embeddings for this dataset. Then conditional on the distribution embedding from a collection of samples on a sampled distribution, they generate new samples of that distribution with score matching. The goal is to do some kind of meta-learning, i.e.: embed a small dataset and then have their model generate more samples of that distribution (unseen in the training, so OOD) in a zero-shot manner. Sadly, the abstract and title makes it very difficult to understand. Their specific choice of average embedding (called prototype) gives them the best generalization on toy experiments (2D Gaussian and paired MNIST digits).
weaknesses: The paper is extremely poorly written. It makes it extra hard to understand. Definitions are unclear and change over the course of the paper. Its a guessing game to understand what is what. I had an horrible time trying to understand the paper. Readers who are not reviewers will clearly just skip this paper given how badly written it is. 

The paper title and abstract needs refactoring to be comprehensible when reading for the first time. The paper needs to actually define and explains things properly.

Experiments are nice but very toy-ish, 2D Gaussians and pairs of MNIST digits. You need something like ImageNet with some classes never seen. CIFAR-100 also could be a start, remove like 5 classes from it for test time. This is whole idea of out-of-distribution generalization. There are existing benchmarks for that, but its more for classification. But you could just reuse them for generation.See Domainbed, https://github.com/facebookresearch/DomainBed, you train on some domains and not on others. This would be perfect to test your model. Also it has been shown that diffusion models can do zero-shot classification (https://arxiv.org/abs/2303.16203), you could even test your approach on the actual DomainBed challenge for classification using your score-model! Sky is the limit, but right now the paper is not ambitious enough and its not looking into interesting applications. With a bit more ambition, this could be a very interesting paper.

It needs more extensive literature review for meta-learning approaches.