Thumb on the Scale: Optimal Loss Weighting in Last Layer Retraining

We thank the reviewer for their comments and suggestions. We would like to address each point presented:

• W1: We included additional results in the supplementary materials B which show similar results for additional class pairs. If accepted, we will update the manuscript with results for additional modalities such as text to expand the applicability of our work. Additionally, while our current results capture only binary classification for class imbalance, we aim to use similar techniques to tackle the more general problem of group imbalance in binary classification. We believe that this work is an important first step in understanding how the loss weighting process is affected by overparameterization.
• W2: By the conventional prior ratio, we meant the inverse ratio of the class priors, $\pi_-/\pi_+$, which is very frequently used in practice. We will ensure that this point is clarified in the next revision of this work.
• W3: We thank the reviewer for pointing out these oversights. We will ensure that any typos are fixed and the above terms are fully expanded earlier in the manuscript.

We thank the reviewer for their careful review of our work and are grateful for the opportunity to clarify these points. We will address each of the presented weakness and questions:

• W1: While many of the dimensions in the last layer go unused or have little impact, the threshold of $\delta=1$ is still critical in theory. At that point, the single optimizer of square loss is not well defined and theory must resort to discussing the minimum norm interpolator of the training data. This threshold, and in particular the regime of $\delta>1$, is also used to argue against the efficacy of loss weighting in large models [1]. In contrast, we study the underparameterized regime which we demonstrate is amenable to loss weighting. We will clarify this point in an update to our manuscript if accepted.
• W2: The CGMT framework allows us to discuss more complex data distributions such as mixtures of Gaussians per class or more general covariance matrices. This work focuses on the identity covariance matrix for simplicity of presentation, but any shared covariance matrix is captured by our theory through a whitening transformation. We hope to extend this framework to weighting for mixtures of Gaussians per class, but this work is a first step towards that goal. Indeed we see that our analysis provides valuable insights in real datasets which are certainly non-Gaussian and non-isotropic. We will ensure that these clarifications are made in the revision.
• Q1: In our analysis of the auxiliary optimization, the only use of the asymptotics is through the Law of Large Numbers and kicks in relatively quickly. Indeed, we report results for n=20 and show that the theoretical weighting scheme is still effective despite the very small number of samples and dimensions.
• Q2: Indeed, one should suspect that the (relative) input complexity should also play a role in the selection of $\rho$. The central aim of our manuscript is to capture the effect of parameterization in the selection of loss weights, and we focus on the identity covariance as a clean basis for this investigation. We hope that this work serves as inspiration for future work exploring other factors that should be considered when selecting loss weight, including relative input complexity.
• Q3: One nice result of our work is that it captures the population case ($\delta\rightarrow0$) as a special case. We point out in [Section 3.1] that our weighting scheme $\tilde\rho$ recovers the ratio of priors which is optimal in the population setting. Additionally our results on downsampling in [Section 3.2] align with [2] when $\delta\rightarrow0$. Our results offer the flexibility to account for the parameterization ratio and are complementary to population results.

We hope that these clarifications enhance the reviewer's understanding of our work and will ensure that these clarifications are included in the revision.

References:

[1] J. Byrd and Z. C. Lipton, “What is the Effect of Importance Weighting in Deep Learning?,” presented at the International Conference on Machine Learning, Dec. 2018.

[2] M. Welfert, N. Stromberg, and L. Sankar, “Theoretical Guarantees of Data Augmented Last Layer Retraining Methods,” May 09, 2024, arXiv: arXiv:2405.05934. doi: 10.48550/arXiv.2405.05934.

We thank the reviewer for their careful consideration of our work. We will consider each of the weakness in turn:

• W1: The construction of the optimal weighting scheme $\tilde\rho$ relies on the linearity of the last layer retraining, but one could conjecture that the same trends would likely hold: namely the more parameters that are being finetuned (with respect to the number of samples and their dimension), the stronger of a weight is needed to achieve the same balance. A key challenge in settings rather than the last layer retraining would be quantifying the effective dimension as our current approach relies on a fixed latent distribution. Spreading weights across different layers would complicate this significantly.
• W2: A key challenge in characterizing losses other than the square loss is in computing the explicit form of the expected Moreau envelope; in other words, such closed form expressions do not exist for most losses making the optimization harder in theory. While the reviewer correctly points out that the scalarization holds for any loss, the reduction beyond that point is dependent on the form of the Moreau envelope. Recent work [1] has found that finetuning models using square loss can achieve equal or better classification performance to using standard cross entropy, which we discuss in the manuscript as motivation for studying the optimal weighting scheme for square loss.
• W3: This assumption may appear strong, but recent work [2,3] has argued that such assumptions in the last layer can be sufficient to guarantee performance for distributions sharing only the first and second order statistics. Additionally, while the nature of the CGMT framework requires that the latent features have some variant of a Gaussian distribution, there is the possibility of extending this work to mixtures of Gaussians which may more closely resemble the true latent distribution of a deep model. Finally, our empirical results align with our theoretical findings, suggesting that the assumptions we make might not be too strong to be useful in practice.

References:

[1] A. Achille, A. Golatkar, A. Ravichandran, M. Polito, and S. Soatto, “LQF: Linear Quadratic Fine-Tuning,” Dec. 21, 2020, arXiv: arXiv:2012.11140. doi: 10.48550/arXiv.2012.11140.

[1] R. Ghane, D. Akhtiamov, and B. Hassibi, “Universality in Transfer Learning for Linear Models,” Feb. 23, 2025, arXiv: arXiv:2410.02164. doi: 10.48550/arXiv.2410.02164.

[2] Q. Han and Y. Shen, “Universality of regularized regression estimators in high dimensions,” Jun. 27, 2022, arXiv: arXiv:2206.07936. doi: 10.48550/arXiv.2206.07936.

We thank the reviewer for their thorough reading of our work. We are very glad that you found the paper relevant and engaging. We hope to address each of the weaknesses and questions mentioned.

• W1: Thank you for pointing this out. The key to showing uniqueness (and therefore convergence to the asymptotic discussed in this work) is to show the joint strict convexity-concavity of the objective in (37) in the supplement. For square-loss, we can take a direct approach by evaluating the expected Moreau Envelope (eME) functions in (37) leveraging that the square-loss is self-dual, hence its Moreau envelope is an easy-to-work-with quadratic. After easy calculation of the expectations over the (scalar) Gaussian random variables, and after combining the two individual eME functions (corresponding to majorities/minorities respectively) we can show that the combination (which is of the form ($\alpha^2+\gamma^2+b^2+1)/\lambda$) is jointly strictly convex (in $\alpha,\gamma,b,\lambda$) by direct evaluation of the Hessian. We note that strict convexity-concavity holds only when both eME terms are considered together, as each individual eME is not jointly strictly convex-concave. Beyond square-loss, we can follow the steps of App. A.6.1 in [1]. The only key change, following the “lesson” learnt by the explicit evaluation of the square-loss above, is to combine the two eME terms. We will include this analysis in our final version and update corollary 1 to reflect this. We will also clarify in Thm. 1 that the convergence to the deterministic parameters holds in-probability (as previous works e.g. [1,24 in manuscript]).
• W2: First, we highlight the fact that asymptotics kick in at relatively small dimensions and samples is a positive feature of our analysis. Non-asymptotic analyses are more challenging. While the CGMT is non-asymptotic in nature, the technical challenge comes in the analysis of the Auxiliary optimization. Recent literature [2,3,4] has made significant progress in this direction, but given the necessary technical overhead such an extension goes beyond our scope. We recognize the value of such non-asymptotic extension, which for example can yield confidence bounds and will certainly discuss this as current limitation and a possible direction for future work. For our main goal of re-weight optimization it is encouraging that direct optimization of the asymptotic formulas gives insights that match with our empirics.
• Q1: Recent work [5,6] has considered the multi-class case, and our scalarization (Theorem 1) would be amenable to these updated analyses. However, the analysis of the auxiliary optimization and the weighting scheme would need to change. For instance, it is not obvious that the multi-class extension of square loss would be amenable to the same closed-form analysis. Additionally, the optimal weighting scheme in our work is chosen such that $b=0$ which is enough to ensure that both classes are fairly classified. In the multiclass setting, however, this condition would be insufficient without significant additional structure on the class means and variances. We hope to extend our work to this setting in the future, and see this manuscript as an important and valuable step in that direction.
• Q2: The construction of the optimal weighting scheme $\tilde\rho$ relies on the linearity of the last layer retraining, but one could conjecture that the same trends would likely hold: namely the more parameters that are being finetuned (with respect to the number of samples and their dimension), the stronger of a weight is needed to achieve the same balance. A key challenge in settings rather than the last layer retraining would be quantifying the effective dimension as our current approach relies on a fixed latent distribution. Spreading weights across different layers would complicate this significantly.

References:

[1] H. Taheri, R. Pedarsani, and C. Thrampoulidis, “Sharp Asymptotics and Optimal Performance for Inference in Binary Models,” Feb. 26, 2020, arXiv: arXiv:2002.07284. doi: 10.48550/arXiv.2002.07284.

[2] L. Zhou, F. Koehler, P. Sur, D. J. Sutherland, and N. Srebro, “A Non-Asymptotic Moreau Envelope Theory for High-Dimensional Generalized Linear Models,” Oct. 21, 2022, arXiv: arXiv:2210.12082.

[3] G. Wang, K. Donhauser, and F. Yang, “Tight bounds for minimum $\ell_1$-norm interpolation of noisy data,” in Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, PMLR, May 2022, pp. 10572–10602.

[4] Q. Han and Y. Shen, “Universality of regularized regression estimators in high dimensions,” Jun. 27, 2022, arXiv: arXiv:2206.07936.

[5] B. Loureiro, G. Sicuro, C. Gerbelot, A. Pacco, F. Krzakala, and L. Zdeborová, “Learning Gaussian Mixtures with Generalised Linear Models: Precise Asymptotics in High-dimensions”.

[6] D. Bosch and A. Panahi, “A Novel Convex Gaussian Min Max Theorem for Repeated Features”. Proceedings of The 28th International Conference on Artificial Intelligence and Statistics