Implicit regularization of multi-task learning and finetuning in overparameterized neural networks

We thank the reviewer for the thoughtful criticisms of our work and suggestions for improvement. We have updated our manuscript in response to several of the reviewer’s points (major revisions for clarity are highlighted in red) and provide responses below. If the reviewer has further suggestions or comments we would be happy to address them during the discussion period.

— 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.

– 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.

– 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.

– 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).

– Additional analysis of nested feature selection regime in deep networks: Please also see our “response to all reviewers” for description of an additional analysis (Appendix G) that provides evidence that the weight rescaling method does indeed cause the ResNet-18 network to display a particular signature of “nested feature learning” behavior during finetuning (which is also observed in the appropriately scaled shallow teacher-student networks).

– 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.

– 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.

I have read the author's response and the other reviews carefully.

1. I acknowledge that, by the author's response, the theoretical analyses of Section 4 were not intended as main contributions. However, the paper is certainly not positioned accordingly. The abstract and introduction specify the description of solutions learned by two-layer diagonal linear networks as part of the main contributions. The fact that these are slight variations of known results is not mentioned in either places, nor in the related work section. For the finetuning implicit bias analysis this is only mentioned in the proof appendix, and for the multi-task setting it is only specified briefly in parenthesis. I find this to be misleading. The relation to prior work needs to be abundantly clear.

2. Given that the theoretical characterizations are not a main contribution, rather the insights and empirical analysis, the bar of empirical evidence becomes higher. The additional experiments in Appendix E are a good step forward, and should probably take a more central role given that they are the only direct support for the implicit bias towards sparsity claims.

   Furthermore, the experiments in Appendix E consider only ReLU networks. To better substantiate the described nested feature selection and sparsity bias in task-specific weights, it would be beneficial to consider diagonal linear networks as well, in which they can be evaluated directly.

3. Presentation: Although improved, the statement of the theoretical results (Section 4) is still not formal. I strongly recommend abiding to the standard practice of stating formal results in a Theorem/Proposition/Corollary environment from which it should be clear what exactly is proven and under which conditions.

Overall, I believe that the insights derived from existing theory, accompanied by empirical evidence, can be of interest and merit publication. However, in my opinion, the paper requires rather major modifications that will clarify its main contributions with respect to prior works, and improve its focus on the novel contributions (e.g. substantiating claims of the nested feature selection and task-specific weight sparsity). Unfortunately, I do not believe that such a revision is feasible as part of the author-reviewer discussion period.

We thank the reviewer for participating in the discussion. We sincerely appreciate the reviewer’s suggestions and attention to detail; we have updated our manuscript to address these additional comments (sections that have been revised or added are highlighted in blue font) and believe it has made our paper stronger and clarified its contributions. At the same time, we wish to challenge some of the statements the reviewer has made about the work, which we believe overlook key aspects of the manuscript.

We believe that our work provides several insights into the subject of multi-task learning and finetuning that are entirely novel and are of interest to the community. As such, we argue that this work merits timely conference publication if the claims are technically sound and communicated accurately. We would greatly appreciate if the reviewer could reassess the revised manuscript, which we believe meets this bar.

Specific responses / comments addressed:

– 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)

– The reviewer suggests that the empirical evidence for our key claims (in particular the claim of task-specific sparsity bias for PT+FT and MTL and the nested feature selection regime for PT+FT) is inadequate. We respectfully disagree; the empirical evidence we provided in our initial submission is a direct demonstration of these key claims. In particular, the reviewer indicates that the sample complexity plots in Figures 1 and 2 (formerly Figures 2 and 3) are only indirect evidence for the implicit penalties we claim for PT+FT and MTL. However, the sample complexity results are in fact strong evidence, given our intentional selection of teacher-student experiments. For instance, the improved sample complexity for sparser ground-truth in the case of non-overlapping main and auxiliary task features (Fig 1d and 2b) demonstrate an inductive bias towards sparsity in task-specific features. Likewise, the improved sample complexity (or lack thereof) for sparser ground-truth features in the case in which main task features are a subset of auxiliary task features (Fig. 1e+f, 2e) demonstrates the presence (or absence) of the nested feature selection regime. We also would like to draw attention to Fig. 1c (present since the initial submission), which provides direct evidence for the implicit regularization claims in the diagonal linear case (which we have now moved to Appendix E, Fig. 5a to make room for the sparsity analysis requested).

– 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.

– We have integrated the analysis of appendix G, comparing the signatures of the nested feature selection regime between the teacher-student setting and the deep network experiments, into the main text, to reflect the greater emphasis on empirical results in the revised manuscript (Section 6, blue text, Fig. 3c).

We thank the reviewer for the positive evaluation of our work and suggestions for improvement.

– We have updated the CIFAR experiments so that in each trial, the choice of the 2 classes used for the “main task” is randomized. Moreover, we have scaled up the experiment to use all the samples available in the dataset. We have also provided further corroboration (Appendix F) of the usefulness of the post-pretraining rescaling method we proposed – see the “response to all reviewers” for details.

– 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.

We thank the reviewer for the positive evaluation of our work and suggestions for improvement. We have attempted to address all the points raised in the review in our revised version of the manuscript:

– 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.

– Number of samples: in our teacher-student, we use 1024 samples for the auxiliary task, and vary the number of main task samples (the number of main task samples is what is depicted on the x axes). Ultimately the choice of 1024 is arbitrary and was made to allow a large number of experiments to be run efficiently while still being large enough to allow the network to achieve very good test-set performance on the auxiliary task. We have run some experiments with more auxiliary task samples (and also more features in the ground-truth teacher network) which show that the key results of Fig. 3a are robust to scaling these parameters. We have added these results to Appendix D of the revised manuscript.

– 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.