Can the Training Loss be Predictive for Out-of-Distribution Generalization?

Authors: Lorenzo Brigato, Stavroula Mougiakakou
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: deep learning, OOD generalization, signal propagation, hyper-parameter search

Abstract:
Traditional model selection in deep learning relies on carefully tuning several hyper-parameters (HPs) controlling regularization strength on held-out validation data, which can be challenging to obtain in scarce-data scenarios or may not accurately reflect real-world deployment conditions due to distribution shifts.
Motivated by such issues, this paper investigates the potential of using solely the training loss to predict the generalization performance of neural networks on out-of-distribution (OOD) test scenarios.
Our analysis reveals that preserving consistent prediction variance across training and testing distributions is essential for establishing a correlation between training loss and OOD generalization.
We propose architectural adjustments to ensure $\textit{variance preservation}$, enabling reliable model selection based on training loss alone, even in over-parameterized settings with a sample-to-parameter ratio exceeding four orders of magnitude.
We extensively assess the model-selection capabilities of $\textit{variance-preserving}$ architectures on several scarce data, domain-shift, and corruption benchmarks by optimizing HPs such as learning rate, weight decay, batch size, and data augmentation strength.

Decision: reject
Reviewer scores: 6, 3, 3, 3

Reviews and discussion:
<thread>
creation date:
1730533732086

thread_type:
official_review

<thread_object>
creator: reviewer_keea

<content>
summary:
In this paper, the authors explore the research question: "whether the training loss can be used as a reliable predictor of OOD performance". The authors found that neural network architecture may be an important factor for the reliability of predictions. Then, the authors study the conditions and neural architectural designs to control the prediction-variance across distributions. Empirical results on example datasets demonstrate of the effectiveness of this methodology. There are several points the authors may address to further clarify this paper. See questions.

Overall, this is an interesting paper and some points can be addressed or even refocus on the HP optimization to further improve it. Current title seems to be too broad for this paper.

strengths:
Good presentations and writtings.

weaknesses:
See summary.

questions:
1)	What are the potential practical implications for this paper? As this paper focus on the neural architectural design to improve the reliability for prediction. Then,

   1.1	does the proposed architecture brings additional benefits or drawbacks or equivalent performances compared with normal neural architectures.  The authors need to show the performance comparisons before and after architectural modifications.

   1.2	Have the authors explored any techniques to improve OOD performance prediction without modifying existing architectures? If not, what are the key challenges in doing so?"

2)	Comprehensive evaluations on different OOD distribution shifts. The authors are suggested to consider different kinds of OOD distribution shifts. As shown in [1] and [2], there are different types of distribution shifts in real-world datasets, it is suggested to consider two types of distribution shifts including diversity shifts and correlation shifts. In this paper, spurious correlations are not considered. Or alternatively, justifying why focusing on the diversity shift instead if this is the primary focus on this paper.
 
[1] A Critical Analysis of Out-of-Distribution Generalization
[2] Ood-bench: benchmarking and understanding out-of-distribution generalization datasets and algorithms

3) Experiments of this method for hyper-parameter optimization. As this method can potentially serve as a way to help improving HP optimization, the authors can consider introduce this to existing HP optimization pipelines to wider experiments, such as PACS, and ImageNet to see how it works. What advantages will this method demonstrate over other HP optimization baselines.  Besides, it could also (might) be beneficial if the focus of this paper be moved into the HP optimization as this shall have strong practical implications in various fields. 

4) Experiments on standard OOD datasets. The authors could explain why choosing the datasets in this paper over standard OOD benchmark datasets. This paper can consider standard datasets in OOD generalization, such as PACS, ColoredMNIST,  VLCS, etc as well as what is done in most OOD generalization papers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730536372748

thread_type:
official_review

<thread_object>
creator: reviewer_mcks

<content>
summary:
This paper investigates whether training loss could be used to predict OOD performance. It claims that preserving consistent prediction variance across training and testing distributions is essential. This paper analyzes the factors that may affect prediction variance and further proposes variance-preserving architectures.

strengths:
* The attempt to bridge the training loss with the model performance on the testing set, especially the OOD test set is interesting.

* Generally this paper is clear. The figures in this paper are pretty.

weaknesses:
* The relationship between the training testing loss correlation and the correlation between each predicted label seems fra-fetched. As claimed in Eq.4, it requires the correlation between the predicted label of each of the two training testing data samples to be close to 1, which seems impossible.

* The experimental setting that random samples a subset from the CIFAR10 training set and tests on the CIFAR10 testing set does not qualify as an OOD scenario. The distribution of the CIFAR10 training set is the same as that of the CIFAR10 testing set. Randomly sampling a subset does not change the distribution.

* According to the results shown in Fig. 2, 3, 4, 5, and 6, these models show **the typical over-fitting pattern**. It seems that the proposed adjustment increases the correlation because it prevents over-fitting. However, overfitting is caused by the training on a small subset of the original training set, which could be simply dealt with. 

Considering the above-mentioned weaknesses, I have doubts about the analyses in this paper and the proposed variance-preserving architecture. In my understanding, the proposed method limits the ability of the network to prevent overfitting, which also explains the performance drop as shown in Table 2.

questions:
As mentioned in the weakness part, I have doubts about the analysis, the experimental setting, and the proposed method in this paper. I would like to see a rebuttal regarding the mentioned weaknesses. Currently, I will recommend reject for this paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730540395431

thread_type:
official_review

<thread_object>
creator: reviewer_9ygd

<content>
summary:
This work investigates the conditions under which the training loss can predict OOD generalization. The authors find that maintaining consistent prediction variance across training and testing distributions is key to correlating training loss with OOD generalization. They propose architectural adjustments to preserve this variance and demonstrate that these variance-preserving architectures allow for reliable model selection based on training loss in various OOD scenarios.

strengths:
1. The study of OoD generalization and model robustness has always been an important topic. Investigating model selection and generalization ability from the perspective of network architecture is a worthwhile endeavor.

2. The paper is well written and easy to follow.

weaknesses:
1. The authors’ claim that the RQ in line 69 is an unexplored research question is inaccurate, as this topic has already been studied extensively in prior work, such as [1] and [2]. This study's approach, which examines the problem from the perspectives of variance preservation and network architecture, is also not new; similar analyses have been conducted in previous studies, such as [3] and [4]. This work does not sufficiently compare itself to these prior studies and build upon them for further exploration.

2. It is inappropriate to directly ignore the gradient term $\nabla_f \mathcal{L}$ in formulas (2) and (3). It is a function of $𝑓$ and $\boldsymbol{w}$, and choosing different $𝑓$ and $\boldsymbol{w}$ will yield different $\nabla_f \mathcal{L}$, thereby impacting $\text{Cov}(J(\boldsymbol{w}, \boldsymbol{h}), J'(\boldsymbol{w}, \boldsymbol{h}))$ and $\text{Var}(J(\boldsymbol{w}, \boldsymbol{h}))$. Considering only the Cov term in formulas (2) and (3) leads to an incomplete conclusion.

3. The architectural adjustments proposed in this work are already commonly used for training neural networks and do not provide any new insights.

[1] Ishaan Gulrajani, David Lopez-Paz, In Search of Lost Domain Generalization, ICLR 2021.

[2] Ye et al., Towards a Theoretical Framework of Out-of-Distribution Generalization, NeurIPS 2021.

[3] Krueger et al., Out-of-Distribution Generalization via Risk Extrapolation, ICML 2021.

[4] Bai et al., NAS-OoD: Neural Architecture Search for Out-of-Distribution Generalization, ICCV 2021.

questions:
1. With a substantial amount of research already available, why is the RQ in line 69 considered an unexplored research question? Is there a key point I might have overlooked?

2. Why can the gradient term $\nabla_f \mathcal{L}$ in formulas (2) and (3) be ignored?

3. Which of the architectural adjustments proposed in this work are newly introduced?

Please see Weaknesses for details.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692328534

thread_type:
official_review

<thread_object>
creator: reviewer_u9kt

<content>
summary:
This work investigates the use of training loss to predict the generalization performance of neural networks on out-of-distribution scenarios, instead of using the validation data. This work is particularly useful in distribution-shifted scenarios, where standard validation data may be unavailable. This work examines the conditions for establishing linear relationships between training and test losses. Comprehensive empirical analysis is provided to analyze the model-selection capabilities of the proposed framework through extensive experiments setup, including optimizing several hyper-parameters across popular OOD benchmarks.

strengths:
1. This work proposes a novel approach for model selection by relying solely on training loss to predict out-of-distribution generalization. This is different from the traditional validation-based model selection and is particularly suitable for scenarios involving distribution shifts.

2. This paper analyzes the conditions for which training loss can correlate with OOD generalization, and develop a variance-preserving architecture based on these conditions.

3. Various hyper-parameters are thoroughly investigated across different types of OOD benchmarks, including corruption and domain shifts.

weaknesses:
1. It would be valuable to analyze and investigate on different types of OOD shifts, including diversity shifts and correlation shifts [1], and discuss whether the proposed conditions hold across these two dimensions of OOD shifts.

2. Extending the evaluation to more complex real-world natural images, such as the WILDS benchmarks [2].

3. Providing additional theoretical justification for supporting and explain the proposed variance preserving architecture, would enhancing the reasoning and foundation of the proposed framework.

Reference:

[1] OoD-Bench: Quantifying and Understanding Two Dimensions of Out-of-Distribution Generalization. CVPR 2021.

[2] WILDS: A Benchmark of in-the-Wild Distribution Shifts. ICML 2021.

questions:
Refer to the detailed suggestions provided in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733081356871

thread_type:
official_comment

<thread_object>
creator: area_chair_rdrx

<content>
comment:
Dear Reviewers,

If you have not already responded to the latest comments by the authors, please engage with them to clarify any remaining issues as the discussion period is coming to a close in less than a day (2nd Dec AoE for reviewer responses).

Thanks for your service to ICLR 2025.

Best, 
AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1733311688806

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal period

comment:
We thank all the reviewers for the time they invested in evaluating our paper and for their constructive feedback.  Below, we summarize the strengths, our revisions, clarifications, and points of disagreement.

# Strengths

**S1.** Reviewer U9Kt considered our approach novel and particularly suitable for scenarios involving distribution shifts.

**S2.** Reviewer 9YgD appreciated the idea of analyzing the influence of the architecture on model selection and generalization.

**S3.** Reviewer mCKs liked the idea of bridging the training loss with OOD generalization, and Reviewer keeA found our paper interesting as a whole.

**S4.** All reviewers agreed that our paper was clearly written and generally had a good presentation, figures in particular.

# Revisions

**R1.** **Characterization of the experimented distribution shifts (Reviewer U9Kt, keeA).** We followed the [official code](https://github.com/m-Just/OoD-Bench/tree/main) from [OODBench](https://openaccess.thecvf.com/content/CVPR2022/papers/Ye_OoD-Bench_Quantifying_and_Understanding_Two_Dimensions_of_Out-of-Distribution_Generalization_CVPR_2022_paper.pdf) to classify the 13 employed datasets across correlation-/diversitty-shift axes and added results to Appendix C.1.

**R2.** **Lack of novelty of our Research Question (RQ) (Reviewer 9YgD).** We added a description in the Related Work section regarding the works referenced by the reviewer, which better clarifies the novelty of our RQ. Those works perform model selection employing a validation signal while we make model selection from the training loss alone, as also acknowledged in **S1**.

**R3.** **Dependence between hyper-parameters ($\mathbf{h}$) and model parameters ($\mathbf{w}$) (Reviewer 9YgD).** We made more explicit (Section 2.1) our assumption that for a fixed architecture-HP-configuration pair, the learning process converges to a fixed parameter set $\mathbf{w}$, which is reasonable under the condition of a fixed optimizer since the predictions of a neural network do not vary significantly across repeated runs with the same HP configuration.

# Clarifications

**C1.** **Extending the evaluation to more complex real-world natural images (Reviewer U9Kt).** We clarified that several of our benchmarks represent real-world natural images, e.g., the ISIC2018 dataset was part of an imaging challenge in MICCAI 2018.

**C2.** **Novelty of architectural adjustments introduced (Reviewer 9YgD).** We clarified that the novelty does not come from the individual components but the architecture as a whole since architectures preserving variance were never proposed.

**C3.** **The small datasets employed do not qualify as an OOD scenario (Reviewer mCKs).** We clarified that learning from small/insufficient data has been cast as an OOD problem by [(Wad et al., 2022)](https://arxiv.org/pdf/2207.12258) as visible in Figure 2 from their work. Aggressively subsampling a large dataset incurs the learner to face an OOD scenario. We also provided additional evidence of this via **R1** since all the small datasets showed either diversity (more) or correlation (less) shifts.

**C4.** **Performance comparisons before and after architectural modifications (Reviewer keeA).** See Sections 2.2, 3.2, 3.3.
 
**C5.** **Key challenges to predict OOD performance in standard architectures (Reviewer keeA).** Pointed to Section 2.2.

**C6.** **Experiments on standard OOD datasets (Reviewer keeA, mCKs).** We employed popular OOD datasets such as CIFAR10-/100-C and domain-shifted versions of Tiny Imagenet (-C/-V2/-A/-R), following the Imagenet ones (due to computational constraints).

**C7.** **Focus of this paper is to be moved into the HP optimization literature (Reviewer keeA).** Our paper dedicated a Related Work paragraph to HP optimization, where we discussed its position and tested extensive HP optimization setups.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the potential for the training loss to predict out-of-distribution (OOD) generalization performance. It does this through a specially designed neural network architecture that aims to control prediction variance across distributions, motivated by a correlation-based analysis of training loss and test performance. The proposed architecture is then tested on a range of OOD datasets where it is shown to indeed have high correlation between training loss and OOD performance.

Strengths: The paper proposes an interesting and novel approach to an important topic, and is well presented.

Weaknesses: The proposed architecture suffers a large (> 5%) performance drop in many cases, and the focus on small data settings may conflate the effects of "overfitting" due to model estimation error (lack of data) vs domain shift, especially in the case study. There is also limited evaluation on real-world OOD datasets.

Overall, while the paper presents an interesting research direction, the practical impact of the work is somewhat unclear due to the large performance drops of the proposed architecture. Also, as OOD generalization generally requires assumptions, it would be useful to understand the conditions under which the proposed architecture works (in the sense of having correlated training and OOD loss) and does not; analysis of its performance on a broader range of real-world datasets will help address this. The AC believes this is an interesting direction of work but in its current form is not quite ready for publication, and hopes the comments will help strengthen the paper for a future submission.