Debiasing Mini-Batch Quadratics for Applications in Deep Learning

Authors: Lukas Tatzel, Bálint Mucsányi, Osane Hackel, Philipp Hennig
Primary area: optimization
Keywords: quadratic Taylor approximation, mini-batching, second-order optimizers, conjugate gradients, uncertainty quantification, Laplace approximation, stochastic curvature, GGN, KFAC

Abstract:
Quadratic approximations form a fundamental building block of machine learning methods. E.g., second-order optimizers try to find the Newton step into the minimum of a local quadratic proxy to the objective function; and the second-order approximation of a network's loss function can be used to quantify the uncertainty of its outputs via the Laplace approximation. When computations on the entire training set are intractable - typical for deep learning - the relevant quantities are computed on mini-batches. This, however, distorts and biases the shape of the associated *stochastic* quadratic approximations in an intricate way with detrimental effects on applications. In this paper, we (i) show that this bias introduces a systematic error, (ii) provide a theoretical explanation for it, (iii) explain its relevance for second-order optimization and uncertainty quantification via the Laplace approximation in deep learning, and (iv) develop and evaluate debiasing strategies.

Decision: accept (poster)
Reviewer scores: 8, 8, 3, 8

Reviews and discussion:
<thread>
creation date:
1730488682490

thread_type:
official_review

<thread_object>
creator: reviewer_9wsh

<content>
summary:
Often one would like to obtain a second-order Taylor approximation to a deep learning objective based on a minibatch.  This paper shows that the use of a minibatch results in a strong overestimation of the top Hessian eigenvalues.  The basic problem is that the directions of highest curvature for a minibatch tend to slightly 'overfit' (my words) that minibatch, and do not generalize to other minibatches or to the full-batch objective.  This paper shows that this bias can be corrected with a scheme that uses two separate minibatches - one to compute the top eigenvectors, and another to compute the directional curvature along these eigenvectors.  The paper discusses applications in both optimization (conjugate gradient method for computing a newton step), and uncertainty quantification (laplace's method)

strengths:
The paper offers a simple solution to a problem of interest to multiple areas of deep learning.

weaknesses:
In the application to optimization, the paper does not actually carry out an end-to-end optimization run with a second-order method using the proposed conjugate gradient scheme; instead, the paper just uses the proposed scheme to compute a single newton step.

questions:
Regarding the motivation in section 3.1 -- just because the gradient lives in the top subspaces does _not_ mean that all the optimization progress occurs in that space.  To the contrary, the evidence suggests that the vast majority of the motion along the top Hessian eigenvectors is oscillatory, and 'cancels out.'  Long-term loss reduction is due to the components of the gradient that lie along the low-curvature Hessian eigenvectors.  For example, see "Does SGD really happen in tiny subspaces?" by Minhak Song, Kwangjun Ahn, Chulhee Yun: https://arxiv.org/abs/2405.16002.  

I don't think this is fatal for the paper, since you are basically only using your improved estimate of the the top Hessian eigenvalues in order to properly compute the conjugate gradient _step size_, which will depend on the top Hessian eigenvalues.  But if you actually tried to optimize only along those eigenvectors, it wouldn't work (as the linked paper shows).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730491121018

thread_type:
official_review

<thread_object>
creator: reviewer_ddr4

<content>
summary:
This paper demonstrates a systematic bias in stochastic 2nd order approximation of empirical risk for neural networks. The study start from an observation showing that a particular curvature structure on minibatch Hessian is biased towards underestimation. Then, authors proposes a debiasing method.

strengths:
- A computationally-cheap estimate of 2nd order information is an important topic for optimization 
- Ample experimental observations 
- Designing computational-cheap experiments to study 2nd order derivative of training loss of neural networks

weaknesses:
I believe the notion of unbiased estimate is not well defined and studied. Define the Hessian associate with full-batch training loss as  $$H(x) = \frac{1}{n} \sum_{i=1}^n \nabla^2 f_i(x)$$. Similarly, we define the unbiased estimate of this matrix as $$H'(x) = \frac{1}{m} \sum_{k=1}^m \nabla^2 f_{i_k}(x)$$ where $i_k$ are uniformly drawn from $\{1,\dots, n\}$. 
For each fixed vector $d$, we have $$E d^\top H' d^\top = d^\top H d$$. Thus, the **directional curvature**, defined in the paper, is unbiased. But, why this paper observe a bias in the estimate? The main issue is that they consider random directions $d$ depending on minibatches. I do not know why we need to choose random $d$s? 

I believe the quantity of interest is not well defined and motivated, here. It is easy to provide an unbiased estimation of  $q$, directional curvature and directional slop. But, the authors want to estimating another quality that they do not exactly defined and it is not clear how its estimation connect with Newton's method or Laplace approximation. 

If you want to estimate the maximum/minimum eigenvalue of the Hessian or even its condition number, the empirical Hessian provides an asymptotically unbiased estimate (see for example https://arxiv.org/pdf/1912.10754).

questions:
- I recommend to replace $u_1$ and $u_2$ in Figure 1 by leading eigenvectors of full-batch Hessian to better grasp my comment in weaknesses. I expect to see that increasing batch size will lead to a better approximation of the curvature. 
- What is exactly the quantity that you want to estimate and why it is important to estimate? Directional curvature on which random directions?
- What do you exactly mean by bias? To prove an estimate has bias, we need to take average. But, experiments in Figure 1 are not computing an average. How can I conclude from these experiments that the estimate is biased?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705094456

thread_type:
official_review

<thread_object>
creator: reviewer_3mnz

<content>
summary:
This paper addresses the systematic bias in mini-batch quadratic approximations used for second-order optimization and uncertainty quantification in deep learning. Mini-batch methods distort the quadratic shape, making uncertainty estimates unreliable. The authors provide a theoretical explanation, discuss the impact on optimization, and propose debiasing strategies to improve accuracy.

strengths:
This is an interesting theoretical investigation collaborated with numerical justifications. It teaches me something that I would expect with a nice and convincing narrative. The overall clarity and language of this paper is good. The insights are well-delivered. The theoretical analyses are mostly backed up by numerical experiments.

weaknesses:
I did not identify any obvious weaknesses.

I did not check the proofs.

questions:
How well do the proposed debiasing strategies scale for large foundation models with millions or billions of parameters?

Can debiased mini-batch quadratics improve uncertainty quantification in foundation models, especially in tasks requiring precise uncertainty, like language modeling or image synthesis?

Can the theoretical insights and debiasing techniques be effectively adapted for transformer-based foundation models?

Could reduced bias in mini-batch quadratics enhance transfer learning by improving local loss landscape approximations during fine-tuning?

How might the debiasing methods apply to multi-modal foundation models, where different data modalities add complexity to curvature estimation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730717632580

thread_type:
official_review

<thread_object>
creator: reviewer_swr8

<content>
summary:
This paper investigates biases in mini-batch quadratic approximations, for second-order optimization in deep learning. Authors show that when computing quadratic on mini-batches rather than the full dataset, it introduces curvature distortions, leading to overly narrow approximations. The bias, assigned to “regression to the mean”, affects Newton steps in optimizers and Laplace approx for uncertainty quantification. They provide a “two-batch” strategy to decouple direction selection from magnitude estimation, to align mini batch approx more closely with full-batch behaviour with minimal computational overhead. They demonstrate the improve stability of their strategy on CG and K-FAC experiments.

strengths:
Authors study a well-known but often overlooked issue in mini-batch approx within second-order optimisations and uncertainty quantification for deep learning: the bias in curvature estimates due to mini-batch sampling. While the issue is not new, the approach of a two-batch debasing strategy is an effective and novel (to my knowledge) way to address it and improve the accuracy of the approximates without suffering too much of a computational overhead. This strategy could generalise well across various settings in ML where the full-batch quadratics are intractable. The paper is supported by nice empirical experiments on CG and K-FAC demonstrating the effectiveness of their solution. The authors provide a theoretical intuition for the bias arising from that “regression to the mean” adding an interesting insight.  The experiments are systematically presented and show clear performance difference between single-batch and debasing method . While very technical, authors make an effort to make it as digestible as possible, with the inclusion of figures to visualise the bias and the impact of their debasing strategy.

weaknesses:
The experimental scope is pretty limited, with only ablations on CNN architecture and CIFAR-10/100 datasets. Testing on a wider variety of architectures, such as transformers or larger datasets, would strengthen a lot the paper’s claims on general applicability. Another limitations lies in the lack of ablations to investigate the sensitivity of the two-batch approach to mini-batch size, as well as the composition of mini-batches used for debasing. Without these, it’s unclear how the two-batch batch method would perform under different resource constraints or mini-batch sampling strategies. Last, the paper discusses the computational trade-offs, but does not provide sufficient detail to me on actual runtime and memory costs for different debasing configuration, leaving with an incomplete understanding of the true method’s scalability.

questions:
— How do you anticipate the method to scale or adapt with different architecture and larger dataset, which have different curvature profiles and optimisation challenges?
— Could you clarify the criteria for selecting mini-batches in the two-batch debasing process? Is there any benefit to use overlapping vs. non-overlapping mini-batches, or would a mini-batch with a specific loss profile provide better results? 
— Can you provide more details on runtime/memory costs associated with the two-batch approach? While authors claim it roughly doubles, it would be helpful to understand how this really scales with batch size, and whether costs remain manageable as model size and dataset complexity grow
</content>
</thread_object>

</thread>
<thread>
creation date:
1733296565587

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

Thank you so much for your constructive feedback and thoughtful questions! 

We are delighted that three of the four Reviewers (3MNZ, 9wSH and SWR8) clearly
recommend accepting our paper with a score of 8. We hope that our thorough rebuttal for
Reviewer ddR4 helped resolve the misunderstandings and addressed all concerns
satisfactorily.

In response to your suggestions, we have made the following updates to our paper (the
references we provide below refer to this updated version):
- We extended our experiment on debiased CG by two additional mini-batch sizes to
  investigate its sensitivity with respect to mini-batch size (see Figure 13). While the
  single-batch runs quickly diverge, the debiased approaches maintain stability across
  all mini-batch sizes.
- We have also added a runtime analysis to the manuscript (see Figure 18). It shows that
  the runtime overhead of the debiased Laplace approach is negligible compared to the
  single-batch version. 
- We have conducted an additional experiment on the ViT Little architecture on ImageNet
  data (see Figure 17). The results show that (i) the curvature biases also exist for
  these types of models, (ii) our proposed KFAC Laplace approach scales to larger models
  and datasets, and (iii) that the debiased approach effectively mitigates the biases.

These updates, we believe, have further strengthened the paper, addressing key points
raised during the review process. We appreciate the constructive and positive
evaluations and feel encouraged by the positive feedback!
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates systematic biases in mini-batch quadratic approximations used for second-order optimization and uncertainty quantification in deep learning. The authors demonstrate that using mini-batches to estimate curvature leads to distortions in the quadratic shape, particularly overestimating top eigenvalues. They provide theoretical explanations for this phenomenon and propose practical debiasing strategies using a two-batch approach. The paper shows how these biases affect applications like conjugate gradient optimization and Laplace approximation for uncertainty quantification, and demonstrates that their debiasing methods can significantly improve performance in these areas.

The paper's strengths lie in its identification of an important issue in deep learning optimization and uncertainty quantification, providing both theoretical insights and empirical evidence for the observed biases. The proposed debiasing strategies are practical and computationally efficient, showing clear improvements in concrete applications. The paper is well-presented with clear explanations and visualizations.

Initially, the experimental scope was somewhat limited, focusing mainly on CNN architectures and CIFAR datasets. However, the authors addressed this concern in their rebuttal by expanding their experiments to include larger models and datasets.

The paper addresses a significant and previously underexplored issue in deep learning, with potential impact across multiple areas including optimization and uncertainty quantification. The proposed debiasing strategies are novel, practical, and show clear improvements over existing methods. The majority of reviewers strongly support acceptance, and the authors have adequately addressed reviewer concerns and expanded their experiments in the rebuttal.