Versatile Neural Processes for Learning Implicit Neural Representations

Some NP methods encode the context points to produce local feature representations at target coordinates. The early work (Garnelo et al., 2018b) uses simple MLP layers to learn features of the context set and suffers from underfitting problem. Attentive NP (Kim et al., 2019) and Transformer NP (Nguyen & Grover, 2022) enhances the feature representation by introducing point-wise self-attention with each context point taken as a token. However, when the number of context points is large (e.g.in order to model an image signal with many details), the computational burden is heavy and such a framework is impractical.

We address this issue by simply introducing a set tokenization module (set tokenizer) followed by self-attention layers. The set tokenizer is instantiated by a set convolution layer (Zaheer et al., 2017) which transforms the neighboring context points to a token. Taking a 2D image as an example, with the kernel size of $k\times k$ and stride of $k$, the set tokenizer can reduce the number of sample points by $k^{2}$ times (assume the image resolution is an integer multiple of $k$). By using the set tokenizer, the complexity of attention layers is reduced from the first term to the second term as below:

where $L_{s}$ and $L_{c}$ are the number of the self-attention layers and the cross-attention layers, $N_{C}$ and $N_{T}$ are the number of context points and target points. The difference between the set convolution and the conventional convolution is that the former can handle missing points (e.g., in the application of inpainting) and the data that live "off the grid" (e.g.  time series data that observed irregularly at any time). The combination of set tokenizer and attention would empower our framework the flexibility in processing 3D signals, which is infeasible for ConvNP (Gordon et al., ; Foong et al., 2020) because ConvNP requires to preserve all the 3D grids, which is very inefficient in sparse 3D space.

In the previous NP methods (Garnelo et al., 2018b; Kim et al., 2019), they learn a single global latent Gaussian variable from the observed context set to model the distributions of a function. However, it potentially limits the expressiveness of the model. Intuitively, increasing the dimension of the latent variables may increase such flexibility, but in practice, this is not sufficient (Lee et al., 2020).

We design a hierarchical latent modulated decoder in order to better model the complex distribution of a function. We sequentially learn multiple global latent variables to modulate the parameters of MLP blocks. Figure 2 shows the details of the designed hierarchical structure.

The decoder consists of $L_{c}$ cross-attention blocks and $L_{K}$ modulated MLP blocks. With the target location $X_{T}$ as query and the context tokens from encoder as the keys and values, the cross-attention blocks output the target location features $\hat{Y}_{0}\in\mathbb{R}^{T\times d}$, where $T$ is the number of target coordinates and $d$ is the feature dimension. The sequential modulated MLP blocks enable the exploitation of hierarchical global latent variables for the approximation of a complex signal. The output of the final ($L_{K}^{th}$) modulated MLP block goes through two MLP layers to estimate the probability of $Y_{T}$.

Figure 2(b) illustrates the detailed network structure of a modulated MLP block. We build a Modulated MLP block by stacking two modulated MLP layers and two unmodulated MLP layers. This block refines $\hat{Y}_{k-1}$ to output features $\hat{Y}_{k}$, which is the input of the next Modulated MLP block. At the heart of each block is a latent variable $\mathbf{z}\in\mathbb{R}^{1\times d_{\mathbf{z}}}$ modeled by a Gaussian distribution. The generation process of $\mathbf{z}_{k}$ (of the $k^{th}$ block) can be formulated as follows:

where MLPs refer to two MLP layers and an intermediate ReLU activation layer. The prediction results $\hat{Y}_{k-1}$ from the previous block (i.e., the $(k-1)^{th}$ block) are used for calculating the conditional prior distribution of $\mathbf{z}_{k}$, i.e., $p_{\psi}(\mathbf{z}_{k}|\hat{Y}_{<k},D_{C},X_{T})$. During training, the conditional posterior distribution $q_{\phi}(\mathbf{z}_{k}|\hat{Y}_{<k},D_{C},D_{T})$ can be calculated as well, by incorporating the ground truth target signal values $Y_{T}$:

Marked with dashed lines, $Y_{T}$ only participates in training and cannot be accessed during inference.

After sampling the low-dimensional latent variable $\mathbf{z}_{k}$, we use the modulated fully-connected (ModFC) layer (Karras et al., 2020; 2021) to adjust the parameters of modulated MLP layers, taking a sampled realization of $\mathbf{z_{k}}$ as the style vector input. Unlike previous NP methods that concatenate the latent variable to every target coordinate, modulating the MLP parameters from the low-dimensional latent variables is a more flexible way to represent continuous functions (Sitzmann et al., 2020b). Please see Appendix A for more details for the mechanism of the modulated MLP.

We model the function representations by multiple latent variables $(\mathbf{z}_{1},\mathbf{z}_{2},...,\mathbf{z}_{L_{K}})$. The KL term in Eq. 3, which measures the mismatch between the approximated distribution and ground truth distribution, takes the hierarchical format as

where $q_{\phi}(\mathbf{z}_{<k}|D_{C},D_{T})=\prod_{i=1}^{k-1}q_{\phi}(\mathbf{z}_{i}|\hat{Y}_{<i},D_{C},D_{T})$ is the approximate posterior of $\mathbf{z}_{<k}$. This decomposed KL term and the reconstruction term formulate the objective:

where $\beta$ denotes a weight to balance the importance between the two terms. In our experiments on 2D and 3D signals, we multiply the KL term with a small weight $\beta$ to better capture the uncertainty of function distributions (Higgins et al., ). In addition, we emphasize that although there are some prior works designing hierarchical architecture for VAEs (Sønderby et al., 2016; Vahdat & Kautz, 2020; Child, 2021) or double latent variable models for NP (Wang & Van Hoof, 2020), our proposed hierarchical architecture is different in principle. Here, every latent variable is a low-dimensional vector obtained after average pooling. Therefore, the designed hierarchical architecture can deal with arbitrary target coordinates and thereby can be used for boosting the performance of approximating the global structure of continuous functions.

We implement the proposed VNP to learn the implicit neural representations for 1D signal regression. This classical 1D regression aims at predicting the function values at given target locations, conditioned on several observations of the samples from the function. Following (Kim et al., 2019), we measure the performance by considering the context set likelihood and target set likelihood, which reflects the context reconstruction error and target prediction error, respectively.

Settings. We train the models on synthetic functions drawn from prior function distributions synthesized with different kernels (RBF, Matern) by following (Gordon et al., ; Kim et al., 2022). For the evaluated methods, we employ importance weighted sampling (Burda et al., 2016) to evaluate the log likelihood, where the last four methods in Table 1 are measured by sampling latent variables from the posterior distribution and calculating the tighter ELBO with importance weighted sampling. For fair comparison, we manage to keep the network size comparable with that of other methods, which is completed by adjusting channel number. Please refer to Appendix B for detailed settings.

Comparison with the State-of-the-Arts. We compare our method with conditional neural process (CNP) (Garnelo et al., 2018a), (stacked) attentive neural process (ANP) (Kim et al., 2019) (with stacked self-attention layers), bootstrapping attentive neural process (BANP) (Lee et al., 2020), and convolutional neural process (ConvNP) (Foong et al., 2020). Table 1 shows the results in terms of log likelihood. Most of the models can provide satisfactory reconstruction results for context points, except CNP which suffers from underfitting problem (Garnelo et al., 2018b). On the context points, ours also achieves comparable performance. However, the prediction results of other methods on the unseen target points are much worse than that on the context sets. In contrast, our final model, VNP, outperforms previous approaches by a large margin at the target points on both of the function distributions. Note that although BANP attempts to increase the expressiveness of function representations by using latent variable bootstrapping from the perspective data resampling.

Ablation Study. We conduct a group of ablation study on this 1D regression task to investigate the effectiveness of different components. Based on ANP, we first build a Hierarchical Neural Process (HNP) with hierarchical latent variables. Note that similar to ANP, the implemented HNP still concatenates every target coordinate with the global latent variables. We observe significant improvements in Table 1 by comparing HNP with ANP, demonstrating that the hierarchical global latent variables boost the performance of function approximation. Secondly, the modulated MLP layer for implicit function parameterization can be further equipped in HNP, referred to as HNP-Mod. Compared with the concatenation of latent variable, using low-dimensional latent variables to modulate the network parameters enables more flexible function approximating. Based on HNP-Mod, we finally add the bottleneck encoder to learn conditional network inputs that are adaptive to target locations, building our final model VNP. In Table 1, we can find this powerful pre-processing encoder can further improve the performance, because the set tokenizer (set convolution here) can preserve locality to learn more appropriate features.

In Table 2, we further ablate on the detailed hierarchical structures in our model. As we can observe, when we increase the number of Modulated MLP blocks, the prediction performance (at target points) is improved until saturated with 6 blocks (i.e., $L_{K}=6$). In addition, the column with "6 / single $\mathbf{z}$" means we use exactly the same network structures as $L_{K}=6$, but instead only impose the KL constraint to the final latent variable $\mathbf{z}_{6}$. It is found that the performance will also drop dramatically, which verifies that the gains come from the hierarchical design instead of the increased network capacity. This ablation study guides us to set the number of Modulated MLP blocks as 6 to compare with other methods, which can keep a balance between the performance and the complexity.

Visualizations. We visualize the obtained function distributions from stacked ANP (Kim et al., 2019) and our VNP, by sampling the latent variables 20 times. Given partial observations of a signal, there exists uncertainty on the continuous function since there are many possible ways to interpret these observations (i.e., context set). An excellent Neural Process model should be able to model such uncertainty through the fitted functions. In other words, the generated functions conditioned on the context set should approximate the data distribution and cover the target set points. As shown in Figure 3, we can see that the generated functions from the stacked ANP cannot predict the target points accurately, e.g., those regions marked by orange. In contrast, our method provide good approximations for the continuous functions, where the distribution covers the groundtruth function.

A 2D image can be modeled by a continuous function that maps the 2D pixel coordinates to the color values. We implement our VNP framework for learning the continuous function of 2D images, which supports tasks such as image completion and super-resolution to arbitrary size.

Settings. We conduct experiments on CelebA dataset (Liu et al., 2015), mainly with the resized resolution of $64\times 64$. Due to the limited representation capability or the high requirement on computation resources, most previous NP methods (Kim et al., 2019; Lee et al., 2020) are in general trained and evaluated on relatively low resolution such as $32\times 32$. Our framework enables the processing of complex signal with higher resolution. We train a single model that supports image completion and super-resolution tasks. During training, we control the context ratio as 0.03, which means 3% pixels are taken as context set. The target ratio is set as 0.15 and the target set has partial intersection with the context set. More detailed experimental settings can be found in Appendix B.

Visualizations. We compare our VNP with stacked ANP (Kim et al., 2019) tested with three different context ratios respectively as shown in Figure 4. Our VNP generates diverse and realistic image completion results with fine details. In comparison, the results of Stacked ANP are blurred.

In Figure 5, we show that the proposed VNP achieves satisfactory results for image completion and superresolution on CelebA64 dataset. As a unified framework, with a single model, it can support image completion, superresolution, and the concurrent of completion and superresolution.

Influence of Patch/Kernel Size. We conduct an ablation study to investigate the influence of the patch/kernel size in our set tokenizer. We set the stride the same as the kernel size. The results are shown in Table 3. It is observed that using larger kernel size does not degrade the performance obviously. In addition, using hierarchical structure brings gains.

Complexity Comparisons. Thanks to our bottleneck encoder design, the proposed VNP is a practical framework for complex signal modeling, where the number of context points as condition is usually large. For comparison, we calculate the GFLOPs of stacked ANP and our VNP when target ratio is 0.25 on CelebA64 and CelebA178 respectively. Here, the kernel size of set tokenizer is 4 on CelebA64 and 10 on CelebA178. The computational complexity comparisons in terms of GFLOPs are shown in Table 4. Since we use set tokenizer to reduce the number of tokens and process the context set in image grids, the GFLOPs of VNP would be smaller than that of Stacked ANP, especially when the context ratio is high.

Implicit neural representations excel at representing 3D signals. However, since 3D signals are usually much more complex, most previous works (Mildenhall et al., 2020; Martin-Brualla et al., 2021) require expensive optimization to fit the neural networks to represent the scenes. We focus on the task of view synthesis in this section to evaluate our proposed VNP.

Settings. We follow the spirits of NeRF (Mildenhall et al., 2020) that fits a network to map the world coordinate into the corresponding RGB values and volume density. We use bottleneck encoder (tokenization from the image patches) to calculate the adaptive input of the MLPs, queried by target world coordinates. Then the hierarchically learned global latent variable will modulate the parameters of MLPs to predict the RGB values and volume density. Note that there is an extra volume rendering process inside each decoding block before the pooling module, because it requires transferring the world coordinate to the image coordinate to compute the latent distribution of $\mathbf{z}_{k}$. We conduct experiments on ShapeNet (Chang et al., 2015) objects, including three sub-datasets: cars, lamps, and chairs. More details on the network structures and hyper parameters can be found in Appendix B.

Comparisons. Our VNP model can generate the implicit neural representations w.r.t. the previously unseen scene in a single forward pass. In this 3D task, one prior work NeRF-VAE (Kosiorek et al., 2021) can also generate the scene from the randomly sampled latent variable. We make quantitative and qualitative comparisons with NeRF-VAE. In addition, we compare with (Chen et al., 2021b) and (Chen & Wang, 2022) which can also learn implicit neural representations of the previously unseen scene, although they are not in the family of probabilistic models and thus cannot model function distributions. As shown in Table 5, our method achieves the best performance in terms of Peak Signal-to-Noise Ratio (PSNR) on Cars and Lamps, and is comparable with Trans-INR (Chen & Wang, 2022) on Chairs. The visualizations in Figure 6(a) show that VNP produces much better novel-view prediction results than NeRF-VAE, which is also a probabilistic generative model.