SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jihoon Tack<sup>1</sup>, Subin Kim<sup>1</sup>, Sihyun Yu<sup>1</sup>, Jaeho Lee<sup>2</sup>, Jinwoo Shin<sup>1</sup>, Jonathan Schwarz<sup>3</sup>
$^{1}$ Korea Advanced Institute of Science and Technology
$^{2}$ Pohang University of Science and Technology
<sup>3</sup>University College London
{jihoontack, subin-kim, sihyun.yu, jinwoos}@kaist.ac.kr
jaeho.lee@postech.ac.kr, schwarzjn@gmail.com
We introduce an efficient optimization-based meta-learning technique for large-scale neural field training by realizing significant memory savings through automated online context point selection. This is achieved by focusing each learning step on the subset of data with the highest expected immediate improvement in model quality, resulting in the almost instantaneous modeling of global structure and subsequent refinement of high-frequency details. We further improve the quality of our meta-learned initialization by introducing a bootstrap correction resulting in the minimization of any error introduced by reduced context sets while simultaneously mitigating the well-known myopia of optimization-based meta-learning. Finally, we show how gradient re-scaling at meta-test time allows the learning of extremely high-quality neural fields in significantly shortened optimization procedures. Our framework is model-agnostic, intuitive, straightforward to implement, and shows significant reconstruction improvements for a wide range of signals. We provide an extensive empirical evaluation on nine datasets across multiple multiple modalities, demonstrating state-of-the-art results while providing additional insight through careful analysis of the algorithmic components constituting our method. Code is available at https://github.com/jihoontack/GradNCP
Neural fields (NFs), also known as implicit neural representations (INRs), have emerged as a new paradigm for representing complex signals as continuous functions (e.g., the implicit pixel coordinate to feature mapping  $(x,y)\mapsto (r,g,b)$  of a 2D-image) parameterized by neural networks [43, 60]. In contrast to multi-dimensional array representations, this approach has the benefits of allowing querying at arbitrary resolutions or from arbitrary viewpoints [39, 52], removing modality-agnostic architectural choices for downstream tasks (e.g., compression [10, 48, 49], classification [11, 2] or generative modeling [55, 68]), representing a more compact parameterization and providing an elegant framework for dealing with data unsuitable for discretization (e.g., vector fields). Thus, the approach has shown great promise in modeling diverse forms of data, including images [38], videos [5], 3D scenes [43], voxels [11], climate data [25], and audio [12].
Unfortunately, fitting even a single NF to a signal is often prohibitively expensive (e.g., more than a half GPU day for a single HD video [30]), limiting the field's ability to learn NFs at scale for a large set of signals. To alleviate this issue, optimization-based meta-learning has recently emerged [53, 61], demonstrating rapid improvements in learning efficiency and thus quickly becoming a standard technique in NF research. Meta-learning allows signals to be represented at good quality within only a handful of optimization steps by replacing tabula rasa optimization from a hand-picked
Figure 1: Online context point selection learns interpretable curricula: Visualization of selected context points (top), residuals relative to the original signal (middle), and the reconstruction (bottom) of GradNCP trained on ImageNet. Selected coordinates are highlighted in red (top) where our selection scheme focuses on only  $25\%$  of available points at each of the  $k$  adaptation steps. GradNCP first focuses on modeling global structure while subsequently learning high-frequency details.
with an optimized initialization closer to minima reachable within a few steps of gradient descent. This is achieved by defining a second-order optimization problem over the initialization and other meta-parameters, computed through the iterative learning process of individual tasks, all started from a shared initialization [1, 15]. In the context of this paper, the inner learning task is the fitting of an NF to a specific signal (e.g., a single audio clip), while the outer problem corresponds to learning an initialization suitable for a whole set of signals of the same type (e.g., an entire speech dataset).
While promising in principle, existing optimization-based meta-learning schemes suffer from problems of their own, primarily poor scaling with the dataset size (often called context set) used in an inner learning task. This issue is particularly severe for NFs, with a single  $1024 \times 1024$  image interpreted as a context set of well over a million function evaluations (i.e.,  $((x,y),(r,g,b))$  pairs), forming a dataset comparable in scale to ImageNet [8]. This is problematic since firstly, existing techniques are designed for the entire dataset be used in each inner optimization step, and secondly, the number of necessary inner steps required for adequate performance is known to increase with the context size. As both the computational graph and intermediate computations (e.g., gradients and activations) are stored in memory for back-propagation w.r.t. to the parameters of the outer problem, meta-learning at scale quickly becomes prohibitively expensive, particularly in terms of memory requirements. Moreover, default optimization-based meta-learning is known to suffer from myopia [16, 66], i.e., the tendency for plasticity to drastically decrease when the inner learning problem is run for more than the number of steps the initialization was optimized for. This strongly decreases the appeal of meta-learning when extremely high-quality NF representations are desired.
In this paper, we investigate whether this high computational burden can be alleviated by changing the form of the inner learning problem. In particular, we draw inspiration from curriculum learning [21] and recent interest in data pruning [44, 58], selecting a relevant subset of data points non-uniformly at each inner learning iteration to reduce the memory, and more crucially, reducing myopia through long horizon optimization made by the saved memory.
Contributions. We present a novel optimization-based meta-learning framework, coined Gradient Norm-based Context Pruning (GradNCP), designed for large-scale problems and thus particularly suitable for learning complex NFs. Specifically, GradNCP involves online data pruning of the context set based on a new score, making memory efficiency the primary objective of our algorithmic design. By focusing learning on the most important samples, GradNCP effectively implements an automated curriculum, rapidly extracting the global structure of the signal before dedicating additional iterations to the representation of high-frequency details, resembling a hand-crafted technique for efficient NF training in previous literature [33] (see Figure 1 for an illustration of this phenomenon). As this allows the computational graph to be released for all unused context points, the resulting memory savings enable longer adaptation horizons leading to improved meta-initializations and crucially, make it possible to train on higher-dimensional signals.
To realize efficient context pruning, we argue for two key considerations: (a) the minimization of information loss caused by the context pruning by careful selection based on per-example scoring statistics (b) the correction of any remaining performance degradation and avoidance of the afore
Figure 2: Computational diagram of GradNCP. A meta-learned initialization  $\theta_0$  is adapted for  $K$  steps to obtain  $\theta_K$ , re-ranking and pruning the context set  $\mathcal{C}_{\mathrm{high}}$  at each step for memory efficiency. Subsequently, we obtain a bootstrapped target  $\theta_{K + L}^{\mathrm{boot}}$  through  $L$  additional steps of (reduced-memory) optimization using the full context set  $\mathcal{C}_{\mathrm{full}}$ . Meta-optimization then minimizes both reconstruction error and the distance between the two parameter states.  $\theta_0$  is thus updated to allow minimization of this distance in the original  $K$  steps, correcting for context pruning and reducing myopia.
mentioned myopia phenomenon, allowing for longer optimization horizons using the full context at test time (when memory requirements are lower). To this end, we propose the following techniques:
1. Online automated context pruning: As large datasets involve a range of information content in its examples (e.g., in long-tail distributions), we hypothesize that learning can be greatly accelerated by focusing each training step on the examples with highest expected model improvement. To achieve this, we approximate the overall effect of updating the model on each example by the norm of the last layer gradient, an extremely cheap operation that results in virtually indistinguishable performance when compared to significantly more expensive methods while outperforming random sub-sampling by a large margin. Furthermore, we show how this technique is particularly suitable for meta-learning, demonstrating how the learning dynamics increase the correlation between the full gradient and our last layer approximation when compared to the random initialization.
2. Bootstrap correction: To counteract possible information loss introduced by the context pruning, we propose to nudge parameters obtained through optimization with a pruned context set to be as close as possible to an informative target obtained using the full set. We achieve this by constructing the target through continued adaptation using the full context and then defining a meta-loss involving the minimization of a distance between the two parameters. This acts both as a correction term while also reducing myopia, enabling advanced test-time adaptation. Importantly, this correction introduces significantly less computational overhead than simply extending meta-learning, as the target is deliberately generated without the need to save any intermediate gradients (i.e., meta-optimization is oblivious to the bootstrap generating procedure).
3. Test-time gradient scaling: Finally, we introduce a gradient scaling technique that allows the use of full context sets at meta-test time (when memory consumption is a significantly smaller concern) enabling the learning of the highest quality NFs while continuing to benefit from significant speed-ups afforded by meta-learning.
We verify the efficacy of GradNCP through extensive evaluations on multiple data modalities, including image, video, audio, and manifold datasets. We exhibit that our GradNCP consistently and significantly outperforms previous meta-learning methods. For instance, measured by peak signal-to-noise ratio (PSNR), GradNCP improves prior state-of-the-art results by  $38.28 \rightarrow 40.60$  on CelebA [36], and  $28.86 \rightarrow 35.28$  on UCF-101 [57]. GradNCP can also lead to accelerated training, reaching the same performance as existing meta-learning methods whiled reducing wall-clock time by  $\approx 45\%$ . Moreover, we show that GradNCP can meta-learn on extremely high-dimensional signals, e.g.,  $256 \times 256 \times 32$  videos, where prior works suffer from a memory shortage on the same machine.
2 GradNCP: Meta-Learning via Gradient Norm-based Context Pruning
2.1 Meta-Learning Neural Fields
We consider the problem of learning neural fields (NFs) for  $N$  signals  $\mathbf{s}_1, \ldots, \mathbf{s}_N$  by taking a functional view of each signal  $\mathbf{s}_i$ , represented by a context set  $\mathcal{C}^{(i)} := \{(\mathbf{x}_j, \mathbf{y}_j)\}_{j=1}^M$  consisting of  $M$  coordinate-value pairs  $(\mathbf{x}_j, \mathbf{y}_j)$  with  $\mathbf{x}_j \in \mathbb{R}^C$  and  $\mathbf{y}_j \in \mathbb{R}^D$ . For each signal  $\mathbf{s}_i$ , we aim to
find parameters  $\theta^{(i)}$  for the NF  $f_{\theta^{(i)}}: \mathbb{R}^C \to \mathbb{R}^D$  which reconstructs  $\mathbf{s}_i$  when evaluated on the full coordinate set (we also consider the generalization setup in Section 3.3). To learn NF, we minimize the mean-squared error (MSE) loss, i.e.,  $\ell(\theta; \mathcal{C}^{(\cdot)}) = \mathcal{L}_{\mathrm{MSE}}(\theta; \mathcal{C}^{(\cdot)}) := \frac{1}{M} \sum_{j=1}^{M} ||f_{\theta}(\mathbf{x}_j) - \mathbf{y}_j||_2^2$ . However, this requires a large number of optimization steps (on the order of 10ks for even moderately complex signals) thus limiting the development of NFs for large databases.
A popular approach for this acceleration [61] relies on model-agnostic meta-learning (MAML, [15]) by defining an optimization process over the initialization  $\theta_0$  with a given step size  $\alpha > 0$ :
$$
\theta_ {0} = \underset {\theta_ {0}} {\arg \min } \frac {1}{N} \sum_ {i = 1} ^ {N} \ell \left(\theta_ {0} - \alpha \nabla_ {\theta_ {0}} \ell \left(\theta_ {0}; \mathcal {C} ^ {(i)}\right); \mathcal {C} ^ {(i)}\right), \tag {1}
$$
where the inner optimization is iterated for  $K$  steps (typically  $K < 100$ ) before a meta-optimization step w.r.t  $\theta_0$  is taken. While mini-batch SGD is commonly used for the optimization w.r.t  $\theta_0$ , the entire context set  $|\mathcal{C}^{(i)}| = M$  (for NFs in the order of 100ks) is required for the update in each of the inner  $K$  steps (a legacy of MAML's predominant use in few-shot learning) with both intermediate computations and the computational graph maintained in memory, causing significant memory concerns. Unfortunately, both simple sub-sampling of  $\mathcal{C}^{(i)}$  at each step or common first-order simplifications of the MAML objective [15, 41] have been shown to significantly under-perform in the NF context [11, 12], thus leaving practitioners with no easy solution, forcing a sub-optimal reduction of  $K$  to single digits or abandonment of the meta-learning paradigm altogether.
2.2 Gradient Norm-based Online Context Pruning
To alleviate this memory burden, we propose an online context pruning strategy, resulting in the choice of a context subset  $\mathcal{C}_{\mathrm{high}}^k \subset \mathcal{C}_{\mathrm{full}}$  at each step  $k$  according to some scoring function  $R_k(\mathbf{x}, \mathbf{y})$  evaluated on each example  $(\mathbf{x}, \mathbf{y}) \in \mathcal{C}_{\mathrm{full}}$ . We denote this pruned context set  $\mathcal{C}_{\mathrm{high}}^k$  to distinguish from the full context  $\mathcal{C}_{\mathrm{full}}$  and define  $\mathcal{C}_{\mathrm{high}}^k \coloneqq \operatorname{TopK}(\mathcal{C}_{\mathrm{full}}; R_k, \gamma)$  where  $\operatorname{TopK}$  is a function returning the  $\gamma | \mathcal{C}_{\mathrm{full}}|$  elements with highest  $R_k$  score. Here  $\gamma \in (0, 1)$  is a hyper-parameter controlling the trade-off between (expected) performance and memory demands. We postulate that  $R_k(\mathbf{x}, \mathbf{y})$  rank examples according to their effect on the overall loss when being used to update the model in the inner loop of Eq. (1), hypothesizing that a suitable  $R_k$  ought to improve on random sub-sampling.
Estimating sample importance via one-step adaptation. The most direct but highly impractical method full-filing the requirement posed on  $R_{k}$  is the curriculum-learning metric target-prediction gain (TPG) [21], defined as  $R_{k}^{\mathrm{TPG}}(\mathbf{x},\mathbf{y}) = \ell (\theta_{k},\mathcal{C}\setminus (\mathbf{x},\mathbf{y})) - \ell (\theta_{k}^{\prime},\mathcal{C}\setminus (\mathbf{x},\mathbf{y}))$ , i.e. the direct measurement of the reduction in loss on all other training examples after updating the full network  $\theta_{k}$  to  $\theta_k^\prime$  using  $(\mathbf{x},\mathbf{y})$ , thus requiring  $\mathcal{O}(2M + M(M - 1)) = \mathcal{O}(M^2)$  operations (A forward and backward pass per example and evaluation on all other  $M - 1$  examples) rendering this metric inconsistent with the primary objectives of this work. A simpler alternative is the related self-prediction gain (SPG)  $R_{k}^{\mathrm{SPG}}(\mathbf{x},\mathbf{y}) = \ell (\theta_{k},\{\mathbf{x},\mathbf{y}\}) - \ell (\theta_{k}^{\prime},\{\mathbf{x},\mathbf{y}\})$  approximating  $R_{k}^{\mathrm{TPG}}$  by measuring learning progress only on the example used to update the model, reducing the operations to linear complexity  $\mathcal{O}(3M)$  (forward pass, model update, another forward pass). While undoubtedly a simplification, this metric has increased appeal for reconstructing NFs where we do not require generalization to previously unseen data in a fashion. Nevertheless, despite significant complexity reduction relative to  $R_{k}^{\mathrm{TPG}}$ , the  $3M$  operations of  $R_{k}^{\mathrm{SPG}}$  remain impractical.
Motivated by recent studies that emphasize the critical role of the last output layer in meta-learning NFs [70, 48, 49] and recognizing the relative inexpensiveness of the operation, we instead focus on an update  $\theta_k'$  from  $\theta_k$  involving only the last layer, leading to  $\mathcal{O}(M)$  complexity. We then further enhance efficiency by using a first-order Taylor approximation, showing SPG under this update to be well approximated by the squared gradient norm of the last layer. Specifically, letting  $\theta_k^{\mathrm{base}}$  be the set of parameters excluding the final layer and  $g_k = [\mathbf{0}, \nabla_{W_k}\ell, \nabla_{b_k}\ell]$  a padded gradient involving zeros for all but the last layer's weights  $W_k$  and biases  $b_k$ , we start from the SPG metric:
$$
\begin{array}{l} \ell \left(\theta_ {k}; \{(\mathbf {x}, \mathbf {y}) \}\right) - \ell \left(\theta_ {k} ^ {\prime}; \{(\mathbf {x}, \mathbf {y}) \}\right) \\ \approx \ell \left(\theta_ {k}; \left\{\left(\mathbf {x}, \mathbf {y}\right) \right\}\right) - \ell \left(\theta_ {k} - \alpha g _ {k}; \left\{\left(\mathbf {x}, \mathbf {y}\right) \right\}\right) \quad (\text {L a s t}) \\ \approx \ell (\theta_ {k}; \{(\mathbf {x}, \mathbf {y}) \}) - \left(\ell (\theta_ {k}; \{(\mathbf {x}, \mathbf {y}) \}) - \alpha g _ {k} ^ {\top} \nabla_ {\theta_ {k}} \ell (\theta_ {k}; \{(\mathbf {x}, \mathbf {y}) \})\right) (T a y l o r \\ = \alpha g _ {k} ^ {\top} \nabla_ {\theta_ {k}} \ell \left(\theta_ {k}; \left\{\left(\mathbf {x}, \mathbf {y}\right) \right\}\right) = \alpha \| g _ {k} \| _ {2} ^ {2}. \\ \end{array}
$$
Algorithm 1 Meta-training of GradNCP
Input: Initial  $\theta_0, \{\mathbf{s}_i\}_{i=1}^N, \gamma, \alpha, \beta, \lambda, K, L$   
1: while not converge do  
2: Sample batch  $\{\mathbf{s}_1, \dots, \mathbf{s}_B\}$ .  
3: for all  $b = 1$  to  $B$  do  
4: Extract context  $\mathcal{C}_{\mathrm{full}}$  from  $\mathbf{s}_b$ .  
5: for all  $k = 0$  to  $K - 1$  do  
6: # Online Context pruning  
7:  $\mathcal{C}_{\mathrm{high}}^k = \mathrm{TopK}(\mathcal{C}_{\mathrm{full}}; R_k, \gamma)$   
8:  $\theta_{k+1} \gets \theta_k - \alpha \nabla_{\theta_k} \mathcal{L}_{\mathrm{MSE}}(\theta_k; \mathcal{C}_{\mathrm{high}}^k)$   
9: end for  
10: # Generate target in  $L$  steps  
11:  $\theta_{K+1}^{\mathrm{boot}} \gets \theta_K - \alpha \nabla_{\theta_K} \mathcal{L}_{\mathrm{MSE}}(\theta_K; \mathcal{C}_{\mathrm{full}})$   
12: ...  
13:  $\mathcal{L}_{\mathrm{total}}^k = \mathcal{L}_{\mathrm{MSE}}(\theta_K; \mathcal{C}_{\mathrm{full}}) + \lambda \mu(\theta_K, \theta_{K+L}^{\mathrm{boot}})$   
14: end for  
15:  $\theta_0 \gets \theta_0 - \beta \frac{1}{B} \sum_{b=1}^{B} \nabla_{\theta_0} \mathcal{L}_{\mathrm{total}}^b$   
16: end while  
Output:  $\theta_0$
Algorithm 2 Meta-testing of GradNCP
Input: Test signal s, learned initialization  $\theta_0$ ,  $K_{\mathrm{test}}$   
1: Extract context  $\mathcal{C}_{\mathrm{full}}$  from s.  
# Where typically  $K_{\mathrm{test}} > K + L$   
2: for all  $t = 0$  to  $K_{\mathrm{test}} - 1$  do  
3: # Context pruning  
4:  $\mathcal{C}_{\mathrm{high}} = \mathrm{TopK}(\mathcal{C}_{\mathrm{full}}; R_t, \gamma)$   
# Compute gradient scaling  
 $g_t^{\mathrm{test}} = \frac{\|\nabla_{\theta_t}\mathcal{L}_{\mathrm{MSE}}(\theta_t,\mathcal{C}_{\mathrm{high}})\|}{\|\nabla_{\theta_t}\mathcal{L}_{\mathrm{MSE}}(\theta_t,\mathcal{C}_{\mathrm{full}})\|}\nabla_{\theta_t}\mathcal{L}_{\mathrm{MSE}}(\theta_t,\mathcal{C}_{\mathrm{full}})$   
# Adaptation with full context  
5:  $\theta_{t+1} \gets \theta_t - \alpha \cdot g_t^{\mathrm{test}}$   
6: end for  
Output:  $\mathcal{L}_{\mathrm{MSE}}(\theta_T,\mathcal{C}_{\mathrm{full}}), \theta_T$
where the equivalence in the last layer results from all gradients w.r.t  $\theta_{k}^{\mathrm{base}}$  in  $\nabla_{\theta_k}\ell (\theta_k;\{(\mathbf{x},\mathbf{y})\}$  being multiplied with zero by their corresponding entry in  $g_{k}$ . As the gradient for the final layer takes the standard form for a linear model (see non-linear version and the full derivation in Appendix B), we can define our sample importance score  $R_{k}^{\mathrm{GradNCP}}$  under  $\mathcal{L}_{\mathrm{MSE}}$  at adaptation step  $k$  as:
$$
R _ {k} ^ {\operatorname {G r a d N C P}} (\mathbf {x}, \mathbf {y}) := \left\| \left(\mathbf {y} - f _ {\theta_ {k}} (\mathbf {x})\right) \left[ \phi_ {\theta_ {k} ^ {\text {b a s e}}} (\mathbf {x}), \mathbf {1} \right] ^ {\mathsf {T}} \right\| \tag {2}
$$
where  $\|\cdot\|$  is the Frobenius norm and  $\phi_{\theta_k^{\mathrm{base}}}(\cdot)$  the penultimate feature vector, such that  $f_{\theta_k}(\mathbf{x}) \coloneqq W_k \phi_{\theta_k^{\mathrm{base}}}(\mathbf{x}) + b_k$ . This score only requires a single forward pass (and an additional small matrix multiplication) without the need for full backpropagation, making it straightforward and scalable to use for high-resolution signals. This is particularly suitable for an extensions of MAML which we use in practice [35], that considers the step size at each iteration  $K$  of the inner learning problem part of the meta-optimization parameters, adding increased flexibility in terms of both choosing the most suitable data and update size for each iteration. We further emphasize that while motivated in the NF context, this approximation can be generally applicable in curriculum learning settings [44, 58].
2.3 Reducing Bias and Myopia through Bootstrap Correction
Despite the careful selection of  $\mathcal{C}_{\mathrm{high}}^k$ , a large reduction in the size of the context set may inevitably introduce information loss. To mitigate this issue, we suggest counteracting any bias by reducing a distance to a parameter state obtained using optimization carried with the full context set. Specifically, after adapting the meta-learner for  $K$  steps with the pruned context set, we generate a bootstrapped target  $\theta_{K + L}^{\mathrm{boot}}$  through optimization for  $L$  additional steps starting from  $\theta_K$ . Crucially, we consider meta-optimization to be oblivious to the target generating process, i.e., do not backpropagate through its updates, saving memory otherwise required for second-order optimization. In automated differentiation frameworks, this can be implemented with a stop-gradient operation on  $\theta_{K + L}^{\mathrm{boot}}$ , leading to each of the  $L$  additional steps being significantly cheaper as well as improved target performance.