Improving Generalization and Safety of Deep Neural Networks with Masked Anchoring

Thank you for the positive assessment of our work. If our responses sufficiently address all the concerns raised, we request you to improve the score and champion our paper.

Choice of hyper-parameter $\alpha$:

• The parameter $\alpha$ directly controls the schedule of residual corruption during training. Simply put, arbitrarily increasing $\alpha$ will start to adversely affect the training convergence and eventually the ID performance itself. In our study, we selected it solely based on two factors: We wanted to use a high value of $\alpha$ (for improved generalization), such that it can converge in the same number of epochs as that of anchoring w/o RAM, and can produce validation accuracies comparable (or better) than the non-anchored model.

• Based on our experiments with CIFAR-10, CIFAR-100 and ImageNet across multiple architectures (RegNet, ResNet, WRN40-2, DEIT, ViTb and SWINv2) $\alpha = 0.2$ is a recommended setting. Similar to any existing regularization strategy (e.g., Mixup) adopted for improving generalization, RAM with high $\alpha$ might have the risk of producing models with reduced sensitivity towards anomalous data (i.e., compared to anchoring w/o RAM). However, we find that with $\alpha = 0.2$, the anomaly rejection trade-off is minimal and interestingly, it can even lead to higher rejection AUROC scores using networks with higher capacity (e.g., SWINv2-B, updated Table 2 in the paper).

How does RAM help anchoring?

Please check our detailed answer to this central question in the summary response. For completion, we include a summary here.

• To clarify this, we begin by considering generalization to an OOD test sample. In the tuple $(\mathrm{r}, \Delta)$ for the test sample $\mathrm{x}_t$, the anchor $\mathrm{r} \in P(R)$ and it is possible that $\Delta = \mathrm{x}_t - \mathrm{r}$ can be novel to P($\Delta$). Thus, by exposing the model to more diverse combinations of $(\mathrm{r}, \Delta)$ during training, generalization can be improved. A naive way to improve the diversity of P($\Delta$) (or equivalently P($X, \Delta$)) is to consider a wider anchor distribution (i.e., $P(R) \supset P(X)$). However, it is non-trivial to characterize the anchor distribution and arbitrarily improve its diversity; and increasing anchor diversity will lead to a combinatorially larger space (anchor, residual) pairs, which can in turn require larger number of epochs to effectively converge.

• We adopt an alternative approach to increasing diversity of P($\Delta$) by making the residuals noisy, and propose a simple implementation in the form of RAM. Given the inherent challenge in defining a suitable residual noise distribution (and inferring its hyperparameters), we define the noise distribution to be same as the anchor distribution itself, i.e., P(R) = P(N) = P(X), and implement it efficiently using the RAM regularizer. Formally, for a given tuple $(\mathrm{r}, \mathrm{x} - \mathrm{r})$, anchor masking zeroes out the anchor while keeping the residual fixed, i.e., $(\mathrm{0}, \mathrm{x} - \mathrm{r})$. Generally, the tuple for making a prediction for $\mathrm{x}$ with a zero anchor (note: zero vector is a valid anchor in our anchor distribution) will be written as $(\mathrm{0}, \mathrm{x})$. However, in anchor masking this can be interpreted as making a prediction using the zero anchor, but with a noised residual $\mathrm{x} + \mathrm{\epsilon}$ where $\mathrm{\epsilon} = -\mathrm{r}$ and $\mathrm{\epsilon} \in P(N)$.

Thank you for your positive feedback. Here is our detailed response to all your questions.

1. Choice of Acronym RAM:

We understand how this terminology can be confusing. Following reviewer advice, we will use MAR (Masking Anchors Randomly) as the new acronym. We have not changed it yet so that the other reviewers do not get confused.

2. Implementation of RAM:

The reviewer is right in pointing out that the parameter $\alpha$ controls the schedule of residual corruption during training -- across an entire batch. However, this is only an implementation choice and one can consider alternative approaches, and the observations will not change. For example, the residual corruption can be applied to $\alpha \%$ of samples from each batch or this can be included as an additional loss objective with a pre-specified penalty.

3. Complexity of Anchored Training w/o and w/ RAM:

From an optimization standpoint, anchored training becomes increasingly challenging as the combinatorial space of (anchor, residual) pair grows. For example, with CIFAR-10 or CIFAR-100, anchored training converges effectively with the same training recipe in the same number of epochs as the standard training. However, with ImageNet (regardless of the architecture), for the same training recipe, anchored training requires $20$ additional epochs to converge to the same level of validation loss. However, the training efficiency itself is not significantly different from non-anchored models. Note that, including RAM regularization ($\alpha = 0.2$ in all experiments) does not impact the training behavior and requires the same number of epochs as standard anchoring. Here we show the comparisons for the SWINv2-B model trained on ImageNet.

4. Improving captions of Tables and Figures:

Thanks for this comment. We have expanded the figure and table captions for better clarity.

5. Hyperparameter Selection:

For all experiments reported in the paper, we follow the training recipes from the torchvision (https://pytorch.org/vision/stable/models.html) and directly used the same hyper-parameter configurations. The only difference was increasing the number of epochs for the case of imagenet ($20 more epochs).

For the hyperparameter $\alpha$, we selected it solely based on two factors: We wanted to use a high value of $\alpha$ (for improved generalization), such that it can converge in the same number of epochs as that of anchoring w/o RAM, and can produce validation accuracies comparable (or better) than the non-anchored model. Based on our experiments with CIFAR-10, CIFAR-100 and ImageNet across multiple architectures (RegNet, ResNet, WRN40-2, DEIT, ViTb and SWINv2) $\alpha = 0.2$ is a recommended setting.

6. What is the $\pm$ shown in results?

We apologize for this oversight. $\pm$ in the calibration metric indicates the standard deviation across all OOD datasets (ImageNet-C, Imagenet-$\bar{\text{C}}$, ImageNet-v2, ImageNet-R and ImageNet-S) and we have made this clear in the updated version.

7. SoTA Comparison:

The key claim of this work is that, across dataset sizes (CIFAR-10 to Imagenet), architectures (ResNet to ViT) and network sizes (5M to 88M parameters), anchored training can provide significant gains in OOD generalization, anomaly rejection and adaptation, compared to vanilla training. We do not attempt to establish SoTA generalization performance, since top results are often obtained via model souping[1] or by fine-tuning large scale pre-trained models[2]. However, we emphasize that, when the train recipe includes high-capacity architectures or advanced training strategies (Mixup, EMA, label-smoothing, cutmix), anchoring continues to improve over the base models in all these cases, showing promise to be effective even with SoTA recipes involving large-scale pre-training and ensembling. In the updated conclusions, we have highlighted that integrating anchoring (w/ RAM) into these SoTA approaches is an important future direction of work.

[1] M. Wortsman et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, ICML 2022

[2] S.Goyal et al. Finetune like you pretrain: Improved finetuning of zero-shot vision models, CVPR 2023.

Thank you for a positive assessment of our paper. If our response sufficiently addresses your questions, we request you to improve the score and champion our paper.

1. Why does anchored training and the RAM regularizer lead to better generalization?

Please check our detailed answer to this central question in the overall response. For completion, we include a summary here. While we have not theoretically characterized the behavior of RAM, our hypothesis is rooted in existing theory and our expanded empirical results provide clear evidence for the hypothesis. However, motivated by the empirical success of anchoring, we plan to carry out a theoretical study on generalization in anchored models as part of our future work (as specified in the updated Conclusion section).

• In essence, anchoring involves representing an input $\mathrm{x}$ as a combination of an anchor $\mathrm{r}$ and the residual $(\mathrm{x} - \mathrm{r})$, while maintaining all other aspects of the conventional neural network. This method capitalizes on the lack of shift invariance in Neural Tangent Kernels (NTKs) induced by common neural networks. Anchored training employs different anchors for the same sample across epochs to marginalize the effect of anchor choice during inference. While it may seem akin to data augmentation, anchoring is fundamentally different in that it combinatorially expands the space of (anchor, residual) pairs with different anchor choices and enables exploration of a richer class of hypotheses, and does not manipulate the inputs. However, as our experiments show, anchoring can be combined with any additional data augmentation technique.

• Now, let us consider generalizing to an OOD test sample. In the tuple $(\mathrm{r}, \Delta)$ for the test sample $\mathrm{x}_t$, the anchor $\mathrm{r} \in P(X)$ and it is possible that $\Delta = \mathrm{x}_t - \mathrm{r}$ can be novel to $P(\Delta)$. Thus, by exposing the model to more diverse combinations of $(\mathrm{r}, \Delta)$ during training, generalization can be improved. A naive way to improve the diversity of $P(\Delta)$ (or equivalently $P(X, \Delta)$) is to consider a wider anchor distribution (i.e., $P(C) \supset P(X)$). However, it is non-trivial to characterize the anchor distribution and arbitrarily improve its diversity; and increasing anchor diversity will lead to a combinatorially larger space (anchor, residual) pairs, which can in turn require larger number of epochs to effectively converge.

• We adopt a simpler alternative by making the residuals noisy, and propose a simple implementation in the form of RAM. Given the inherent challenge in defining a suitable noise distribution (and inferring its hyperparameters), we define the noise distribution to be same as the anchor distribution itself, i.e., $P(C) = P(N) = P(X)$, and implement it efficiently using the RAM regularizer. Formally, for a given tuple $(\mathrm{r}, \mathrm{x}-\mathrm{r})$, anchor masking can be reinterpreted as $(\mathbf{0}, \mathrm{x} + \epsilon)$ where $\epsilon \in P(N)$. In other words, we want to identify the true label for $\mathrm{x}$ with the zero anchor, given the noisy residual.

• Finally, network capacity plays a significant role in effectively leveraging the increased diversity produced by RAM. For example, with ImageNet, as we move from RegNet (5M) to SWINv2-B (88M), we witness larger performance improvements over both anchoring w/o RAM as well as standard training (see updated Table 1 in the paper).

2. Choice of $\alpha$ and the Generalization-Anomaly Rejection Trade-off:

• The parameter $\alpha$ directly controls the schedule of residual corruption during training. Simply put, arbitrarily increasing $\alpha$ will start to adversely affect the training convergence and eventually the ID performance itself. Hence, we recommend a nominal value of $\alpha=0.2$, which we empirically find to converge in the same number of epochs as anchoring w/o RAM across all architectures, and to result in much improved generalization.

• Similar to any existing regularization strategy (e.g., Mixup) adopted for improving generalization, RAM with high $\alpha$ has the risk of producing models with reduced sensitivity towards anomalous data (i.e., compared to anchoring w/o RAM) [1]. We find that with $\alpha = 0.2$, the anomaly rejection trade-off is minimal and interestingly, it can even lead to higher rejection AUROC scores with networks with higher capacity (e.g., SWINv2-B, updated Table 2 in the paper).

[1] Hendrycks, Dan, et al. "Pixmix: Dreamlike pictures comprehensively improve safety measures.", CVPR 2022.

Thank you for your positive comments and the questions. If our responses sufficiently address all the concerns raised, we request you to improve the score and champion our paper.

How does anchoring work and why does RAM lead to better generalization?

In order to clarify how anchoring works and why the proposed RAM regularization leads to improved generalization, we have provided a detailed explanation in the overall response. Here we include an expanded version of that with some additional comments to address your questions. While we do not theoretically establish the utility of RAM, our hypothesis is rooted in existing theory and our expanded empirical results provide clear evidence for the hypothesis.

• In essence, anchoring involves representing an input $\mathrm{x}$ as a combination of an anchor $\mathrm{r}$ and the residual $(\mathrm{x} - \mathrm{r})$, while maintaining all other aspects of the conventional neural network. This method capitalizes on the lack of shift invariance in Neural Tangent Kernels (NTKs) induced by common neural networks. Anchored training employs different anchors for the same sample across epochs to marginalize the effect of anchor choice during inference.

• While it may seem akin to data augmentation, anchoring is fundamentally different in that it combinatorially expands the space of (anchor, residual) pairs with different anchor choices and does not manipulate the inputs. However, as our experiments show, anchoring can be combined with any additional data augmentation technique.

• For the first time, we systematically study the viability of anchoring as a training protocol across dataset sizes (CIFAR-10 to Imagenet), architectures (ResNet to ViT) and network sizes (5M to 88M parameters). In this context, we also study other key properties like calibration, anomaly rejection and adaptation (both ID and OOD evaluation). We have entirely reorganized the results section for clarity.

• Now, let us consider generalizing to an OOD test sample. In the tuple $(\mathrm{r}, \Delta)$ for the test sample $\mathrm{x}_t$, the anchor $\mathrm{r} \in P(X)$ and it is possible that $\Delta = \mathrm{x}_t - \mathrm{r}$ can be novel to $P(\Delta)$. Thus, by exposing the model to more diverse combinations of $(\mathrm{r}, \Delta)$ during training, generalization can be improved. A naive way to improve the diversity of $P(\Delta)$ (or equivalently $P(X, \Delta)$) is to consider a wider anchor distribution (i.e., $P(C) \supset P(X)$). However, it is non-trivial to characterize the anchor distribution and arbitrarily improve its diversity; and increasing anchor diversity will lead to a combinatorially larger space (anchor, residual) pairs, which can in turn require larger number of epochs to effectively converge.

• We adopt an alternative approach to increasing diversity of $P(\Delta)$ by making the residuals noisy, and propose a simple implementation in the form of RAM. Given the inherent challenge in defining a suitable residual noise distribution (and inferring its hyperparameters), we define the noise distribution to be same as the anchor distribution itself, i.e., $P(C) = P(N) = P(X)$. Formally, for a given tuple $(\mathrm{r}, \mathrm{x}-\mathrm{r})$, anchor masking zeroes out the anchor while keeping the residual fixed, i.e., $(0, \mathrm{x}-\mathrm{r})$. Generally, the tuple for making a prediction for $\mathrm{x}$ with a zero anchor (note: zero vector is a valid anchor) will be written as $(\mathbf{0}, \mathrm{x})$. However, in anchor masking this can be interpreted as making a prediction using the zero anchor, but with a noised residual $\mathrm{x} + \epsilon$ where $\epsilon \in P(N)$.

2. Trade-off between OOD Generalization and Anomaly Rejection:

The parameter $\alpha$ controls the schedule of residual corruption during training. Arbitrarily increasing $\alpha$ can adversely affect the training convergence and eventually the ID performance itself. Hence, we recommend a nominal value of $\alpha=0.2$, which we empirically find to converge in the same number of epochs as anchoring w/o RAM across all architectures, and to result in much improved generalization. Similar to any existing approach (e.g., Mixup) adopted for improving generalization, RAM with high $\alpha$ could lead to reduced sensitivity towards anomalous data [1]. However, we find that with $\alpha = 0.2$, the anomaly rejection trade-off is minimal and interestingly, it can even lead to higher rejection AUROC scores with networks with higher capacity (e.g., SWINv2-B, updated Table 2).

[1] Hendrycks, Dan, et al. "Pixmix: Dreamlike pictures comprehensively improve safety measures.", CVPR 2022.

Thank you for your comments and questions. Following your feedback, we have expanded the background, reorganized the results and clearly stated our findings. If our responses sufficiently address all the concerns raised, we request you to improve the score and champion our paper.

1. How does anchoring work and why does RAM lead to better generalization?

In order to clarify how anchoring works and why the proposed RAM regularization leads to improved generalization, we have provided a detailed explanation in the overall response. Here we include an expanded version. While we have not theoretically characterized the behavior of RAM, our hypothesis is rooted in existing theory and our expanded empirical results provide clear evidence for the hypothesis. However, motivated by the empirical success of anchoring, we plan to carry out a theoretical study on generalization in anchored models as part of our future work (see updated Conclusion section).

• In regard to theoretical foundations for anchoring, we directly refer to two key results: (a) In [1], it was shown that centering a dataset using different constant inputs will lead to different solutions, due to inherent lack of shift invariance in NTKs induced by commonly adopted neural networks. Anchored training uses different anchors for the same sample across different epochs with the goal of marginalizing out the effect of anchor choice at inference time. But in this process, it implicitly explores a large class of hypotheses and enforces prediction consistency; (b) In [2], it was theoretically showed that an anchored model can demonstrate improved extrapolation behavior at test time, when the anchor $\mathrm{x} \in P(X)$ and the residual for the unseen sample $\Delta \in P(\Delta)$. However, neither of these works rigorously studied the generalization behavior of anchored models on large-scale datasets, architectures or even practical distribution shifts. We have also accordingly updated the background section.

• While the idea of enforcing prediction consistency across different anchor choices might appear similar to data augmentation, we want to clarify that anchoring does not impose any invariance to data characteristics, but only expands the combination of (anchor, residual) pairs with each additional anchor.

• For the first time, we systematically study the viability of anchoring as a training protocol across dataset sizes (CIFAR-10 to Imagenet), architectures (ResNet to ViT) and network sizes (5M to 88M parameters). In this context, we also study other key properties such as calibration, anomaly rejection and adaptation (both ID and OOD evaluation). We have entirely reorganized the results section for clarity.

• Let us consider generalizing to an OOD test sample. In the tuple $(\mathrm{r}, \Delta)$ for the test sample $\mathrm{x}_t$, the anchor $\mathrm{r} \in P(X)$ and it is possible that $\Delta = \mathrm{x}_t - \mathrm{r}$ can be novel to $P(\Delta)$. Thus, by exposing the model to more diverse combinations of $(\mathrm{r}, \Delta)$ during training, generalization can be improved. A naive way to improve the diversity of $P(\Delta)$ (or equivalently $P(X, \Delta)$) is to consider a wider anchor distribution (i.e., $P(C) \supset P(X)$). However, it is non-trivial to characterize the anchor distribution (e.g., for ImageNet) and arbitrarily improve its diversity; and increasing anchor diversity will lead to a combinatorially larger space of (anchor, residual) pairs, thus requiring a larger number of epochs to effectively converge.

• We adopt an alternative approach to increasing diversity of $P(\Delta)$ by making the residuals noisy, and propose a simple implementation in the form of RAM. Given the challenge in defining a suitable residual noise distribution, we define the noise distribution to be same as the anchor distribution itself, i.e., $P(C) = P(N) = P(X)$. Formally, for a given tuple $(\mathrm{r}, \mathrm{x}-\mathrm{r})$, anchor masking zeroes out the anchor while keeping the residual fixed, i.e., $(\mathbf{0}, \mathrm{x}-\mathrm{r})$. Generally, the tuple for making a prediction for $\mathrm{x}$ with a zero anchor (note: zero vector is a valid anchor) will be written as $(\mathbf{0}, \mathrm{x})$. However, anchor masking can be interpreted as making a prediction using the zero anchor, but with a noised residual $\mathrm{x} + \epsilon$ where $\epsilon \in P(N)$.

• Finally, network capacity plays an important role in effectively leveraging the increased diversity produced by RAM. For example, with ImageNet, as we move from RegNet (5M) to SWINv2-B (88M), we witness larger generalization performance improvements over both anchoring w/o RAM as well as standard training (see updated Table 1).

[1] J.J. Thiagarajan et al. Single model uncertainty estimation via stochastic data centering, Neurips, 2022.

[2] A. Netanyahu et al. "Learning to Extrapolate: A Transductive Approach." ICLR 2023.

We thank you for the positive feedback. If our responses sufficiently address all the concerns raised, we request you to improve the score and champion our paper.

1. Why does Anchoring help?

Kindly check our detailed answer to this central question in the overall response. For completion, we include a summary here.

In essence, anchoring involves representing an input $\mathrm{x}$ as a combination of an anchor $\mathrm{r}$ and the residual $\mathrm{x} - \mathrm{r}$, while maintaining the other training conditions of conventional neural networks. This method capitalizes on the lack of shift invariance in Neural Tangent Kernels induced by common neural networks. Anchored training employs different anchors for the same sample across epochs to marginalize the effect of anchor choice during inference. While it may seem akin to data augmentation, anchoring is fundamentally different in that it combinatorially expands the space of (anchor, residual) pairs with different anchor choices and enables exploration of a richer class of hypotheses, and does not manipulate the inputs. However, as our experiments show, anchoring can be combined with any additional data augmentation technique.

2. How does RAM work and Risk of Overfitting:

Kindly check our detailed answer to this central question in the overall response. For completion, we include a summary here.

To clarify this, we begin by considering generalization to an OOD test sample. In the tuple $(\mathrm{r}, \Delta)$ for the test sample $\mathrm{x}_t$, the anchor $\mathrm{r} \in P(X)$ and it is possible that $\Delta = \mathrm{x}_t - \mathrm{r}$ can be novel to P($\Delta$). Thus, by exposing the model to more diverse combinations of $(\mathrm{r}, \Delta)$ during training, generalization can be improved. A naive way to improve the diversity of $P(\Delta)$ (or equivalently $P(X, \Delta)$) is to consider a wider anchor distribution (i.e., $P(R) \supset P(X)$). However, it is non-trivial to characterize the anchor distribution and arbitrarily improve its diversity; and increasing anchor diversity will lead to a combinatorially larger space (anchor, residual) pairs, thus requiring larger number of epochs to effectively converge.

We adopt an alternative approach to increasing diversity of $P(\Delta)$ by making the residuals noisy, and propose a simple implementation in the form of RAM. Given the inherent challenge in defining a suitable residual noise distribution, we define the noise distribution to be same as the anchor distribution itself, i.e., $P(C) = P(N) = P(X)$, and implement it efficiently using the RAM regularizer. Formally, for a given tuple $(\mathrm{r},\mathrm{x} - \mathrm{r})$, anchor masking zeroes out the anchor while keeping the residual fixed, i.e., $(0, \mathrm{x} - \mathrm{r})$. Generally, the tuple for making a prediction for $\mathrm{x}$ with a zero anchor (note: zero vector is a valid anchor in our anchor distribution) will be written as $(0, \mathrm{x})$. However, in anchor masking this can be interpreted as making a prediction using the zero anchor, but with a noised residual $\mathrm{x} + \epsilon$ where $\epsilon = -\mathrm{r}$ and $\epsilon \in P(N)$.

3. Anchor Selection at Test Time:

During both training and testing, the anchors are sampled from the anchor distribution $P(R)$, which is set to the training distribution $P(X)$ itself. However, given that anchoring explicitly marginalizes the choice of anchor, any random training sample (or a small number of them) can be used to obtain predictions at inference time. We also emphasize that, during anchoring, the anchor's label is never used and the anchor sample only contributes to the residual computation.

4. Comparison with SoTA:

The key claim of this work is that, across dataset sizes (CIFAR-10 to Imagenet), architectures (ResNet to ViT) and network sizes (5M to 88M parameters), anchored training can provide significant gains in OOD generalization, anomaly rejection and adaptation, compared to vanilla training. We do not attempt to establish SoTA generalization performance, since top results are often obtained via model souping [1] or by fine-tuning large scale pre-trained models [2]. However, we emphasize that, when the train recipe includes high-capacity architectures or advanced training strategies (Mixup, EMA, label-smoothing, cutmix), anchoring continues to improve over the base models in all these cases, showing promise to be effective even with SoTA recipes involving large-scale pre-training and ensembling. In the updated conclusions, we have highlighted that integrating anchoring (w/ RAM) into these SoTA approaches is an important future direction of work.

[1] M. Wortsman et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, ICML 2022

[2] S.Goyal et al. Finetune like you pretrain: Improved finetuning of zero-shot vision models, CVPR 2023.