Paper: Implicit regularization of multi-task learning and finetuning in overparameterized neural networks (ICLR-2024)

Reviewer comment (Novelty):
The main contributions of Sections 4.1 and 4.2 are slight variations of known results. In particular, the characterization for the finetuning setting (Section 4.1) is a slight variation of Theorem 1 in [1], and the characterization for multi-task learning (Section 4.2) is a slight variation of Theorem 2 in [2]. Thus, beyond existing knowledge, the paper mostly provides interpretations of the sparsity-inducing implicit regularization for diagonal linear neural networks. Such interpretations can be valuable, yet the setting considered is quite limited.

Author response:
—  **Regarding our theoretical contribution**: we would like to emphasize that we do not believe the primary contribution of our paper is in proving new theorems.  As the reviewer notes, our main theoretical claims are applications/interpretations of existing results in the setting of pretraining+finetuning and multi-task learning; we have endeavored to make this more clear in the revised manuscript.  However, we believe that applying this theoretical framework to PT+FT and MTL is novel, and the insights gained are non-obvious and will be of interest to the community.  We would like to highlight some of these insights here.  (1) The intermediate regime between rich and lazy learning that we highlight (i.e. “nested feature selection”) is not emphasized in prior work, but we show is key to understanding the inductive biases of PT+FT.  (2) Our thorough suite of teacher-student experiments demonstrates consequences of L1/L2-like norm biases which may not be obvious — for instance, the fact that models can display a “rich” bias toward sparsity in new features learned during finetuning while maintaining a “lazy” lack of bias toward sparsity in use of features inherited from pretraining. (3) One of our key findings, not addressed in this review, is not captured at all by the linear theory, as it arises only in the nonlinear case due to the possibility of tasks involving correlated-but-not-identical nonlinear features.

Reviewer comment (Theory):
The analysis disregards the role of the auxiliary task data and objective, which are crucial in practice. The considered multi-task setting is equivalent to a single task with multiple outputs, meaning the analysis inherently cannot provide insight into how the implicit regularization of training with auxiliary tasks differs from single-task training. Due to the substantial overlap between the results of Section 4 and prior work, a promising direction forward would be to investigate how the effects of auxiliary task data and objectives come into play.

Author response:
– **Regarding investigating the role of auxiliary task data / objective:** We agree this is an interesting and worthwhile direction for future work and are excited to pursue it.  However, we also think that our existing contributions constitute a paper’s worth of results on their own.

Reviewer comment (Theory):
The comparison between the implicit regularization of finetuning and multi-task learning is not an “apples to apples” comparison. For finetuning, the implicit regularizer is derived for gradient flow with the square loss, whereas for multi-task learning the L2 representation cost is considered, although it is not clear whether gradient descent/flow indeed leads to minimal L2 norm solutions in this case.

Author response:
– **Regarding the PT+FT / MTL comparison:** we agree that the theoretical results are not directly comparable and have further highlighted this in the revised manuscript.  However, we believe our teacher-student experiments demonstrate that the norms we derive for PT+FT (under gradient flow) and MTL (under L2 minimization) are useful descriptions in that they make accurate predictions about the relative performance of the two methods across a diverse suite of task setups.

Reviewer comment (Experiments):
The empirical demonstration of the proposed heuristic (scaling down weights) is insufficient. The evaluation includes only a single model (ResNet18) and a single dataset (variant of CIFAR100) with an oddly small number of training samples. Although the paper is mostly theoretical, this does not meet the extent of evaluation expected when proposing a practical guideline.

Author response:
– **Additional results for deep networks / realistic tasks:** We agree with the reviewer’s requests for more convincing empirical results in realistic networks/tasks and have done our best to meet them. We scaled up our CIFAR-100 experiment to use more samples.  See the “response to all reviewers” for a description of additional validation of the proposed rescaling method with different architectures and tasks (Appendix F).

Reviewer comment (Evaluation):
Claims of an implicit regularization towards sparsity are measured in experiments only through sample complexity, observed when plotting the loss against the number of samples. It would be more informative to measure the sparsity levels directly.

Author response:
– **Additional evidence for sparsity-incentivizing implicit regularization:** we agree that our paper would benefit from a more direct characterization of sparsity in learned solutions, particularly in the nonlinear networks where the applicability of the theory is less assured.  We have provided additional analyses that corroborate this description in the revised manuscript, provided in Appendix E.  See the “response to all reviewers” for a description of the key results of this analysis.  

– That said, we agree with the reviewer that a direct analysis of the sparsity of learned solutions is helpful and complements our existing analyses.  We have integrated the highlights of the analysis of sparsity in solutions in the nonlinear case (Appendix E)  into the main text (Fig. 2).  Moreover, as requested, we have conducted a similar analysis in the diagonal linear case, the highlights of which are now included in Fig. 1 (and more details on the linear case are now provided in Appendix E.1).  These analyses directly illustrate the bias towards sparsity in learned task-specific features (Fig. 1c, Appendix E fig. 5e for linear case, Fig. 2c, Appendix E fig. 6c+d for nonlinear case), and sparsity in features learned in the “nested feature selection” regime (Fig. 1h for linear case, Fig. 2h for nonlinear case) as a function of the weight rescaling factor.

Reviewer comment (Writing):
The presentation severely lacks clarity. Claims are made without laying out their exact statements and under which assumptions they are proven. Specifically, in Section 4.1 it is not specified what the minimized objective is, what the optimization algorithm is, nor what the auxiliary weights stand for and how they are obtained. The appendix does not provide much clarity either, though after close examination it appears to prove a slight variation of Theorem 1 in Azulay et al. 2021 for the implicit regularization of diagonal linear networks optimized with gradient flow from a certain initialization.

Author response:
– **Regarding clarity:** we thank the reviewer for the helpful suggestions for improved clarity in the presentation of theoretical results and notation, which we have tried to address in the revised manuscript (see red text).  We have also changed “feature-learning” to “feature selection” where appropriate (i.e. in the linear setting), as suggested.  

– We have updated the abstract, introduction, and related work section to make more clear the close relationship between the theoretical statements in our work and those in prior work on diagonal linear networks.  

– We have rewritten the main theoretical claims more clearly and formally (see Corollary 1 and Corollary 2 in the revised draft)

Reviewer comment (Experiments):
The main focus of paper is on MSE loss, I think adding the CE loss focus could improve the paper.

Author response:
We agree with the reviewer that an analysis of the cross-entropy loss case would be of interest and this is an active follow-up direction for us.  We provided some comments on the cross-entropy case in Appendix C of the original draft.  Unfortunately, we are unlikely to be able to provide a comprehensive analysis of the cross-entropy loss case during the discussion period – in particular, it requires a different theoretical analysis of PT+FT, which we leave to future work.

Reviewer comment (Experiments):
Sample Size and Repetition: The paper would benefit from addressing the potential limitation regarding sample size and repetition in the experiments. It is essential to clarify whether the results, as shown in Figure 3(f), are robust across different samples. The authors should consider repeating the experiments with various samples multiple times and reporting the mean performance along with confidence intervals. As it stands, it appears that the experiments were conducted only once, which may raise concerns about the reliability and generalizability of the findings. A more comprehensive analysis of the impact of different samples on the results would strengthen the paper's credibility and reliability.

Author response:
– **Error bars / number of repetitions run:** each experimental result shown is an average over many trials, using randomly sampled datasets and/or weight initializations, as appropriate. We meant to illustrate the trial-to-trial variability in our initial draft with error bars; however, this was not clear in our initial draft, as the error bars were small and the thickness of the lines obscured them.  We have revised the figures to make the confidence intervals more clear, and where they are not visible, it is because they are smaller than the thickness of the lines/points in the plots – we hope that the low variability addresses the reviewer’s concerns that the results may not be robust.  Moreover, the updated figures make it more clear which aspects of the plots are actual experimental results vs. interpolating lines, as suggested by the reviewer.

Reviewer comment (Evaluation):
Lack of Visualization and Evaluation Metrics: The paper could benefit from a more comprehensive presentation of results. It lacks visualization techniques to visually illustrate these biases. Visual representations can provide a clearer and more intuitive understanding of the observed phenomena.

Author response:
– **Visualizations / evaluations:** we have provided several new analyses and figures in the revised manuscript that help illustrate the inductive biases we describe in more concrete terms.  Please see our “response to all reviewers” for a description of these analyses.

Reviewer comment (Experiments):
Regarding the experimental repetition, it's not clear why the author relies on results from a single experiment. Is this choice due to concerns about potential bias in the sampled data, or are there other reasons for not conducting multiple experiments to ensure robustness and reliability of the findings?

Author response:
– **Error bars / number of repetitions run:** each experimental result shown is an average over many trials, using randomly sampled datasets and/or weight initializations, as appropriate. We meant to illustrate the trial-to-trial variability in our initial draft with error bars; however, this was not clear in our initial draft, as the error bars were small and the thickness of the lines obscured them.  We have revised the figures to make the confidence intervals more clear, and where they are not visible, it is because they are smaller than the thickness of the lines/points in the plots – we hope that the low variability addresses the reviewer’s concerns that the results may not be robust.  Moreover, the updated figures make it more clear which aspects of the plots are actual experimental results vs. interpolating lines, as suggested by the reviewer.

Reviewer comment (Presentation):
The figures corresponding to samples seem to lack clarity. The x-axis is labeled with "32 samples" and "256 samples," but it's not entirely clear which parts of the figures represent actual experimental results. Could the authors consider visually highlighting the data points with dots and connecting them with lines to provide a more explicit representation of the experimental results, making it easier for readers to interpret the findings?

Author response:
– **Error bars / number of repetitions run:** each experimental result shown is an average over many trials, using randomly sampled datasets and/or weight initializations, as appropriate. We meant to illustrate the trial-to-trial variability in our initial draft with error bars; however, this was not clear in our initial draft, as the error bars were small and the thickness of the lines obscured them.  We have revised the figures to make the confidence intervals more clear, and where they are not visible, it is because they are smaller than the thickness of the lines/points in the plots – we hope that the low variability addresses the reviewer’s concerns that the results may not be robust.  Moreover, the updated figures make it more clear which aspects of the plots are actual experimental results vs. interpolating lines, as suggested by the reviewer.