SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ling Yang $^{1*}$  Jingwei Liu $^{1\dagger}$  Shenda Hong $^{1}$  Zhilong Zhang $^{1}$  Zhilin Huang $^{2}$  Zheming Cai $^{1}$  Wentao Zhang $^{1}$  Bin Cui $^{1\dagger}$
$^{1}$ Peking University  $^{2}$  Tsinghua University
yangling0818@163.com, jingweiliu1996@163.com, zhilong.zhang@bjmu.edu.cn
{hongshenda, wentao.zhang, bin.cui}@pku.edu.cn
Diffusion models are a new class of generative models, and have dramatically promoted image generation with unprecedented quality and diversity. Existing diffusion models mainly try to reconstruct input image from a corrupted one with a pixel-wise or feature-wise constraint along spatial axes. However, such point-based reconstruction may fail to make each predicted pixel/feature fully preserve its neighborhood context, impairing diffusion-based image synthesis. As a powerful source of automatic supervisory signal, context has been well studied for learning representations. Inspired by this, we for the first time propose CONPREDIFF to improve diffusion-based image synthesis with context prediction. We explicitly reinforce each point to predict its neighborhood context (i.e., multi-stride features/tokens/pixels) with a context decoder at the end of diffusion denoising blocks in training stage, and remove the decoder for inference. In this way, each point can better reconstruct itself by preserving its semantic connections with neighborhood context. This new paradigm of CONPREDIFF can generalize to arbitrary discrete and continuous diffusion backbones without introducing extra parameters in sampling procedure. Extensive experiments are conducted on unconditional image generation, text-to-image generation and image inpainting tasks. Our CONPREDIFF consistently outperforms previous methods and achieves a new SOTA text-to-image generation results on MS-COCO, with a zero-shot FID score of 6.21.
Recent diffusion models [99, 5, 63, 4, 47, 10, 22] have made remarkable progress in image generation. They are first introduced by Sohl-Dickstein et al. [76] and then improved by Song & Ermon [78] and Ho et al. [28], and can now generate image samples with unprecedented quality and diversity [24, 68, 67]. Numerous methods have been proposed to develop diffusion models by improving their empirical generation results [53, 77, 79] or extending the capacity of diffusion models from a theoretical perspective [80, 81, 47, 46, 108]. We revisit existing diffusion models for image generation and break them into two categories, pixel- and latent-based diffusion models, according to their diffusing spaces. Pixel-based diffusion models directly conduct continuous diffusion process in the pixel space, they incorporate various conditions (e.g., class, text, image, and semantic map) [29, 70, 51, 2, 66] or auxiliary classifiers [80, 14, 27, 54, 40] for conditional image generation.
On the other hand, latent-based diffusion models [65] conduct continuous or discrete diffusion process [87, 30, 1] on the semantic latent space. Such diffusion paradigm not only significantly reduces the
*Contact: Ling Yang, yangling0818@163.com.
$^\dagger$  Contributed equally.
$^{\ddagger}$ Corresponding Authors: Wentao Zhang, Bin Cui.
computational complexity for both training and inference, but also facilitates the conditional image generation in complex semantic space [62, 38, 58, 19, 98]. Some of them choose to pre-train an autoencoder [41, 64] to map the input from image space to the continuous latent space for continuous diffusion, while others utilize a vector quantized variational autoencoder [88, 17] to induce the token-based latent space for discrete diffusion [24, 75, 114, 85].
Despite all these progress of pixel- and latent-based diffusion models in image generation, both of them mainly focus on utilizing a point-based reconstruction objective over the spatial axes to recover the entire image in diffusion training process. This point-wise reconstruction neglects to fully preserve local context and semantic distribution of each predicted pixel/feature, which may impair the fidelity of generated images. Traditional non-diffusion studies [15, 45, 32, 50, 110, 8] have designed different context-preserving terms for advancing image representation learning, but few researches have been done to constrain on context for diffusion-based image synthesis.
In this paper, we propose CONPREDIFF to explicitly force each pixel/feature-token to predict its local neighborhood context (i.e., multi-stride features/tokens/pixels) in image diffusion generation with an extra context decoder near the end of diffusion denoising blocks. This explicit context prediction can be extended to existing discrete and continuous diffusion backbones without introducing additional parameters in inference stage. We further characterize the neighborhood context as a probability distribution defined over multi-stride neighbors for efficiently decoding large context, and adopt an optimal-transport loss based on Wasserstein distance [21] to impose structural constraint between the decoded distribution and the ground truth. We evaluate the proposed CONPREDIFF with the extensive experiments on three major visual tasks, unconditional image generation, text-to-image generation, and image inpainting. Notably, our CONPREDIFF consistently outperforms previous diffusion models by a large margin regarding generation quality and diversity.
Our main contributions are summarized as follows: (i): To the best of our knowledge, we for the first time propose CONPREDIFF to improve diffusion-based image generation with context prediction; (ii): We further propose an efficient approach to decode large context with an optimal-transport loss based on Wasserstein distance; (iii): CONPREDIFF substantially outperforms existing diffusion models and achieves new SOTA image generation results, and we can generalize our model to existing discrete and continuous diffusion backbones, consistently improving their performance.
Diffusion Models for Image Generation Diffusion models [99, 76, 78, 28] are a new class of probabilistic generative models that progressively destruct data by injecting noise, then learn to reverse this process for sample generation. They can generate image samples with unprecedented quality and diversity [24, 68, 67], and have been applied in various applications [99, 9, 6]. Existing pixel- and latent-based diffusion models mainly utilize the discrete diffusion [30, 1, 24] or continuous diffusion [87, 65] for unconditional or conditional image generation [80, 14, 27, 54, 40, 68]. Discrete diffusion models were also first described in [76], and then applied to text generation in Argmax Flow [30]. D3PMs [1] applies discrete diffusion to image generation. VQ-Diffusion [24] moves discrete diffusion from image pixel space to latent space with the discrete image tokens acquired from VQ-VAE [88]. Latent Diffusion Models (LDMs) [87, 65] reduce the training cost for high resolution images by conducting continuous diffusion process in a low-dimensional latent space. They also incorporate conditional information into the sampling process via cross attention [89]. Similar techniques are employed in DALLE-2 [62] for image generation from text, where the continuous diffusion model is conditioned on text embeddings obtained from CLIP latent codes [59]. Imagen [68] implements text-to-image generation by conditioning on text embeddings acquired from large language models (e.g., T5 [60]). Despite all this progress, existing diffusion models neglect to exploit rich neighborhood context in the generation process, which is critical in many vision tasks for maintaining the local semantic continuity in image representations [111, 45, 32, 50]. In this paper, we firstly propose to explicitly preserve local neighborhood context for diffusion-based image generation.
Context-Enriched Representation Learning Context has been well studied in learning representations, and is widely proved to be a powerful automatic supervisory signal in many tasks. For example, language models [52, 13] learn word embeddings by predicting their context, i.e., a few words before and/or after. More utilization of contextual information happens in visual tasks, where spatial context
Figure 1: In training stage, CONPREDIFF first performs self-denoising as standard diffusion models, then it conducts neighborhood context prediction based on denoised point  $\boldsymbol{x}_{t-1}^i$ . In inference stage, CONPREDIFF only uses its self-denoising network for sampling.
is vital for image domain. Many studies [15, 111, 45, 32, 50, 110, 8, 106, 95, 94, 44] propose to leverage context for enriching learned image representations. Doersch et al. [15] and Zhang et al. [110] make predictions from visible patches to masked patches to enhance the self-supervised image representation learning. Hu et al. [32] designs local relation layer to model the context of local pixel pairs for image classification, while Liu et al. [45] preserves contextual structure to guarantee the local feature/pixel continuity for image inpainting. Inspired by these studies, in this work, we propose to incorporate neighborhood context prediction for improving diffusion-based generative modeling.
Discrete Diffusion We briefly review a classical discrete diffusion model, namely Vector Quantized Diffusion (VQ-Diffusion) [24]. VQ-Diffusion utilizes a VQ-VAE to convert images  $x$  to discrete tokens  $x_0 \in \{1, 2, \dots, K, K + 1\}$ ,  $K$  is the size of codebook, and  $K + 1$  denotes the [MASK] token. Then the forward process of VQ-Diffusion is given by:
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = \boldsymbol {v} ^ {\top} \left(x _ {t}\right) \boldsymbol {Q} _ {t} \boldsymbol {v} \left(x _ {t - 1}\right) \tag {1}
$$
where  $\pmb{v}(x)$  is a one-hot column vector with entry 1 at index  $x$ . And  $\pmb{Q}_t$  is the probability transition matrix from  $x_{t-1}$  to  $x_t$  with the mask-and-replace VQ-Diffusion strategy. In the reverse process, VQ-Diffusion trains a denoising network  $p_\theta(\pmb{x}_{t-1}|\pmb{x}_t)$  that predicts noiseless token distribution  $p_\theta(\tilde{\pmb{x}}_0|\pmb{x}_t)$  at each step:
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) = \sum_ {\tilde {\boldsymbol {x}} _ {0} = 1} ^ {K} q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}, \tilde {\boldsymbol {x}} _ {0}\right) p _ {\theta} \left(\tilde {\boldsymbol {x}} _ {0} \mid \boldsymbol {x} _ {t}\right), \tag {2}
$$
which is optimized by minimizing the following variational lower bound (VLB) [76]:
$$
\mathcal {L} _ {t - 1} ^ {\text {d i s}} = D _ {K L} \left(q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} _ {0}\right) \mid p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right)\right). \tag {3}
$$
Continuous Diffusion A continuous diffusion model progressively perturbs input image or feature map  $\pmb{x}_0$  by injecting noise, then learn to reverse this process starting from  $\pmb{x}_T$  for image generation. The forward process can be formulated as a Gaussian process with Markovian structure:
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t}; \sqrt {1 - \beta_ {t}} \boldsymbol {x} _ {t - 1}, \beta_ {t} \mathbf {I}\right), \tag {4}
$$
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {0}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t}; \sqrt {\bar {\alpha} _ {t}} \boldsymbol {x} _ {0}, (1 - \bar {\alpha} _ {t}) \mathbf {I}\right),
$$
where  $\beta_{1},\ldots ,\beta_{T}$  denotes fixed variance schedule with  $\alpha_{t}:= 1 - \beta_{t}$  and  $\overline{\alpha}_t\coloneqq \prod_{s = 1}^t\alpha_s$ . This forward process progressively injects noise to data until all structures are lost, which is well approximated by  $\mathcal{N}(0,\mathbf{I})$ . The reverse diffusion process learns a model  $p_{\theta}(\boldsymbol{x}_{t - 1}|\boldsymbol{x}_t)$  that approximates the
true posterior:
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t - 1}; \mu_ {\theta} \left(\boldsymbol {x} _ {t}\right), \Sigma_ {\theta} \left(\boldsymbol {x} _ {t}\right)\right), \tag {5}
$$
Fixing  $\Sigma_{\theta}$  to be untrained time dependent constants  $\sigma_t^2 I$ , Ho et al. [28] improve the diffusion training process by optimizing following objective:
$$
\mathcal {L} _ {t - 1} ^ {\text {c o n}} = \underset {q (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {t - 1})} {\mathbb {E}} \left[ \frac {1}{2 \sigma_ {t} ^ {2}} \| \mu_ {\theta} (\boldsymbol {x} _ {t}, t) - \hat {\mu} (\boldsymbol {x} _ {t}, \boldsymbol {x} _ {0}) \| ^ {2} \right] + C, \tag {6}
$$
where  $C$  is a constant that does not depend on  $\theta$ .  $\hat{\mu}(\boldsymbol{x}_t, \boldsymbol{x}_0)$  is the mean of the posterior  $q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_0, \boldsymbol{x}_t)$ , and  $\mu_\theta(\boldsymbol{x}_t, t)$  is the predicted mean of  $p_\theta(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t)$  computed by neural networks.
4 The Proposed CONPREDIFF
In this section, we elucidate the proposed CONPREDIFF as in Figure 1. In Sec. 4.1, we introduce our proposed context prediction term for explicitly preserving local neighborhood context in diffusion-based image generation. To efficiently decode large context in training process, we characterize the neighborhood information as the probability distribution defined over multi-stride neighbors in Sec. 4.2, and theoretically derive an optimal-transport loss function based on Wasserstein distance to optimize the decoding procedure. In Sec. 4.3, we generalize our CONPREDIFF to both existing discrete and continuous diffusion models, and provide optimization objectives.
4.1 Neighborhood Context Prediction in Diffusion Generation
We use unconditional image generation to illustrate our method for simplicity. Let  $\pmb{x}_{t - 1}^{i}\in \mathbb{R}^{d}$  to denote  $i$ -th pixel of the predicted image,  $i$ -th feature point of the predicted feature map, or  $i$ -th image token of the predicted token map in spatial axes. Let  $\mathcal{N}_i^s$  denote the  $s$ -stride neighborhoods of  $\pmb{x}_{t - 1}^{i}$ , and  $K$  denotes the total number of  $\mathcal{N}_i^s$ . For example, the number of 1-stride neighborhoods is  $K = 8$ , and the number of 2-stride ones is  $K = 24$ .
S-Stride Neighborhood Reconstruction Previous diffusion models make point-wise reconstruction, i.e., reconstructing each pixel, thus their reverse learning processes can be formulated by  $p_{\theta}(\pmb{x}_{t - 1}^{i}|\pmb{x}_{t})$ . In contrast, our context prediction aims to reconstruct  $\pmb{x}_{t - 1}^{i}$  and further predict its  $s$ -stride neighborhood contextual representations  $H_{\mathcal{N}_i^s}$  based on  $\pmb{x}_{t - 1}^{i}$ :  $p_{\theta}(\pmb{x}_{t - 1}^{i},H_{\mathcal{N}_i^s}|\pmb{x}_t)$ , where  $p_{\theta}$  is parameterized by two reconstruction networks  $(\psi_p,\psi_n)$ .  $\psi_{p}$  is designed for the point-wise denoising of  $\pmb{x}_{t - 1}^{i}$  in  $\pmb{x}_t$  and  $\psi_{n}$  is designed for decoding  $H_{\mathcal{N}_i^s}$  from  $\pmb{x}_{t - 1}^{i}$ . For denoising  $i$ -th point in  $\pmb{x}_t$ , we have:
$$
\boldsymbol {x} _ {t - 1} ^ {i} = \psi_ {p} (\boldsymbol {x} _ {t}, t), \tag {7}
$$
where  $t$  is the time embedding and  $\psi_p$  is parameterized by a U-Net or transformer with an encoder-decoder architecture. For reconstructing the entire neighborhood information  $H_{\mathcal{N}_i^s}$  around each point  $x_{t - 1}^i$ , we have:
$$
\boldsymbol {H} _ {\mathcal {N} _ {i} ^ {s}} = \psi_ {n} \left(\boldsymbol {x} _ {t - 1} ^ {i}, t\right) = \psi_ {n} \left(\psi_ {p} \left(\boldsymbol {x} _ {t}, t\right)\right), \tag {8}
$$
where  $\psi_n \in \mathbb{R}^{Kd}$  is the neighborhood decoder. Based on Equation (7) and Equation (8), we unify the point- and neighborhood-based reconstruction to form the overall training objective:
$$
\mathcal {L} _ {\mathrm {C O N P R E D I F F}} = \sum_ {i = 1} ^ {x \times y} \left[ \underbrace {\mathcal {M} _ {p} \left(\boldsymbol {x} _ {t - 1} ^ {i} , \hat {\boldsymbol {x}} ^ {i}\right)} _ {\text {p o i n t d e n o i s i n g}} + \underbrace {\mathcal {M} _ {n} \left(\boldsymbol {H} _ {\mathcal {N} _ {i} ^ {s}} , \hat {\boldsymbol {H}} _ {\mathcal {N} _ {i} ^ {s}}\right)} _ {\text {c o m t e x t p r e d i c t i o n}} \right], \tag {9}
$$
where  $x, y$  are the width and height on spatial axes.  $\hat{\pmb{x}}^i (\hat{\pmb{x}}_0^i)$  and  $\hat{H}_{\mathcal{N}_i^s}$  are ground truths.  $\mathcal{M}_p$  and  $\mathcal{M}_n$  can be Euclidean distance. In this way, CONPREDIFF is able to maximally preserve local context for better reconstructing each pixel/feature-token.
Interpreting Context Prediction in Maximizing ELBO We let  $\mathcal{M}_p, \mathcal{M}_n$  be square loss,  $\mathcal{M}_n(\boldsymbol{H}_{\mathcal{N}_i^s}, \hat{\boldsymbol{H}}_{\mathcal{N}_i^s}) = \sum_{j \in \mathcal{N}_i} (\boldsymbol{x}_0^{i,j} - \hat{\boldsymbol{x}}_0^{i,j})^2$ , where  $\hat{\boldsymbol{x}}_0^{i,j}$  is the  $j$ -th neighbor in the context of  $\hat{\boldsymbol{x}}_0^i$  and  $\boldsymbol{x}_0^{i,j}$  is the prediction of  $\boldsymbol{x}_0^{i,j}$  from a denoising neural network. Thus we have:
$$
\boldsymbol {x} _ {0} ^ {i, j} = \psi_ {n} \left(\psi_ {p} \left(\boldsymbol {x} _ {t}, t\right) (i)\right) (j). \tag {10}
$$
Compactly, we can write the denoising network as:
$$
\Psi \left(x _ {t}, t\right) (i, j) = \left\{ \begin{array}{l l} \psi_ {n} \left(\psi_ {p} \left(\boldsymbol {x} _ {t}, t\right) (i)\right) (j), & j \in \mathcal {N} _ {i}, \\ \psi_ {p} \left(\boldsymbol {x} _ {t}, t\right) (i), & j = i. \end{array} \right. \tag {11}
$$
We will show that the DDPM loss is upper bounded by ConPreDiff loss, by reparameterizing  $\pmb{x}_0(\pmb{x}_t,t)$ . Specifically, for each unit  $i$  in the feature map, we use the mean of predicted value in its neighborhood as the final prediction:
$$
\boldsymbol {x} _ {0} \left(\boldsymbol {x} _ {t}, t\right) (i) = 1 / \left(| \mathcal {N} _ {i} | + 1\right) * \sum_ {j \in \mathcal {N} _ {i} \cup \{i \}} \Psi \left(\boldsymbol {x} _ {t}, t\right) (i, j). \tag {12}
$$
Now we can show the connection between the DDPM loss and ConPreDiff loss:
$$
\begin{array}{l} | | \hat {\pmb {x}} _ {0} - \pmb {x} _ {0} (\pmb {x} _ {t}, t) | | _ {2} ^ {2} = \sum_ {i} (\hat {\pmb {x}} _ {0} ^ {i} - \pmb {x} _ {0} (\pmb {x} _ {t}, t) (i)) ^ {2}, \\ = \sum_ {i} (\hat {\boldsymbol {x}} _ {0} ^ {i} - \sum_ {j \in \mathcal {N} _ {i} \cup \{i \}} \Psi (\boldsymbol {x} _ {t}, t) (i, j) / (| \mathcal {N} _ {i} | + 1)) ^ {2}, \\ = \sum_ {i} \left(\sum_ {j \in \mathcal {N} _ {i} \cup \{i \}} \left(\Psi \left(\boldsymbol {x} _ {t}, t\right) (i, j) - \hat {\boldsymbol {x}} _ {0} ^ {i})\right) ^ {2} / \left(| \mathcal {N} _ {i} | + 1\right) ^ {2}, \right. \tag {13} \\ \end{array}
$$
(Cauchy Inequality)  $\leq \sum_{i}\sum_{j\in \mathcal{N}_{i}\cup \{i\}}(\Psi (\pmb{x}_{t},t)(i,j) - \hat{\pmb{x}}_{0}^{i})^{2} / (|\mathcal{N}_{i}| + 1),$
$$
= 1 / (| \mathcal {N} _ {i} | + 1) \sum_ {i} [ (\hat {\boldsymbol {x}} _ {0} ^ {i} - \psi_ {p} (\boldsymbol {x} _ {t}, t) (i)) ^ {2} + \sum_ {j \in \mathcal {N} _ {i}} (\hat {\boldsymbol {x}} _ {0} ^ {i, j} - \boldsymbol {x} _ {0} ^ {i, j}) ^ {2} ]
$$
In the last equality, we assume that the feature is padded so that each unit  $i$  has the same number of neighbors  $|\mathcal{N}|$ . As a result, the ConPreDiff loss is an upper bound of the negative log likelihood.
**Complexity Problem** We note that directly optimizing the Equation (9) has a complexity problem and it will substantially lower the efficiency of CONPREDIFF in training stage. Because the network  $\psi_{n}:\mathbb{R}^{d}\to \mathbb{R}^{Kd}$  in Equation (8) needs to expand the channel dimension by  $K$  times for large-context neighborhood reconstruction, it significantly increases the parameter complexity of the model. Hence, we seek for another way that is efficient for reconstructing neighborhood information.
We solve the challenging problem by changing the direct prediction of entire neighborhoods to the prediction of neighborhood distribution. Specifically, for each  $\boldsymbol{x}_{t-1}^i$ , the neighborhood information is represented as an empirical realization of i.i.d. sampling  $Q$  elements from  $\mathcal{P}_{\mathcal{N}_i^s}$ , where  $\mathcal{P}_{\mathcal{N}_i^s} \triangleq \frac{1}{K} \sum_{u \in \mathcal{N}_i^s} \delta_{h_u}$ . Based on this view, we are able to transform the neighborhood prediction  $\mathcal{M}_n$  into the neighborhood distribution prediction. However, such sampling-based measurement loses original spatial orders of neighborhoods, and thus we use a permutation invariant loss (Wasserstein distance) for optimization. Wasserstein distance [23, 21] is an effective metric for measuring structural similarity between distributions, which is especially suitable for our neighborhood distribution prediction. And we rewrite the Equation (9) as:
$$
\mathcal {L} _ {\mathrm {C O N P R E D I F F}} = \sum_ {i = 1} ^ {x \times y} \left[ \underbrace {\mathcal {M} _ {p} \left(\boldsymbol {x} _ {t - 1} ^ {i} , \hat {\boldsymbol {x}} ^ {i}\right)} _ {\text {p o i n t d e n o i s i n g}} + \underbrace {\mathcal {W} _ {2} ^ {2} \left(\psi_ {n} \left(\boldsymbol {x} _ {t - 1} ^ {i} , t\right) , \mathcal {P} _ {\mathcal {N} _ {i} ^ {s}}\right)} _ {\text {n e i g h b o r h o o d d i s t r i b u t i o n p r e d i c t i o n}} \right], \tag {14}
$$
where  $\psi_{n}(\pmb{x}_{t - 1}^{i},t)$  is designed to decode neighborhood distribution parameterized by feedforward neural networks (FNNs), and  $\mathcal{W}_2(\cdot ,\cdot)$  is the 2-Wasserstein distance as defined below. We provide a more explicit formulation of  $\mathcal{W}_2^2 (\psi_n(\pmb{x}_{t - 1}^i,t),\mathcal{P}_{\mathcal{N}_i^s})$  in Sec. 4.2.
Definition 4.1. Let  $\mathcal{P}$ ,  $\mathcal{Q}$  denote two probability distributions with finite second moment defined on  $\mathcal{Z} \subseteq \mathbb{R}^m$ . The 2-Wasserstein distance between  $\mathcal{P}$  and  $\mathcal{Q}$  defined on  $\mathcal{Z}$ ,  $\mathcal{Z}' \subseteq \mathbb{R}^m$  is the solution to the optimal mass transportation problem with  $\ell_2$  transport cost [90]:
$$
\mathcal {W} _ {2} (\mathcal {P}, \mathcal {Q}) = \left(\inf  _ {\gamma \in \Gamma (\mathcal {P}, \mathcal {Q})} \int_ {\mathcal {Z} \times \mathcal {Z} ^ {\prime}} \| Z - Z ^ {\prime} \| _ {2} ^ {2} d \gamma (Z, Z ^ {\prime})\right) ^ {1 / 2} \tag {15}
$$
where  $\Gamma (\mathcal{P},\mathcal{Q})$  contains all joint distributions of  $(Z,Z^{\prime})$  with marginals  $\mathcal{P}$  and  $\mathcal{Q}$  respectively.
4.2 Efficient Large Context Decoding
Our CONPREDIFF essentially represents the node neighborhood  $\hat{H}_{\mathcal{N}_i^s}$  as a distribution of neighbors' representations  $\mathcal{P}_{\mathcal{N}_i^s}$  (Equation (14)). We adopt Wasserstein distance to characterize the distribution reconstruction loss because  $\mathcal{P}_{\mathcal{N}_i^s}$  has atomic non-zero measure supports in a continuous space, where the family of  $f$ -divergences such as KL-divergence cannot be applied. Maximum mean discrepancy may be applied but it needs to choose a specific kernel function.
We define the decoded distribution  $\psi_{n}(\pmb{x}_{t - 1}^{i},t)$  as an FNN-based transformation of a Gaussian distribution parameterized by  $\pmb{x}_{t - 1}^{i}$  and  $t$ . The reason for choosing this setting stems from the fact that the universal approximation capability of FNNs allows to (approximately) reconstruct any distributions in 1-Wasserstein distance, as formally stated in Theorem 4.2, proved in Lu & Lu [48]. To enhance the empirical performance, our case adopts the 2-Wasserstein distance and an FNN with  $d$  -dim output instead of the gradient of an FNN with 1-dim outout. Here, the reparameterization trick [42] needs to be used:
$$
\psi_ {n} \left(\boldsymbol {x} _ {t - 1} ^ {i}, t\right) = \operatorname {F N N} _ {n} (\xi), \xi \sim \mathcal {N} \left(\mu_ {i}, \Sigma_ {i}\right),
$$
$$
\mu_ {i} = \operatorname {F N N} _ {\mu} \left(\boldsymbol {x} _ {t - 1} ^ {i}, \Sigma_ {i} = \operatorname {d i a g} \left(\exp \left(\operatorname {F N N} _ {\sigma} \left(\boldsymbol {x} _ {t - 1} ^ {i}\right)\right)\right). \right. \tag {16}
$$