SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Doyup Lee*
Kakao Brain
doyup.lee@kakaobrain.com
Chiheon Kim*
chiheon.kim@kakaobrain.com
Minsu Cho†
mscho@postech.ac.kr
Wook-Shin Han
wshan@dblab.postech.ac.kr
Generalizable implicit neural representation (INR) enables a single continuous function, i.e., a coordinate-based neural network, to represent multiple data instances by modulating its weights or intermediate features using latent codes. However, the expressive power of the state-of-the-art modulation is limited due to its inability to localize and capture fine-grained details of data entities such as specific pixels and rays. To address this issue, we propose a novel framework for generalizable INR that combines a transformer encoder with a locality-aware INR decoder. The transformer encoder predicts a set of latent tokens from a data instance to encode local information into each latent token. The locality-aware INR decoder extracts a modulation vector by selectively aggregating the latent tokens via cross-attention for a coordinate input and then predicts the output by progressively decoding with coarse-to-fine modulation through multiple frequency bandwidths. The selective token aggregation and the multi-band feature modulation enable us to learn locality-aware representation in spatial and spectral aspects, respectively. Our framework significantly outperforms previous generalizable INRs and validates the usefulness of the locality-aware latents for downstream tasks such as image generation.
Recent advances in generalizable implicit neural representation (INR) enable a single coordinate-based multi-layer perceptron (MLP) to represent multiple data instances as a continuous function. Instead of per-sample training of individual coordinate-based MLPs, generalizable INR extracts latent codes of data instances [13, 14, 40] to modulate the weights or intermediate features of the shared MLP model [8, 11, 19, 35]. However, despite the advances in previous approaches, their performance is still insufficient compared with individual training of INRs per sample.
We postulate that the expressive power of generalizable INRs is limited by the ability of locality-awareness to localize relevant entities from a data instance and control
Figure 1: Learning curves of PSNRs during training on ImageNette  $178 \times 178$ .
their structure in a fine-grained manner. Primitive entities of a data instance, such as pixels in an image, tend to have a higher correlation with each other if they are closer in space and time. Thus, this locality of data entities has been used as an important inductive bias for learning the representations of complex data [3]. However, previous approaches to generalizable INRs are not properly designed to leverage the locality of data entities. For example, when latent codes modulate intermediate features [11, 12] or weight matrices [8, 19, 35] of an INR decoder, the modulation methods do not exploit a specified coordinates for decoding, which restricts the latent codes to encoding global information over all pixels without capturing local relationships between specific pixels.
To address this issue, we propose a novel encoder-decoder framework for locality-aware generalizable INR to effectively localize and control the fine-grained details of data. In our framework, a Transformer encoder [37] first extracts locally relevant information from a data instance and predicts a set of latent tokens to encode different local information. Then, our locality-aware INR decoder effectively leverages the latent tokens to predict fine-grained details. Specifically, given an input coordinate, our INR decoder uses a cross-attention to selectively aggregate the local information in the latent tokens and extract a modulation vector for the coordinate. In addition, our INR decoder effectively captures the high-frequency details in the modulation vector by decomposing it into multiple bandwidths of frequency features and then progressively composing the intermediate features. We conduct extensive experiments to demonstrate the high performance and efficacy of our locality-aware generalizable INR on benchmarks as shown in Figure 1. In addition, we show the potential of our locality-aware INR latents to be utilized for downstream tasks such as image synthesis.
Our main contributions can be summarized as follows: 1) We propose an effective framework for generalizable INR with a Transformer encoder and locality-aware INR decoder. 2) The proposed INR decoder with selective token aggregation and multi-band feature modulation can effectively capture the local information to predict the fine-grained data details. 3) The extensive experiments validate the efficacy of our framework and show its applications to a downstream image generation task.
Implicit neural representations (INRs). INRs use neural networks to represent complex data such as audio, images, and 3D scenes, as continuous functions. Especially, incorporating Fourier features [24, 36], periodic activations [31], or multi-grid features [25] significantly improves the performance of INRs. Despite its broad applications [1, 6, 10, 32, 34], INRs commonly require separate training of MLPs to represent each data instance. Thus, individual training of INRs per sample does not learn common representations in multiple data instances.
Generalizable INRs. Previous approaches focus on two major components for generalizable INRs; latent feature extraction and modulation methods. Auto-decoding [23, 26] computes a latent vector per data instance and concatenates it with the input of a coordinate-based MLP. Given input data, gradient-based meta-learning [4, 11, 12] adapts a shared latent vector using a few update steps to scale and shift the intermediate activations of the MLP. Learned Init [35] also uses gradient-based meta-learning but adapts whole weights of the shared MLP. Although auto-decoding and gradient-based meta-learning are agnostic to the types of data, their training is unstable on complex and large-scale datasets. TransINR [8] employs the Transformer [37] as a hypernetwork to predict latent vectors to modulate the weights of the shared MLP. In addition, Instance Pattern Composers [19] have demonstrated that modulating the weights of the second MLP layer is enough to achieve high performance of generalizable INRs. Our framework also employs the Transformer encoder, but focuses on extracting locality-aware latent features for the high performance of generalizable INR.
Leveraging Locality of Data for INRs Local information in data has been utilized for efficient modeling of INRs, since local relationships between data entities are widely used for effective process of complex data [3]. Given an input coordinate, the coordinate-based MLP only uses latent vectors nearby the coordinate, after a CNN encoder extracts a 2D grid feature map of an image for super-resolution [7] and reconstruction [22]. Spatial Functa [4] demonstrates that leveraging the locality of data enables INRs to be utilized for downstream tasks such as image recognition and generation. Local information in 3D coordinates has also been effective for scene modeling as a hybrid approach using 3D feature grids [18] or the part segmentation [17] of a 3D object. However, previous approaches assume explicit grid structures of latents tailored to a specific data type. Since
Figure 2: Overview of our framework for locality-aware generalizable INR. Given a data instance, Transformer encoder extracts its localized latents. Then, the locality-aware INR decoder uses selective token aggregation and multi-band feature modulation to predict the output for the input coordinate.
we do not predefine a relationship between latents, our framework is flexible to learn and encode the local information of both grid coordinates in images and non-grid coordinates in light fields.
We propose a novel framework for locality-aware generalizable INR which consists of a Transformer encoder to localize the information in data into latent tokens and a locality-aware INR decoder to exploit the localized latents and predict outputs. First, we formulate how generalizable INR enables a single coordinate-based neural network to represent multiple data instances as a continuous function by modulating its weights or features. Then, after we introduce the Transformer encoder to extract a set of latent tokens from input data instances, we explain the details of the locality-aware INR decoder, where selective token selection aggregates the spatially local information for an input coordinate via cross-attention; multi-band feature modulation leverages a different range of frequency bandwidths to progressively decode the local information using coarse-to-fine modulation in the spectral domain.
3.1 Generalizable Implicit Neural Representation
Given a set of data instances  $\mathcal{X} = \{\mathbf{x}^{(n)}\}_{n=1}^{N}$ , each data instance  $\mathbf{x}^{(n)} = \{(\mathbf{v}_i^{(n)}, \mathbf{y}_i^{(n)})\}_{i=1}^{M_n}$  comprises  $M_n$  pairs of an input coordinate  $\mathbf{v}_i^{(n)} \in \mathbb{R}^{d_{\mathrm{in}}}$  and the corresponding output feature  $\mathbf{y}_i^{(n)} \in \mathbb{R}^{d_{\mathrm{out}}}$ . Conventional approaches [24, 31, 36] adopt individual coordinate-based MLPs to train and memorize each data instance  $\mathbf{x}^{(n)}$ . Thus, the coordinate-based MLP cannot be reused and generalized to represent other data instances, requiring per-sample optimization of MLPs for unseen data instances.
A generalizable INR uses a single coordinate-based MLP as a shared INR decoder  $F_{\theta}:\mathbb{R}^{d_{\mathrm{in}}}\to \mathbb{R}^{d_{\mathrm{out}}}$  to represent multiple data instances as a continuous function. Generalizable INR [8, 11, 12, 19, 26] extracts the  $R$  number of latent codes  $\mathbf{Z}^{(n)} = \{\mathbf{z}_k^{(n)}\in \mathbb{R}^d\}_{k = 1}^R$  from a data instance  $\mathbf{x}^{(n)}$ . Then, the latents are used for the INR decoder to represent a data instance  $\mathbf{x}^{(n)}$  as  $\mathbf{y}_i^{(n)} = F_\theta (\mathbf{v}_i^{(n)};\mathbf{Z}^{(n)})$  while updating the parameters  $\theta$  and latents  $\mathbf{Z}^{(n)}$  to minimize the errors over  $\mathcal{X}$ :
$$
\min  _ {\theta , \mathbf {Z} ^ {(n)}} \frac {1}{N M _ {n}} \sum_ {n = 1} ^ {N} \sum_ {i = 1} ^ {M _ {n}} \left\| \mathbf {y} _ {i} ^ {(n)} - F _ {\theta} \left(\mathbf {v} _ {i} ^ {(n)}; \mathbf {Z} ^ {(n)}\right) \right\| _ {2} ^ {2}. \tag {1}
$$
We remark that each previous approach employs a different number of latent codes to modulate a coordinate-based MLP. For example, a single latent vector  $(R = 1)$  is commonly extracted to modulate intermediate features of the MLP [11, 12, 26], while a multitude of latents  $(R > 1)$  are used to modulate its weights [8, 19, 35]. While we modulate the features of MLP, we extract a set of latent codes to localize the information of data to leverage the locality-awareness for latent features.
3.2 Transformer Encoder
Our framework employs a Transformer encoder [37] to extract a set of latents  $\mathbf{Z}^{(n)}$  for each data instance  $\mathbf{x}^{(n)}$  as shown in Figure 2. After a data instance, such as an image or multi-view images, is patched into a sequence of data tokens, we concatenate the patchified tokens into a sequence of  $R$  learnable tokens as the encoder input. Then, the Transformer encoder extracts a set of latent tokens, where each latent token corresponds to an input learnable token. Note that the permutation-equivariance of self-attention in the Transformer encoder enables us not to predefine the local structure of data and the ordering of latent tokens. During training, each latent token learns to capture the local information of data, while covering whole regions to represent a data instance. Thus, whether a data instance is represented on a grid or non-grid coordinate, our framework is flexible to encode various types of data into latent tokens, while learning the local relationships of latent tokens during training.
3.3 Locality-Aware Decoder for Implicit Neural Representations
We propose the locality-aware INR decoder in Figure 2 to leverage the local information of data for effective generalizable INR. Our INR decoder comprises two primary components: i) Selective token aggregation via cross attention extracts a modulation vector for an input coordinate to aggregate spatially local information from latent tokens. ii) Multi-band feature modulation decomposes the modulation vector into multiple bandwidths of frequency features to amplify the high-frequency features and effectively predict the details of outputs.
3.3.1 Selective Token Aggregation via Cross-Attention
We remark that encoding locality-aware latent tokens is not straightforward since the self-attention in Transformer do not guarantee a specific relationship between tokens. Thus, the properties of the latent tokens are determined by a modulation method for generalizable INR to exploit the extracted latents. For example, given an input coordinate  $\mathbf{v}$  and latent tokens  $\{\mathbf{z}_1,\dots,\mathbf{z}_R\}$ , a straightforward method can use Instance Pattern Composers [19] to construct a modulation weight  $\mathbf{W}_{\mathrm{m}} = [\mathbf{z}_1,\dots,\mathbf{z}_R]^{\top} \in \mathbb{R}^{R \times d_{\mathrm{in}}}$  and extract a modulation vector  $\mathbf{m}_{\mathbf{v}} = \mathbf{W}_{\mathrm{m}}\mathbf{v} = [\mathbf{z}_1^{\top}\mathbf{v},\dots,\mathbf{z}_R^{\top}\mathbf{v}]^{\top} \in \mathbb{R}^{R}$ . However, the latent tokens cannot encode the local information of data, since each latent token equally influences each channel of the modulation vector regardless of the coordinate locations (see Section 4.3).
Our selective token aggregation employs cross-attention to aggregate the spatially local latents nearby the input coordinate, while guiding the latents to be locality-aware. Given a set of latent tokens  $\mathbf{Z}^{(n)} = \{\mathbf{z}_k^{(n)}\}_{k=1}^R$  and a coordinate  $\mathbf{v}_i^{(n)}$ , a modulation feature vector  $\mathbf{m}_{\mathbf{v}_i}^{(n)} \in \mathbb{R}^d$  shifts the intermediate features of an INR decoder to predict the output, where  $d$  is the dimensionality of hidden layers in the INR decoder. For the brevity of notation, we omit the superscript  $n$  and subscript  $i$ .
Frequency features We first transform an input coordinate  $\mathbf{v} = (v_{1},\dots ,v_{d_{\mathrm{in}}})\in \mathbb{R}^{d_{\mathrm{in}}}$  into frequency features using sinusoidal positional encoding [31, 36]. We define the Fourier features  $\gamma_{\sigma}(\mathbf{v})\in \mathbb{R}^{d_{\mathrm{F}}}$  with bandwidth  $\sigma >1$  and feature dimensionality  $d_{\mathrm{F}}$  as
$$
\gamma_ {\sigma} (\mathbf {v}) = \left[ \cos \left(\pi \omega_ {j} v _ {i}\right), \sin \left(\pi \omega_ {j} v _ {i}\right): i = 1, \dots , d _ {\text {i n}}, j = 0, \dots , n - 1 \right] \tag {2}
$$
where  $n = \frac{d_{\mathrm{F}}}{2d_{\mathrm{in}}}$ . A frequency  $\omega_{j} = \sigma^{j / (n - 1)}$  is evenly distributed between 1 and  $\sigma$  on a log-scale. Based on the Fourier features, we define the frequency feature extraction  $\mathbf{h}_{\mathrm{F}}(\cdot)$  as
$$
\mathbf {h} _ {\mathrm {F}} (\mathbf {v}; \sigma , \mathbf {W}, \mathbf {b}) = \operatorname {R e L U} \left(\mathbf {W} \gamma_ {\sigma} (\mathbf {v}) + \mathbf {b}\right), \tag {3}
$$
where  $\mathbf{W} \in \mathbb{R}^{d \times d_{\mathrm{F}}}$  and  $\mathbf{b} \in \mathbb{R}^d$  are trainable parameters for frequency features,  $d$  denotes the dimensionality of hidden layers in the INR decoder.
Selective token selection via cross-attention To predict corresponding output  $\mathbf{y}$  to the coordinate  $\mathbf{v}$ , we adopt a cross-attention to extract a modulation feature vector  $\mathbf{m}_{\mathbf{v}} \in \mathbb{R}^{d}$  based on the latent tokens  $\mathbf{Z} = \{\mathbf{z}_k\}_{k=1}^R$ . We first extract the frequency features of the coordinate  $\mathbf{v}$  in Eq (3) as the query of the cross-attention as
$$
\mathbf {q} _ {\mathbf {v}} := \mathbf {h} _ {\mathrm {F}} (\mathbf {v}; \sigma_ {\mathbf {q}}, \mathbf {W} _ {\mathbf {q}}, \mathbf {b} _ {\mathbf {q}}), \tag {4}
$$
where  $\mathbf{W}_{\mathrm{q}}\in \mathbb{R}^{d\times d_{\mathrm{F}}}$  and  $\mathbf{b}_{\mathrm{q}}\in \mathbb{R}^{d}$  are trainable parameters, and  $\sigma_{\mathrm{q}}$  is the bandwidth for query frequency features. The cross-attention in Figure 2 enables the query to select latent tokens, aggregate
its local information, and extract the modulation feature vector  $\mathbf{m}_{\mathbf{v}}$  for the input coordinate:
$$
\mathbf {m} _ {\mathbf {v}} := \text {M u l t i H e a d A t t e n t i o n} (\text {Q u e r y} = \mathbf {q} _ {\mathbf {v}}, \text {K e y} = \mathbf {Z}, \text {V a l u e} = \mathbf {Z}). \tag {5}
$$
An intuitive implementation for selective token aggregation can employ hard attention to select only one latent token for each coordinate. However, in our primitive experiment, using hard attention leads to unstable training and a latent collapse problem that selects only few latent tokens. Meanwhile, multi-head attentions encourage each latent token to easily learn the locality in data instances.
3.3.2 Multi-Band Feature Modulation in the Spectral Domain
After the selective token aggregation extracts a modulation vector  $\mathbf{m}_{\mathbf{v}}$ , we use multi-band feature modulation to effectively predict the details of outputs. Although Fourier features [24, 36] reduce the spectral bias [2, 28] of neural networks, adopting a simple stack of MLPs to INRs still suffers from capturing the high-frequency data details. To address this issue, we use a different range of frequency bandwidths to decompose the modulation vector into multiple frequency features in the spectral domain. Then, our multi-band feature modulation uses the multiple frequency features to progressively decode the intermediate features, while encouraging a deeper MLP path to learn higher frequency features. Note that the coarse-to-fine approach in the spectral domain is analogous to the locally hierarchical approach in the spatial domain [21, 29, 39] to capture the data details.
Extracting multiple modulation features with different frequency bandwidths We extract  $L$  level of modulation features  $\mathbf{m}_{\mathbf{v}}^{(1)},\dots ,\mathbf{m}_{\mathbf{v}}^{(L)}$  from  $\mathbf{m}_{\mathbf{v}}$  using different bandwidths of frequency features. Given  $L$  frequency bandwidths as  $\sigma_{1}\geq \sigma_{2}\geq \dots \geq \sigma_{L}\geq \sigma_{\mathrm{q}}$  , we use Eq (3) to extract the  $\ell$  -th level of frequency features of an input coordinate  $\mathbf{v}$  as
$$
\left(\mathbf {h} _ {\mathrm {F}}\right) _ {\mathbf {v}} ^ {(\ell)} := \mathbf {h} _ {\mathrm {F}} (\mathbf {v}; \sigma_ {\ell}, \mathbf {W} _ {\mathrm {F}} ^ {(\ell)}, \mathbf {b} _ {\mathrm {F}} ^ {(\ell)}) = \operatorname {R e L U} \left(\mathbf {W} _ {\mathrm {F}} ^ {(\ell)} \boldsymbol {\gamma} _ {\sigma_ {\ell}} (\mathbf {v}) + \mathbf {b} _ {\mathrm {F}} ^ {(\ell)}\right), \tag {6}
$$
where  $\mathbf{W}_{\mathrm{F}}^{(\ell)}$  and  $\mathbf{b}_{\mathrm{F}}^{(\ell)}$  are trainable parameters and shared across data instances. Then, the  $\ell$ -th modulation vector  $\mathbf{m}_{\mathbf{v}}^{(\ell)}$  is extracted from the modulation vector  $\mathbf{m}_{\mathbf{v}}$  as
$$
\mathbf {m} _ {\mathbf {v}} ^ {(\ell)} := \operatorname {R e L U} \left(\left(\mathbf {h} _ {\mathrm {F}}\right) _ {\mathbf {v}} ^ {(\ell)} + \mathbf {W} _ {\mathrm {m}} ^ {(\ell)} \mathbf {m} _ {\mathbf {v}} + \mathbf {b} _ {\mathrm {m}} ^ {(\ell)}\right), \tag {7}
$$
with a trainable weight  $\mathbf{W}_{\mathrm{m}}^{(\ell)}$  and bias  $\mathbf{b}_{\mathrm{m}}^{(\ell)}$ . Considering that ReLU cutoffs the values below zero, we assume that  $\mathbf{m}_{\mathbf{v}}^{(\ell)}$  filters out the information of  $\mathbf{m}_{\mathbf{v}}$  based on the  $\ell$ -th frequency patterns of  $(\mathbf{h}_{\mathrm{F}})_{\mathbf{v}}^{(\ell)}$ .
Multi-band feature modulation After decomposing a modulation vector into multiple features with different frequency bandwidths, we progressively compose the  $L$  modulation features by applying a stack of nonlinear operations with a linear layer and ReLU activation. Starting with  $\mathbf{h}_{\mathbf{v}}^{(1)} = \mathbf{m}_{\mathbf{v}}^{(1)}$ , we compute the  $\ell$ -th hidden features  $\mathbf{h}_{\mathbf{v}}^{(\ell)}$  for  $\ell = 2,\dots ,L$  as
$$
\widetilde {\mathbf {h}} _ {\mathbf {v}} ^ {(\ell)} := \mathbf {m} _ {\mathbf {v}} ^ {(\ell)} + \mathbf {h} _ {\mathbf {v}} ^ {(\ell - 1)} \quad \text {a n d} \quad \mathbf {h} _ {\mathbf {v}} ^ {(\ell)} := \operatorname {R e L U} \left(\mathbf {W} ^ {(\ell)} \widetilde {\mathbf {h}} _ {\mathbf {v}} ^ {(\ell)} + \mathbf {b} ^ {(\ell)}\right), \tag {8}
$$
where  $\mathbf{W}^{(\ell)}\in \mathbb{R}^{d\times d}$  and  $\mathbf{b}^{(\ell)}\in \mathbb{R}^d$  are trainable weights and biases of the INR decoder.  $\widetilde{\mathbf{h}}_{\mathbf{v}}^{(\ell)}$  denotes the  $\ell$ -th pre-activation of INR decoder for coordinate  $\mathbf{v}$ . Note that the modulation features with high-frequency bandwidth can be processed by more nonlinear operations than the features with lower frequency bandwidths, considering that high-frequency features contain more complex signals.
Finally, the output  $\hat{y}$  is predicted using all intermediate hidden features of the INR decoder as
$$
\hat {\mathbf {y}} := \sum_ {\ell = 1} ^ {L} f _ {\text {o u t}} ^ {(\ell)} \left(\mathbf {h} _ {\mathbf {v}} ^ {(\ell)}\right), \tag {9}
$$
where  $f_{\mathrm{out}}^{(\ell)}:\mathbb{R}^d\to \mathbb{R}^{d_{\mathrm{out}}}$  are a linear projection into the output space. Although utilizing only  $\mathbf{h}_{\mathbf{v}}^{(L)}$  is also an option to predict outputs, skip connections of all intermediate features into the output layer enhances the robustness of training to the hyperparameter choices.
Figure 3: Reconstructed images of FFHQ with  $512 \times 512$  resolution by TransINR [8] (left), IPC [19] (middle), and our locality-aware generalizable INR (right).
We conduct extensive experiments to demonstrate the effectiveness of our locality-aware generalizable INR on image reconstruction and novel view synthesis. In addition, we conduct in-depth analysis to validate the efficacy of our selective token aggregation and multi-band feature modulation to localize the information of data to capture fined-grained details. We also show that our locality-aware latents can be utilized for image generation by training a generative model on the extracted latents. Our implementation and experimental settings are based on the official codes of Instance Pattern Composers [19] for a fair comparison. We attach the implementation details to Appendix A.
4.1 Image Reconstruction
We follow the protocols in previous studies [8, 19, 35] to evaluate our framework on image reconstruction of CelebA, FFHQ, and ImageNette with  $178 \times 178$  resolution. Our framework also outperforms previous approaches on high-resolution images with  $256 \times 256$ ,  $512 \times 512$ , and  $1024 \times 1024$  resolutions of FFHQ. We compare our framework with Learned Init [35], TransINR [8], and IPC [19]. The Transformer encoder predicts  $R = 256$  latent tokens, while the INR decoder uses  $d_{\mathrm{in}} = 2$ ,  $d_{\mathrm{out}} = 3$ ,  $d = 256$  dimensionality of hidden features,  $L = 2$ ,  $\sigma_{\mathrm{q}} = 16$  and  $(\sigma_1, \sigma_2) = (128, 32)$  bandwidths.
$178 \times 178$  Image Reconstruction Table 1 shows that our generalizable INR significantly outperforms previous methods by a large margin. We remark that TransINR, IPC, and our framework use the same capacity of the Transformer encoder, latent tokens, and INR decoder except for the modulation methods. Thus, the results imply that our locality-aware INR decoder with selective token aggregation and multi-band feature modulation is effective to
Table 1: PSNRs of reconstructed images of  ${178} \times  {178}$  CelebA,FFHQ,and ImageNette.
CelebA	FFHQ	ImageNette
Learned Init [35]	30.37	-	27.07
TransINR	33.33	33.66	29.77
IPC	35.93	37.18	38.46
Ours	50.74	43.32	46.10
capture local information of data and fine-grained details for high-quality image reconstruction.