Super-Resolution Neural Operator

Single image super-resolution (SR) addresses the inverse problem of reconstructing high-resolution (HR) images from their low-resolution (LR) counterparts. In a data-driven way, deep neural networks (DNNs) learn the inversion map from many LR-HR sample pairs and have demonstrated appealing performances [4, 20, 21, 40, 41, 29, 24]. Nevertheless, most DNNs are developed in the configuration of single scaling factors, which cannot be used in scenarios requiring arbitrary SR factors [37, 38]. Recently, implicit neural functions (INF) [5, 18] have been proposed to represent images in arbitrary resolution, and paving a feasible way for continuous SR. These networks, as opposed to storing discrete signals in grid-based formats, represent signals with evaluations of continuous functions at specified coordinates, where the functions are generally parameterized by a multi-layer perceptron (MLP). To share knowledge across instances instead of fitting individual functions for each signal, encoder-based methods [5, 26, 15] are proposed to retrieve latent codes for each signal, and then a decoding MLP is shared by all the instances to generate the required output, where both the coordinates and the corresponding latent codes are taken as input. However, the point-wise behavior of MLP in the spatial dimensions results in limited performance when decoding various objects, particularly for high-frequency components [32, 30].

Neural operator is a newly proposed neural network architecture in the field of computational physics [23, 19, 17] for numerically efficient solvers of partial differential equations (PDE). Stemming from the operator theory, nueral operators learn mappings between infinite-dimensional function spaces, which is inherently capable of continuous function evaluations and has shown promising potentials in various applications [27, 9, 13]. Typically, neural operator consists of three components: 1) lifting, 2) iterative kernel integral, and 3) projection. The kernel integrals operate in the spatial domain, and thus can explicitly capture the global relationship constraining the underlying solution function of the PDE. The attention mechanism in transformers [36] is a special case of kernel integral where linear transforms are first exerted to the feature maps prior to the inner product operations [17]. Tremendous successes of transformers in various tasks [6, 22, 34] have shown the importance of capturing global correlations, and this is also true for SR to improve performance. [8].

In this paper, we propose the super-resolution neural operator (SRNO), a deep operator learning framework that can resolve HR images from their LR counterparts at arbitrary scales. As shown in Fig.1, SRNO learns the mapping between the corresponding function spaces by treating the LR-HR image pairs as continuous functions approximated with different grid sizes. The key characteristics distinguishing SRNO from prior continuous SR works are: 1) the kernel integral in each layer is efficiently implemented via the Galerkin-type attention, which possesses non-local properties in the spatial dimensions and have been proved to be comparable to a Petrov-Galerkin projection [3]; and 2) the multilayer attention architecture allows for the dynamic latent basis update, which is crucial for SR problems to “hallucinate” high-frequency information from the LR image. When employing same encoders to capture features, our method outperforms previous continuous SR methods in terms of both reconstruction accuracy and running time.

In summary, our main contributions are as follows:

• •

  We propose the methodology of super-resolution neural operator that maps between finite-dimensional function spaces, allowing for continuous and zero-shot super-resolution irrespective the discretization used on the input and output spaces.

• •

  We develop an architecture for SRNO that first explores the common latent basis for the whole training set and subsequently refines an instance-specific basis by the Galerkin-type attention mechanism.

• •

  Numerically, we show that the proposed SRNO outperforms existing continuous SR methods with less running time, and even generates better results on the resolutions for which the fixed scale SR networks were trained.

Deep learning based SR methods. [4, 20, 21, 40, 41, 29, 24] have achieved impressive performances, in multi-scale scenarios one has to train and store several models for each scale factor, which is unfeasible when considering time and memory budgets. In recent years, several methods [11, 37, 31] are proposed to achieve arbitrary-scale SR with a single model, but their performances are limited when dealing with out-of-distribution scaling factors. Inspired by INF, LIIF [5] takes continuous coordinates and latent variables as inputs, and employs an MLP to achieve outstanding performances for both in-distribution and out-of-distribution factors. In contrast, LTE [18] transforms input coordinates into the Fourier domain and uses the dominant frequencies extracted from latent variables to address the spectral bias problem [32, 30]. In a nutshell, treating images as RGB-valued functions and sharing the implicit function space are the keys to the success of LIIF-like works [5, 18]. Nevertheless, a purely local decoder, like MLP, is not able to accurately approximate arbitrary images, although it is rather sensitive to the input coordinates.

Neural Operators. Recently, a novel neural network architecture, Neural Operator (NO), was proposed for discretization invariant solutions of PDEs via infinite-dimensional operator learning [19, 10, 23, 17]. Neural operators only need to be trained once and are capable of transferring solutions between differently discretized meshes while keeping a fixed approximation error. A valuable merit of NO is that it does not require knowledge of the underlying PDE, which allows us to introduce it by the following abstract form,

where $u:D\rightarrow\mathbb{R}^{d_{u}}$ is the solution function residing in the Banach space $\mathcal{U}$, and $\mathbf{L}:\mathcal{A}\rightarrow\mathbf{L}(\mathcal{U};\mathcal{U}^{*})$ is an operator-valued functional that maps the coefficient function $a\in\mathcal{A}$ of the PDE to $f\in\mathcal{U}^{*}$, the dual space of $\mathcal{U}$. As in many cases the inverse operator of $\mathbf{L}$ even does not exist, NO seeks a feasible operator $\mathcal{G}:\mathcal{A}\rightarrow\mathcal{U},a\mapsto u$, directly mapping the coefficient to the solution within an acceptable tolerance.

The operator $\mathcal{G}$ is numerically approximated by training a neural network $\mathcal{G}_{\theta}:\mathcal{A}\rightarrow\mathcal{U}$, where $\theta$ are the trainable parameters. Suppose we have $N$ pairs of observations $\{a_{j},u_{j}\}_{j=1}^{N}$ where the input functions $a_{j}$ are sampled from probability measure $\mu$ compactly supported on $\mathcal{A}$, and $u_{j}=\mathcal{G}(a_{j})$ are used as the supervisory output functions. The infinitely dimensional operator learning problem $\mathcal{G}\leftarrow\mathcal{G}_{\theta}$ thus is associated with the empirical-risk minimization problem [35]. In practice, we actually measure the approximation loss using the sampled observations $u_{o}^{(j)}$ and $a_{o}^{(j)}$, which are the direct results of discretization:

Similar to classical feedforward neural networks (FFNs), the NO is of an iterative architecture. For ease of exposition, suppose $\mathcal{A}$ is defined on the bounded domain $D\subset\mathbb{R}^{d}$, and the inputs and outputs of the intermediate layers are all vector-valued functions, with dimension $d_{z}$. Then $\mathcal{G}_{\theta}:\mathcal{A}\rightarrow\mathcal{U}$ can be formulated as follows:

where $\mathcal{L}:\mathbb{R}^{d_{a}+d}\rightarrow\mathbb{R}^{d_{z}}$, and $\mathcal{P}:\mathbb{R}^{d_{z}}\rightarrow\mathbb{R}^{d_{u}}$ are the local lifting and projection functions respectively , mapping the input $a$ to its first layer hidden representation $z_{0}$ and the last layer hidden representation $z_{T}$ back to the output function $u$. $W:\mathbb{R}^{d_{z}}\rightarrow\mathbb{R}^{d_{z}}$ is a point-wise linear transformation, and $\sigma:\mathbb{R}^{d_{z}}\rightarrow\mathbb{R}^{d_{z}}$ is the nonlinear activation function.

Although the PDE in (1) point-wisely defines the behavior of the solution function $u(x)$, the solution operator $\mathcal{G}$ we are seeking should exhibit the non-local property such that $\mathcal{G}(a)$ can approximate $u$ everywhere rather than locally. For this purpose, NO may employ the kernel integral operators $\mathcal{K}_{t}:\{z_{t}:{D}_{t}\rightarrow\mathbb{R}^{d_{z}}\}\mapsto\{{z_{t+1}:{D_{t+1}}\rightarrow\mathbb{R}^{d_{z}}}\}$ to maintain the continuum in the spatial domain, and one of the most adopted forms [19] is defined as

where the kernel matrix $K_{t}:\mathbb{R}^{d+d}\rightarrow\mathbb{R}^{d_{z}\times d_{z}}$ is parameterized by $\Phi$.

All operations of NO are defined in function space and the training data are just samples at coordinates ${x_{i}}$ irrespective of the discretization sizes. This property inspires us to employ the NO’s archetecture for continuous SR. Different from NO that aims at mapping between infinite-dimensional function spaces, we define SRNO in the sense of mappings between two approximation spaces of finite-dimensional but continuous functions. This is reasonable because imaging optics inevitably limit the highest possible frequency in the radiance field under observation. With this configuration, we can design the archetecture of SRNO using the well-developed theoretical tool, the Galerkin-type method that is widely used in the filed of Finite elements [7].

Problem Formulation. Let $(\mathcal{H},\langle\cdot,\cdot\rangle_{\mathcal{H}})$ denote a Hilbert space equipped with the inner-product structure, which is continuously embedded in the space of continuous functions $C^{0}(\Omega)$, with $\Omega\subset\mathbb{R}^{2}$ a bounded domain. An image is defined as a vector-valued function $\mathcal{H}\ni f:\Omega\rightarrow\mathbb{R}^{3}$. Assume we can access the function values of $f$ at the coordinates $\{x_{i}\}_{i=1}^{n_{h}}\subset\Omega$ with the biggest discretization size $h$. Note that $x_{i}$’s are not necessarily equidistant. Our goal is to learn a super-resolution neural operator between two Hilbert spaces with different resolutions: $\mathcal{S}_{\theta}:\mathcal{H}\supset\mathcal{H}(\Omega_{h_{c}})\rightarrow\mathcal{H}(\Omega_{h_{f}})\subset\mathcal{H}$, where $h_{c},h_{f}$ denotes the coarse and the fine grid sizes, respectively. Given $N$ function pairs $\{a^{(k)},u^{(k)}\}_{k=1}^{N}$, where $a^{(k)}\in\mathcal{H}(\Omega_{h_{c}})$ and $u^{(k)}\in\mathcal{H}(\Omega_{h_{f}})$. Our SRNO parameterized by $\theta$ can be solved through the associated empirical-risk minimization problem:

Since our data $a^{(k)}$ and $u^{(k)}$ are functions, to work with them numerically, we assume access only to their point-wise evaluations. Let $\mathbb{A}_{h_{c}}\subset\mathcal{H}(\Omega_{h_{c}})\subset\mathcal{H}$ be an approximation space [3] associated with $\{x_{i}\}^{n_{h_{c}}}_{i=1}$, such that for any $a\in\mathbb{A}_{h_{c}}$, $a(\cdot)=\sum_{i=1}^{n_{h_{c}}}{a(x_{i})\xi_{x_{i}}(\cdot)}$, where $\{\xi_{x_{i}}(\cdot)\}$ form a set of nodal basis for $\mathbb{A}_{h_{c}}$ in the sense that $\xi_{x_{i}}(x_{j})=\delta_{ij}$. Similarly, $\mathbb{U}_{h_{f}}\subset\mathcal{H}(\Omega_{h_{f}})\subset\mathcal{H}$ is another approximation space associate with grid size $h_{f}$. Note that Our observed LR-HR image pairs $\{a,u\}$ may have different grid size pairs $\{h_{c},h_{f}\}$, which means they belongs to different approximation spaces $\{\mathbb{A}_{h_{c}},\mathbb{U}_{h_{f}}\}$.

Lifting. The neural operators from PDEs usually use a simple pointwise function, $\mathcal{L}$ in (3), to expand the input channels. But for the SR problem, a deep feature encoder $E_{\psi}$ (with trainable parameters $\psi$) and spatial interpolations should be considered in the lifting operation, due to the complexity of natural images and the discretization inconsistency between LR and HR image functions. Another critical ingredient in Lifting is the incorporation scheme of coordinates. Prior study [14] demonstrates that deep CNNs can implicitly learn to encode the information about absolute positions. As shown in Fig.2 (c), the position information of the grid points $\{a_{l}\}_{l=1}^{4}$ have been implicitly encoded by $E_{\psi}$ employing CNNs. Therefore, we only need to explicitly construct the coordinate features $\hat{a}(x)$ with the fractional part of coordinate $x$ inside a grid. In order to reduce the blocky artifacts resulting from direct interpolation of the LR feature maps, we propose to concatenate the features weighted by the corresponding bilinear interpolation factors $s_{l}$. Compared to the local ensemble trick in [5], our method takes shorter running time and overcomes the over-smoothing problem. We reformulate the lifting operation in (3) as:

where $x$ are coordinates of HR image functions, $\hat{x}_{l}$ are the coordinates of the four neighbors of $x$, ${\delta}_{l}(x)=x-\hat{x}_{l}$, $\mathbf{c}=(2/r_{x},2/r_{y})$ represents a $r_{x}\times r_{y}$ local area in HR images with $r_{x}$ and $r_{y}$ the scaling factors, and $\mathcal{L}:\mathbb{R}^{4\times(d_{e}+2)+2}\rightarrow\mathbb{R}^{d_{z}}$ is an local linear transformation function, with $d_{e}$ the number of channels after $E_{\psi}$.

Kernel integral. The kernel integral operator, in (4), actually tells that we can identify the hidden representation $z(x)$ of the underlying image function $f(x)$ with distributions [28], inner products in the settings of SRNO. We first define the kernel $K_{t}:\mathbb{R}^{d_{z}+d_{z}}\rightarrow\mathbb{R}^{d_{z}\times d_{z}}$ with respect to the input pair $(z(x),z(y))$, rather than like (6) depending on the spatial variables $(x,y)$. Furthermore, in order to more efficiently explore $z(x)$ we can use multiple sets of test functions defining the distributions, which may immediately remind one of a single-head self-attention [17]. Denote $z_{i}=z(x_{i})\in\mathbb{R}^{d_{z}}$ for $i=1,\dots,n_{h_{f}}$. The kernel integral operator can be approximated by the Monte-Carlo method (omitting the layer index):

where

In the language of transformers, the matrices $W_{q},W_{k},W_{v}\in\mathbb{R}^{d_{z}\times d_{z}}$ correspond to the queries, keys and values functions respectively. By allowing $K$ to be the sum of multiple functions with separate trainable parameters, the multi-head self-attention can also be rewritten like in (9). For every observation $(a,u)$, this operation has a complexity of $O(n_{h_{f}}^{2}d_{z})$, which is unaffordable for the SR problem where the sampling number $n_{h_{f}}$ usually reaches to $10,000$ or beyond. To overcome this issue, we employ the Galerkin-type attention operator which has a linear complexity $O(n_{h_{f}}d_{z}^{2})$. The approximation capacity of a linearized Galerkin-type attention has been proved to be comparable to a Petrov-Galerkin projection [3].

We suppose that a image function $f$ is locally integrable, i.e., $f$ is measurable and $\int_{\Omega}{|f|}<\infty$ for every compact $\Omega\subset\mathbb{R}^{2}$. The idea behind distributions [28] is to identify $f$ with $\int{f\phi}$ by suitablely choosing a “test function” $\phi$, rather than by a series of separate evaluations $f(x)$ at coordinates $x\in\Omega$. Let $v(x)=W_{v}z(x),k(x)=W_{k}z(x)$, and $q(x)=W_{q}z(x)$ be the keys, queries and values functions respectively, which are $d_{z}$ dimensional vector-valued functions, with a subscript denoting the corresponding component. The kernel integral attention then can be written in a component-wise form ($j=1,\dots,d_{z}$):

Denote the evaluations $(z(x_{1}),\dots,z(x_{n_{h_{f}}}))^{T}$ by $\mathbf{z}\in\mathbb{R}^{n_{h_{f}}\times d_{z}}$. Let $Q=\mathbf{z}W_{q},K=\mathbf{z}W_{k},V=\mathbf{z}W_{v}\in\mathbb{R}^{n_{h_{f}}\times d_{z}}$. The columns of $Q/K/V$ contain the vector representations of the learned basis functions, spanning certain subspaces of the latent representation Hilbert spaces respectively.

where $\widetilde{K}=\text{Ln}(K),\widetilde{V}=\text{Ln}(K)$, with $\text{Ln}(\cdot)$ the layer normalization. The $j$-th column of $\widetilde{K}^{T}\widetilde{V}$ contains the coefficients for the linear combination of the basis vectors $\{q_{j}\}_{j=1}^{d_{z}}$ to form the output $\mathbf{z}$. Consequently, the global correlations are reflected through the components of $\widetilde{K}^{T}\widetilde{V}$. As a result, SRNO can utilize all information effectively without the feature unfolding [5] or local convolutions [18], significantly suppressing the discontinuous patterns that appear around feature boundaries.

In addition to providing a global aggregation for the output z at each sampling point, the linear Galerkin-type attention (12) has the ability to obtain a quasi-optimal approximation in the current approximation space spanned with the columns of $Q$ [3]. However, the expressive capability of the current bases in $Q$, $K$, and $V$ solely depends on the current input LR image through the latent representation z, is there any chance that we can enrich the bases which contains some extra and useful information for SR reconstruction, but not included in the input z? For this purpose, the point-wise FFN $\mathcal{O}:\mathbb{R}^{d_{z}}\rightarrow\mathbb{R}^{d_{z}}$ in Fig.2 (b) introduces nonlinearties on one hand, and the positions concatenated in z enhance the bases on the other. In this way, the basis functions are being constantly enriched, and we reformulate the iterative process (See supplementary for a principled discussion):

Network details. The Network architecture is shown in Fig.2. As to the feature encoder $E_{\psi}$, we employ EDSR-baseline [21], or RDN [41], both of which drop their upsampling layers, and their output channel dimensions $d_{e}=64$. The CNNs in $E_{\psi}$ assist SRNO in capturing the common basis functions from the ensemble of training samples, while the Galerkin-type attention layers provide the instance-specific basis enhancement. We employ the multi-head attention scheme in [36] by dividing the queries, keys and values into $n_{heads}$ parts with each of dimension $d_{z}/n_{heads}$. In our implementation, $d_{z}=256,n_{heads}=16$, yielding 16-dimensional output values. We only use two iterations $(T=2)$ in the kernel integral operator, which already outperforms previous works, while keeping the running time advantage. Note that we utilize $1\times 1$ convolutions to replace all the linear layers in SRNO, since they have a GPU-friendly data structure. The detailed network structures are listed in the supplementary.

In this paper, we proposed the Super-Resolution Neural Operator (SRNO) for continuous super-resolution. In SRNO, each image is seen as a function, and our method learns a map between finite-dimensional function spaces, which allows SRNO can be trained and generalize on different levels of discretization. First, in the Lifting, we use CNN-based encoders to capture feature maps from LR images and design a simple but efficient interpolation-free method, addressing the discretization inconsistency problem encountered in SR. Second, to approximate our target function in the Iterative Kernel Integration, we employ a linear attention operator that has been proved to be comparable to the Petrov-Galerkin projection. Finally, we map the last hidden representation to the output function. Experimental results show that our SRNO outperforms other arbitrary-scale SR methods in performance as well as computation time, and particularly in the capability of capturing the global image structures.