Bias in Motion: Theoretical Insights into the Dynamics of Bias in SGD Training

We would like to thank the reviewer for their time reading and assessing our paper. We are happy that the reviewer found our paper well-written and, in the following, we would like to address their main comments touching an important point, i.e. the robustness of our result in more complex settings.

We extended the experiment on MNIST by now training on a network with no ReLU non-linearities –i.e. linear activations– and visualise the result in Figure 2 of the rebuttal PDF. Compared with Figure 5(a) of the paper, we note that removing the non-linearities leads to similar qualitative behaviour and we still observe a crossing. Giving another confirmation of our theory.

The reviewer also asked about situations where the crossing does not occur. We can find an instance in Fig.1a of the main paper and, more broadly, looking at the phase diagram of Fig.3b we can notice that all the dark red and dark blue regions of the generative model’s parameters are instances where the crossing does not occur. Our theory shows that this is because the preferred subpopulation of our classifier stays the same across the different phases of dynamics described in the paper.

Connected to this point, we would like to refer to answer to reviewer CSav where we also see that the crossing is not necessarily happening –Fig.1 and Fig.3 of the rebuttal PDF– and it occurs when the data distributions of the groups show characteristics identified by our theory –that depends on the variance and relative representation of the groups as explained in section 4.2 of the paper. We kindly ask the reviewer to refer to CSav rebuttal for further details.

We hope we were able to address your concern and once again thank you for your valuable time and remarks.

We hope that the clarifications and additional experiments have addressed the reviewer’s concerns. If there are any remaining questions or points that need further discussion, we would be grateful for any additional feedback.

We would like to thank the reviewer for their general appreciation of our work and for their comment.

Impact of the value of $m$

Thank you for your insightful question. We plan to incorporate the following points into the revised manuscript:

• A higher value of $m$ would introduce additional order parameters for the student-teacher overlaps and student-shift overlaps corresponding to the extra clusters. Despite this added complexity, we anticipate a similar system of coupled ODEs to what is described in Appendix D.1, and we are confident that the results would remain analytically tractable.

• We expect that the form of the equations in Theorem 3.4 would remain fundamentally unchanged, with the same type of timescales and dynamics discussed in the paper. The constants, such as $\Delta^{mix}$, would be adjusted to account for the weighted average across a greater number of clusters. Consequently, the same factors to drive bias –degree of spurious correlation, variance, relative representation– and the three-phase behaviour described in the paper would still emerge. We are optimistic that the key takeaways from Section 4 will extend to the $m>2$ setting.

• While we believe most conclusions will hold for higher values or $m$, , this is an interesting direction we are currently investigating. In this work, we focused on $m=2$, which already shows a rich and sometimes counter-intuitive phenomenology. However, the $m>2$ scenario introduces additional complexities, such as varying levels of similarity between groups, which merit further exploration.

We would welcome any further suggestions or discussion to address this concern more comprehensively.

We thank the reviewer for their appreciation of our analysis and their comments. We feel that these are important comments to improve the clarity of our work and we thank the reviewer for pointing them out, we ran additional experiments (discussed below) to address them.

Limitations of the model

Our theoretical analysis is limited to purely linear classifiers to maintain analytical tractability. To address questions about the robustness of the results and whether this may be due to the low capacity of the classifier, we relied on numerical evaluation on deeper networks. We provided some experiments in section 5 and appendix F where we used 2 layer MLPs and finetuned a ResNet-18 on real data sets.

Taking your remark and the related remark of reviewer JTC4 into account, we have now included additional results that demonstrate the behaviour predicted by our theory on deeper and wider MLPs (see the submitted pdf in our global response):

• We have expanded the experiment in section F.2 as follows: We consider a data distribution with $d=100$, $\rho=0.7$, $\Delta_+=0.1$, $\Delta_-=1$, $T_{\pm}=0.9$, $v=4$, $\alpha_+ = 0.471$, $\alpha_- = -0.188$. We train a 2-layer MLP with ReLU activation in the hidden layer(s) and sigmoidal activation at the output with 128 neurons in the hidden layer with online SGD with a learning rate of 0.01 using BCE loss. We refer to this as the ‘standard configuration’ useful for further comparisons. We visualise the loss curves and test error rates on the positive (blue) and negative (red) subpopulations along with standard deviation across 10 seeds in Figure 4 (left panel).

• We further train MLPs of increasing depth (2,3,4,5 layers) with 128 hidden neurons in each layer and visualise the loss and error curves in Figure 6. The figures show the three phase behaviour predicted by our theory - the positive subpopulation (in blue) is learnt faster initially since it exhibits stronger spurious correlation with the shift vector; the negative subpopulation (in red) is learnt faster subsequently due to higher variance and finally the positive subpopulation is asymptotically prioritised once more due to higher relative representation. As such, we also observe the ‘double-crossing’ phenomena discussed in the paper. We stress that this result and the next are obtained by taking exactly the same parameters used in Fig.10 of the paper without exploring the phase space any further.

• We also train MLPs of 2 layers but increasing width (2, 16, 128 and 1024 hidden neurons) and visualise loss and errors curves in Figure 7. We note the presence of the same three phase behaviour across increasing parameterization. Thus, our empirics with over-parameterized students show qualitatively similar behaviour. We hope this further empirical validation helps show the validity of our theory in settings closer to modern regimes.

Hyperparameters used in experiments

Apologise for not including the learning rate in the empirical experiments and will do so in the camera ready. Note that our synthetic model does however explicitly consider the effect of the learning rate. As shown in Fig.3b, discussed in sec. 4.2 and expanded upon in Appendix E.1 of the paper, the learning rate can act as a source of bias by advantaging the cluster with smaller variance.

Following your suggestion, we tested the ‘standard configuration’ with varying learning rate (Figure 5 of the included PDF). As predicted by the theory (see sec. 4.2 and Figure 3b of the paper), for high learning rates, the loss of the higher variance group blows up (left panel of Figure 5 of the included PDF). For smaller learning rates, the dynamics remain the same, it is simply the overall speed that changes as we can see in the included figure (for smaller learning rates, it takes longer for the crossings to occur).

Weight decay is quite an interesting technique to consider and our analysis offers direct predictions based on the theory. As discussed in Appendix E, smaller student magnitude advantages the higher variance group and indeed when we empirically include a weight decay penalty of 0.1 (Figure 4 right panel of the included PDF), we observe that asymptotically the higher variance cluster is now preferred since the penalty encourages smaller weights.

As for Adam, we see very similar behaviour (Figure 4 centre panel of the included PDF) although we had to use a slower learning rate of 0.001 to prevent losses from blowing up since Adam leads to faster optimization manifesting in a higher effective learning rate than SGD. We fully expect all our key findings to translate.

Suggested additional experiment

This is a great suggestion. We provide results (see submitted PDF) from additional experiments to show how modifying the real data changes behaviour exactly as predicted by the theory. Figure 1 shows an extension of the CIFAR10 experiment discussed in the paper where we now vary the relative representation of the groups (darker colours reflect greater imbalance) and we see that as predicted by our theory, crossings only occur when the higher variance group (red) has lower relative representation. Notice that the dynamics of the initial phase seem almost unaffected by the relative representation and are determined only by the variance as predicted by the theory.

We also extend the MNIST experiment discussed in the paper and show how increasing variance of the second group (red), controlled through image brightness, leads to faster learning in the second phase in Figure 3.

In conclusion, we hope these extensions provide strong evidence showing how the factors identified in the theory affect real data experiments exactly as predicted.

We also thank the reviewer for pointing out additional references that we will include in the further work. If there are any remaining concerns or if further clarification is needed, we would be happy to engage in additional discussion to address them.

We hope that the clarifications and additional experiments have addressed the reviewer’s concerns. If there are any remaining questions or points that need further discussion, we would be grateful for any additional feedback.

We thank the reviewer for their appreciation of our characterisation of the non-asymptotics training dynamics and their constructive comments.

Regarding the Use of Linear Classifiers:

As the reviewer points out, our theoretical analysis is limited to linear classifiers to ensure analytical tractability and to obtain, for the first time to our knowledge, closed-form expressions for the generalisation error curves and order parameters.

However, we have also tested the robustness of our theory on more complex architectures. In Section 5 and Appendix F, we provided experiments using 2-layer MLPs and fine-tuned a ResNet-18 on real data sets. Taking your remark and the related remark of reviewer CSav into account, we have now included additional results that demonstrate the behaviour predicted by our theory on deeper and wider MLPs (see the submitted pdf in our global response):

• We have expanded the experiment in section F.2 as follows: We consider a data distribution with $d=100$, $\rho=0.7$, $\Delta_+=0.1$, $\Delta_-=1$, $T_{\pm}=0.9$, $v=4$, $\alpha_+ = 0.471$, $\alpha_- = -0.188$. We train a 2-layer MLP with ReLU activation in the hidden layer(s) and sigmoidal activation at the output with 128 neurons in the hidden layer with online SGD with a learning rate of 0.01 using BCE loss. We refer to this as the ‘standard configuration’ useful for further comparisons. We visualise the loss curves and test error rates on the positive (blue) and negative (red) subpopulations along with standard deviation across 10 seeds in Figure 4 (left panel).

• We further train MLPs of increasing depth (2,3,4,5 layers) with 128 hidden neurons in each layer and visualise the loss and error curves in Figure 6. The figures show the three phase behaviour predicted by our theory - the positive subpopulation (in blue) is learnt faster initially since it exhibits stronger spurious correlation with the shift vector; the negative subpopulation (in red) is learnt faster subsequently due to higher variance and finally the positive subpopulation is asymptotically prioritised once more due to higher relative representation. As such, we also observe the ‘double-crossing’ phenomena discussed in the paper. We stress that this result and the next are obtained by taking exactly the same parameters used in Fig.10 of the paper without exploring the phase space any further.

• We also train MLPs of 2 layers but increasing width (2, 16, 128 and 1024 hidden neurons) and visualise loss and errors curves in Figure 7. We note the presence of the same three phase behaviour across increasing parameterization. Thus, our empirics with over-parameterized students show qualitatively similar behaviour. We hope this further empirical validation helps show the validity of our theory in settings closer to modern regimes.

Regarding Online SGD with Empirical Loss:

We would like to clarify that our theory indeed considers with online SGD using the empirical loss. Specifically, at each time step, we sample a new supervised example, calculate the loss for that example, and then update the model parameters using gradient descent. All our empirical experiments also follow this training procedure.

While this approach is standard in theoretical analyses due to its analytical convenience, we acknowledge the importance of capturing the dynamics of multi-pass/batch SGD, as we discussed in the 2nd paragraph of our conclusion. In future work, we aim to address this challenge by leveraging results from dynamical mean field theory (e.g., [1]).

On tractability beyond m=2:

Our current work focuses on $m=2$ to facilitate the interpretation of the dynamics and to identify key takeaways. However, we are confident that the equations remain analytically tractable for higher $m$ values. Adding more subpopulations would involve taking weighted averages by representation in all expectations, similar to the existing analysis. We anticipate that this lead to a system of coupled ODEs similar to what we presented in Appendix D.1.

If there are any remaining concerns or if further clarification is needed, we would be happy to engage in additional discussion to address them.

References:

[1] F. Mignacco, F. Krzakala, P. Urbani, L. Zdeborová. Dynamical mean-field theory for stochastic gradient in Gaussian mixture classification. NeuriPS 2020.

We thank the reviewer for their insightful comments and careful consideration of our paper.

On the practical significance of our work

Thank you for highlighting this important point. We plan to address the practical implications of our findings by adding the following discussions to the revised manuscript:

• Caution Against Implicit Assumptions: Our results caution against the common implicit assumption in the literature on spurious correlation mitigation that bias remains the same throughout training (e.g., [1, 2]).

• Identifying Key Drivers of Bias: Connecting our theoretical insights to actionable practices, we identify variance–specifically image brightness–as a key factor influencing bias in neural networks, as discussed in our CIFAR 10 experiment (section 5). This finding enriches the literature on simplicity bias by providing a direct explanation for what might make a subpopulation appear ‘easier’ for a model to learn, thereby introducing bias.

• Expanding Bias Mitigation Techniques: Most bias mitigation strategies focus on under/oversampling to adjust the relative representation of groups. Our work suggests that e.g., by normalising image data to have similar brightness across subpopulations, practitioners can address an additional source of bias–group variance. This expands the toolkit available for developing bias-free bias-free classifiers, offering a more nuanced approach to bias mitigation.

We hope these additions will clarify the practical significance of our work, and we are open to further discussion to ensure your concerns are fully addressed.

References

[1] E. Z Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, C. Finn. Just Train Twice: Improving Group Robustness without Training Group Information. ICML 2021.

[2] Y. Yang, E. Gan, G. K. Dziugaite, B. Mirzasoleiman. Identifying Spurious Biases Early in Training through the Lens of Simplicity Bias. arXiv: 2305.19761.