Learning Transferable Features for Implicit Neural Representations

We thank the reviewer for careful and thorough evaluation of our work. We appreciate the reviewer’s kind comments on STRAINER’s originality in learning transferable representations for INRs, extensive evaluation of STRAINER on datasets, and the detailed analysis provided by STRAINER on feature transferability and its training dynamics. We address the comments and suggestions below.

Theoretical insights: Why the shared encoder layers facilitate better learning.

Suppose we have an INR consisting of $L$ fully connected layers and non-linearity $\sigma$ after each layer. Assume that $\sigma$ is a piecewise-linear approximation of the sine function. For any piecewise linear non-linearity, the INR can be considered a continuous piecewise-affine operator, i.e., the INR subdivides the input space into pieces or regions $\omega \in \Omega$, where $\Omega_L$ is the input space partition due to all $L$ layers. Each region/piece of the function is mapped to the output in an affine manner.

The subdivision $\Omega_L$, happens in a layer-wise fashion, e.g., layer $\ell$ contributes $\partial \Omega_\ell$ to the overall subdivision [13]. An important outcome of this layerwise subdivision is that shallower layers learn coarser partitions of the input domain, therefore resulting in shallower layers learning lower frequency features [55]. Deeper layers however can be more localized therefore, learn higher frequency or less symmetric features.

The motivation behind INR is to learn a joint partition of the shallower layers, that can be further subdivided by the deeper layers into an optimal input space partition, resulting in a low error continuous piecewise-affine operator.

Unstable test time optimization for STRAINER. We attribute the instability of STRAINER at test time fitting to the underlying SIREN model. In SP Figure 7 of the supplementary, reconstruction quality (PSNR) plots of Siren and STRAINER models plotted across time reflect the periodic instability. However, this can be fixed by using a learning rate (LR) scheduler. As shown in RT Figure1(b), we apply an exponential multiplicative LR scheduler that decays the learning rate by 0.4 resulting in a smooth and stable learning for STRAINER-10 (shown in green).

Training distribution of images. We conduct experiments in in-domain (ID) and out of domain (OD) image fitting where the shared encoder is trained on CelebA-HQ[16] and tested on CelebA-HQ, AFHQ[5], and OASIS-MRI[12,21]. We notice the superior performance of STRAINER compared to baselines (see RT Table 1) for in-domain image fitting. We also find that STRAINER captures highly transferable features. For OD image fitting on AFHQ, we find that STRAINER is able to reconstruct OD AFHQ almost comparable to ID fitting on AFHQ reconstructions, differing by a single digit in PSNR (MP Table 1,2). While for OASIS-MRI, we find STRAINER trained on CelebA-HQ surprisingly outperform even ID fitting on OASIS-MRI. To further assess the effect of training distribution on STRAINER, we train the shared encoder on Flowers [45] and Stanford Cars [46] as they have different spatial distribution of color and high frequencies compared to AFHQ and OASIS-MRI. We find that STRAINER is able to learn transferable features and achieves comparable PSNR to the shared encoder trained on CelebA-HQ[16]. As shown in RT Table 1, in-domain fitting for AFHQ[5] has the highest reconstruction PSNR, followed by STRAINER trained on CelebA-HQ, Flowers, and Stanford Cars (in decreasing order, presented in RT Table 1). The results suggest that STRAINER trained on natural images yields highly transferable representations, that inturn generalize well across natural images. We propose this as a direction for future work.

Comparison with other transfer learning techniques. We compare STRAINER with other metalearning and transformer baseline models such as Meta-Learned 5K[37], TransINR[47] and IPC[48] (with and without test time optimization) for image fitting. We evaluate STRAINER and baselines on CelebA-HQ (in-domain) , AFHQ(out of domain) and OASIS-MRI(out of domain). As seen in RT Table 1, STRAINER significantly outperforms IPC[48], TransINR[47] and other baselines by ~ 3-5+ db in PSNR. We see that pretrained TransINR and IPC yield good PSNR at iteration 0, STRAINER catches up within little time and then consistently outperforms up until the end of 2000 iterations.

We refer the reviewer to a more detailed summary of baseline results in our response to Reviewer FYQ8, “Comparison of baselines”, and “Metalearning baseline for Kodak”.

Additionally, please also refer to our response to Reviewer HDgr, “New Out-of-Distribution generalization experiments” for detailed evaluation of out of domain generalizability.

Due to limited time, our experiments have been carefully chosen to address the concerns of all reviewers and are indicative of STRAINER learning highly transferable features for tasks such as image representation, with minimum training time and as little as 10 images. A comparison of training statistics and model footprints is mentioned in RT Table 3. We are happy to engage more and provide more findings during the discussion period.

We sincerely thank the reviewer for their time, effort and their positive feedback in evaluating our work! We request the reviewer to kindly update their scores if our rebuttal experiments and analysis have satisfactorily addressed their comments and questions.

We thank the reviewer for careful and thorough evaluation of our work. We appreciate the kind comments made on the novelty of our method and performance over Siren/meta-learning baselines. We welcome the reviewers suggestion on comparing with more baselines, and are thankful for the references provided and we will update them in Section 2, Background.

Comparison of baselines: Prompted by the primary concern of comparisons to baselines, we further conduct thorough image fitting comparison of STRAINER with IPC[48] and TransINR[47] with and without test time optimization as detailed below. As seen in RT Table 1, STRAINER significantly outperforms IPC[48], TransINR[47] and other baselines by ~ 3-5+ db in PSNR. We attribute the superior performance of STRAINER to better input space partitioning learned by STRAINER as a consequence of sharing the initial layers, coupled with the instance specific decoder which provides an added degree of freedom allowing STRAINER to fit an unseen image with different morphology better. We see that pretrained TransINR and IPC yield good PSNR at iteration 0, STRAINER catches up within little time and then consistently outperforms up until the end of 2000 iterations.

We also note that STRAINER features are more transferable for out of domain image fitting as shown for AFHQ and OASIS-MRI, even when the STRAINER encoder has been pre-trained on domains such as Flowers[45] and Stanford Cars[46]. In-domain fitting for AFHQ has the best reconstruction quality, followed by STRAINER trained on CelebA-HQ, Flowers, and Stanford Cars (in decreasing order, presented in RT Table 1.).

More details on our baselines. We use the publicly available implementation of IPC[48] which also includes the baselines for TransINR, both of which use FFNets[38]. While STRAINER uses Siren[31], we remark that FFNets and Siren models offer comparable representation capabilities for image fitting. We train TransINR[47] and IPC[48] baselines on 14,000 images from CelebA-HQ for 300 epochs[48]. To verify successful training of the model, we ensure that the PSNR on the test set is in near agreement with that reported in their respective papers barring changes due to the number of data samples. We also ensure that all baselines and STRAINER MLPs have 6 layers and the same number of parameters. All models are evaluated on our test images for 2000 iterations. Lastly, for the OASIS-MRI dataset, we run all baselines using 3 channel MRI images and report metrics in RT Table 1.

STRAINER and IPC. In MP Figure 6.iv, we visualize the layerwise subdivisions / partitions in the input space of INRs. We observe that partitions evolve from coarse to fine as we go deeper in the INR. Initial layers provide for more global (low-frequency) transferable features. The deeper layers of the INR result in more complex subdivision of the input space giving rise to more local features and pertaining to high frequency detail in the image. As mentioned in MP Sec. 3.1, we motivate sharing layers in STRAINER to learn representations that give rise to partitions that generalize well across samples. IPC[48] proposes that the second layer in an INR be instance-specific while keeping all the remaining layers instance-agnostic. Modulating the second layer's weights would result in non-local changes resulting in suboptimal reconstruction. Further, IPC’s second layer weights obtained from a Transformer model may also affect the location of the optimal instance-specific layer.

We further adapt IPC’s design of an instance specific second layer in STRAINER and compare it to other STRAINER models. As shown in RT Figure 1(b), we see the best reconstruction given by STRAINER-10(ours), then STRAINER (instance-specific second layer) and then SIREN - which validates our above intuition on input space division for STRAINER and IPC.

Metalearning baseline for Kodak: STRAINER learns higher quality reconstruction as indicated by image metrics (PSNR, SSIM[54], and LPIPS[53]) compared to Learnit baseline as shown in RT Table 2. We train all models on CelebA-HQ images of the same resolution as Kodak[1] images. STRAINER performs better than SIREN and Meta-Learned 5K[37] with +2db in PSNR for network width of 256, and ~+3db for network width of 512. STRAINER(width=256) performs similar to SIREN(width=512), further attesting to the representation power of STRAINER’s learned features. We also remark that training time for STRAINER’s shared encoder on high resolution is negligible compared to meta learning models for high resolution image fitting - highlighting that STRAINER can be easily adapted to high resolution and out of domain images.

Ablation study on key hyperparameters: We refer the reviewer to SP Figure 7. We show an ablation study on the number of layers shared in the STRAINER encoder and its ability to reconstruct the signal. We observe progressive increase in reconstruction quality (PSNR) as we increase the number of layers shared.

Evaluating STRAINER on complex (3D) signals. Please refer to our response to Reviewer HDgr, “Evaluating STRAINER on 3D signals”

Due to limited time, our experiments have been carefully chosen to address the concerns of all reviewers and are indicative of STRAINER learning highly transferable features for tasks such as image representation, with minimum training time and as little as 10 images. A comparison of training statistics and model footprints is mentioned in RT Table 3. We are happy to engage more and provide more findings during the discussion period.

We thank the reviewer for the careful and thorough evaluation of our work. We appreciate the kind comments on STRAINER being able to learn fast at test time, efficient recovery of high frequency components, and its ability to learn more transferable input space subdivision. We address the comments and suggestions below.

Citing additional papers on meta-learning INR. We thank the reviewer for the additional references. We find them relevant to our literature review and will include them in Section 2, Background.

Evaluating STRAINER on 3D signals. STRAINER is a general purpose transfer learning framework which can be used to initialize INRs for regressing 3D data like occupancy maps, radiance fields or video. To demonstrate the effectiveness of STRAINER on 3D data, we have performed the following OOD generalization experiment. We pre-train STRAINER on 10 randomly selected ‘Chair’ objects from the ShapeNet[44] dataset. At test time, we fit the ‘Thai Statue’ 3D object[52]. STRAINER achieves a 12.3% relative improvement in IOU compared to random initialization for a SIREN architecture – in 150 iterations STRAINER-10 obtains an IOU of 0.91 compared to an IOU of 0.81 without STRAINER-10 initialization. We present visualizations of the reconstructed Thai Statue in RT Figure 1(c). Upon qualitative evaluation, we see that STRAINER-10 is able to capture ridges and edges better and faster than compared to SIREN. We will add the results in the main text in Sec 4. and compare with baselines LearnIt[37], and IPC[48].

Fitting signals (2D, 3D shapes, video) to STRAINER is straightforward and intuitive. However, extending STRAINER to inverse problems in 3D such as learning Neural Radiance Fields from limited observations and 3D Gaussian Splatting is non-trivial. We consider this an exciting future direction for us and thank the reviewer for their kind suggestions.

New Out-of-Distribution generalization experiments. Following the recommendations by the reviewer, we have trained STRAINER’s shared encoder on Flowers [45] and Stanford Cars [46] datasets and report reconstruction quality (PSNR) when testing on AFHQ [5] dataset in RT Table 1. We specifically chose Flowers [45] and Stanford Cars [46] as they have different spatial distribution of color and high frequencies compared to AFHQ and OASIS-MRI which would allow us to evaluate out of domain image fitting more thoroughly. We find that STRAINER is able to learn transferable features and achieves comparable PSNR to the shared encoder trained on CelebA-HQ[16]. As shown in RT Table 1, in-domain fitting for AFHQ has the highest reconstruction PSNR, followed by STRAINER trained on CelebA-HQ, Flowers, and Stanford Cars (in decreasing order, presented in RT Table 1). The results suggest that STRAINER is capable of capturing transferable representations that generalize well across natural images. We propose this as a direction for future work.

We also emphasize in RT Table 2 (and MP Table 3 ) where we train on CelebA-HQ but evaluate out of domain on high resolution Kodak Images[1] on categories of airplane and statue where STRAINER outperforms baseline models substantially or is very comparable. Further, we also compare STRAINER with recent work on generalizable INRs [47, 48]. We show in RT Table 1, that STRAINER outperforms meta-learning as well as transformer baselines such as TransINR[47] and IPC[48], for out of domain fitting for AFHQ and OASIS-MRI datasets even when STRAINER is trained on unrelated source domains such as stanford cars.

Due to limited time, our experiments have been carefully chosen to address the concerns of all reviewers and are indicative of STRAINER learning highly transferable features for tasks such as image representation, with minimum training time and as little as 10 images. A comparison of training statistics and model footprints is mentioned in RT Table 3. We are happy to engage more and provide more findings during the discussion period.

Generalization ability of STRAINER for Inpainting. We thank the reviewer for raising this important point. Successful recovery for inverse problems such as single-image inpainting relies on the inductive bias of the model. STRAINER’s rich representation encoded in the shared initial layers of an INR captures a prior over the data allowing it to converge rapidly fast. However, since it’s based on SIREN, its inductive bias is the same as siren - which may result in similar performance as Siren. We expect results similar to denoising and super-resolution as reported in MP Table 5 depending on the fidelity of the available signal. Further, for a non-trivial problem such as inpainting, conventional INRs may require support from added regularization terms if lacking strong inductive biases that favor solving inpainting, as shown by INRR [49]. Additionally, convolutional architectures such as the Deep Image Prior [50] leverage self-similarity and are successful at inpainting. While INRs do not exhibit locality bias, models such as Wire[51] have strong inductive bias due to the Gabor Wavelet which favors solving inverse problems. We’re curious to assess if STRAINER can capture any favorable properties which can benefit the task of inpainting and will explore it as future efforts to understand STRAINER deeply.