Patch Attentive Neural Processes (PANP)

• The paper introduces PANP, which uses patch representations to reduce attention complexity from O(N²) to O(P²) for scalable high-resolution image processing.
• It employs dual paths – deterministic and latent – via Transformer encoders to capture both context-specific and global generative features.
• Empirical evaluations on datasets like CIFAR-10 and CelebA demonstrate improved MSE and reconstruction quality, confirming PANP’s efficacy in high-res meta-regression.

Patch Attentive Neural Processes (PANP) extend the Attentive Neural Process (ANP) architecture to high-resolution image meta-regression by introducing learnable patch representations and scalable Transformer-based attention mechanisms. PANP applies a patchwise processing approach inspired by Vision Transformer (ViT) and Masked Auto-Encoder (MAE), allowing Neural Process (NP) models to handle large images by substantially reducing attention sequence length, while maintaining flexible deterministic and latent modeling pathways (Yu et al., 2022).

1. Problem Setting and Motivation

PANP addresses meta-learning regression, in which each task involves reconstructing a sample from a distribution of two-dimensional functions $f: \mathbb{R}^2 \to \mathbb{R}$, often modeled as image regression under a random function family such as a Gaussian process with a fixed kernel. Given a context set of observed input-output pairs $(X_C, Y_C) = \{ (x_c, y_c) \}_{c \in C}$ and target inputs $X_T = \{ x_t \}_{t \in T}$, the objective is to predict target outputs $Y_T = \{ y_t \}_{t \in T}$.

Standard Neural Process models aggregate information by mean pooling, which limits expressiveness, while ANP augments this with multi-head attention over the context set to allow target-specific conditioning. However, ANP’s cross-attention step incurs $O(N^2)$ time and space complexity, where $N$ is the number of context tokens. For images, $N$ can easily exceed $10^5$ for moderate resolutions, making attention intractable at the per-pixel level. PANP addresses this limitation by partitioning the image into $P \ll N$ non-overlapping square patches and processing these as the basic tokens, making cross-attention tractable even for high-dimensional image inputs.

2. Model Architecture

The PANP architecture consists of two parallel processing paths (deterministic and latent), each based on attention over patch representations.

2.1 Patch Extraction and Encoding

An image $I \in \mathbb{R}^{H \times W \times C}$ is divided into $P = (H/p) \cdot (W/p)$ square patches $\{P_i\}_{i=1}^P$, each of size $p \times p \times C$. Each patch $P_i$ is encoded by a shared convolutional encoder $\mathrm{Conv}_e$ to produce an embedding $e_i = \mathrm{Conv}_e(P_i) \in \mathbb{R}^d$. After adding a trainable positional embedding $\mathrm{pos}_i \in \mathbb{R}^d$ to $e_i$, one obtains the initial patch token $h_i^{(0)} = e_i + \mathrm{pos}_i$.

2.2 Deterministic Path

A stack of $L$ Transformer encoder blocks processes the set of context patch tokens, with each token updated via multi-headed self-attention (MSA) and layer normalization (LN):

$$h_i^{(\ell+1)} = h_i^{(\ell)} + \mathrm{MSA}(\mathrm{LN}(h_i^{(\ell)})), \quad \ell=0,\dots,L-1$$

After $L$ layers, $r_i = h_i^{(L)}$ for $i \in C$. These are mean-pooled to form a global context summary $r_C = (1/|C|)\sum_{i\in C} r_i$.

For each target patch $t \in T$, a query token $q_t = e_t + \mathrm{pos}_t$ is constructed and used in a cross-attention operation against the set $\{ r_i \}_{i \in C }$:

$$r^*_t = \mathrm{CrossAttn}(Q = q_t, K = \{ r_i \}, V = \{ r_i \})$$

2.3 Latent Path

The latent path employs a (shared or separate) Transformer encoder stack to yield patchwise representations $s_i$ for $i \in C$. These are mean-pooled:

$s_C = \frac{1}{|C|} \sum_{i \in C} s_i$

A variational Gaussian posterior over a global latent $z \in \mathbb{R}^m$ is parameterized as:

$$q(z | s_C) = \mathcal{N}(z; \mu(s_C), \mathrm{diag}[\sigma(s_C)]^2)$$

At training time, a mean-pooled feature over $C \cup T$, $s_T$, is used to estimate the "true" posterior $q(z | s_T)$.

2.4 Decoder

A multilayer perceptron (MLP) with GeLU activations decodes each target patch:

$$\hat{P}_t = D([z;\,r^*_t;\,\mathrm{pos}_t]) \in \mathbb{R}^{p \times p \times C}$$

where $[z;\,r^*_t;\,\mathrm{pos}_t]$ denotes concatenation. The reconstructions $\{ \hat{P}_t \}$ are reassembled to form the target output.

3. Mathematical Formulation and Learning Objective

The PANP objective follows the standard NP framework with modifications for patch-based representations:

• For patch encodings, $e_i = \mathrm{Conv}_e(P_i)$, $h_i^{(0)} = e_i + \mathrm{pos}_i$.
• For each target $t$, cross-attend to context representations as above.
• The decoder predicts a likelihood per patch conditioned on the sampled latent and attended context:

$$p(P_t | z, r^*_t, \mathrm{pos}_t) = \mathcal{N}\big(P_t; \mu_D(z,r^*_t, \mathrm{pos}_t), \mathrm{diag}[\sigma_D(z,r^*_t, \mathrm{pos}_t)]^2 \big)$$

The evidence lower bound (ELBO) per task is:

$$\log p(P_T | P_C) \geq \mathbb{E}_{q(z|s_T)} \left[ \sum_{t\in T} \log p(P_t | z, r^*_t, \mathrm{pos}_t) \right] -\mathrm{KL}[q(z|s_T) \parallel q(z|s_C)]$$

This structure retains both global stochasticity via the latent $z$ and context-specific adaptation through cross-attentive deterministic paths.

4. Computational Complexity and Scalability

A key innovation in PANP is the reduction in attention bottleneck by substituting $N \approx \#$pixels with $$P = \left(\frac{H}{p}\right)\left(\frac{W}{p}\right)$$. With typical patch width $p$ in the range 8–16, $P$ is smaller by multiple orders of magnitude than the pixel count, reducing attention cost from $O(N^2 d)$ in ANP to $O(P^2 d)$ in PANP. Mean aggregation, as in NPs, remains $O(Nd)$, but attention is tractable for large images only when working with patch-level tokens—ANP is not feasible for images beyond $\sim 64 \times 64$ pixels, while PANP can scale to $128 \times 128$ and $256 \times 256$ with only a modest resource increase.

5. Experimental Evaluation

PANP’s effectiveness is demonstrated by experiments on synthetic Gaussian process images and real-world datasets such as CIFAR-10 and CelebA at resolutions up to $256 \times 256$. Comparisons are drawn with NP, pixel-level ANP, ViT-style reconstruction baselines, and MAE:

• On CIFAR-10 ($32 \times 32$), PANP reduces mean squared error (MSE) by approximately 10% relative to pixel-level ANP and by 20% relative to NP.
• PANP matches or slightly outperforms pixel-level ANP at low resolutions. At higher resolutions, where ANP becomes computationally prohibitive, PANP remains tractable, making it suitable for high-resolution tasks.
• Qualitative results show improved preservation of global image structure (shapes and colors) and fine textures within patches.

Performance metrics in these evaluations include test log-likelihood, MSE on held-out patches, and image reconstruction metrics such as PSNR and SSIM.

6. Strengths, Limitations, and Future Directions

PANP overcomes the core scalability bottleneck of ANP by using patch-wise attention, enabling attention-based Neural Processes for high-resolution images and other large structured data. The architecture’s two-path design maintains both global generative ability (via the latent $z$) and target-specific adaptation (via cross-attention). However, some loss of detail may occur at patch boundaries unless patch size $p$ is chosen sufficiently small, and while $P \ll N$, attention cost is still $O(P^2)$. This suggests a tradeoff: very fine patches recover per-pixel detail but increase costs.

Potential extensions include:

• Hierarchical patch attention, spanning multiple patch scales.
• Sparse or localized attention mechanisms to further improve scalability.
• Generalization to video or 3D domains via spatio-temporal or volumetric patches.
• Integrating convolutional inductive biases directly into the Transformer blocks.

In summary, PANP generalizes Attentive Neural Processes to high-resolution, patch-wise image regression with efficient attention, combining the strengths of NP/ANP meta-learning and ViT/MAE-style patch representation to produce a tractable, flexible regressor for visual tasks (Yu et al., 2022).