Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse

Summary

This work investigates the implicit bias of both multi-task learning (MTL) and fine tuning pre-trained models (PT+FT) in diagonal linear networks and shallow ReLU neural networks. The contributions are as follows: 1. The authors prove that both MTL and PT+FT bias the networks to *reuse features*. This holds for both diagonal linear networks and shallow ReLU networks (Corollary 1, Corollary 2 and Proposition 3). 2. The authors prove that during finetuning networks interpolate between a "lazy" and "rich" regime. (Proposition 4). 3. For PT+FT it is observed empirically that the choosen features are a sparse subset of the features learned during pretraining. 4. The paper provides empirical evidence that for ReLU networks only PT+FT benefits from features that are correlated between the auxilary and primary task. 5. Finally, they present a practical technique to improve PT+FT by weight rescaling after pretraining.

Strengths

I think the setting is interesting and very relevant to modern deep learning. This is also one of the few papers studying the inductive bias of multi-task networks. In particular: - The paper offers a number of results with many experiments on toy models. - The finding that MTL and PT+FT induces a regularizer that interpolates between $\ell_1$ and $\ell_2$ regularization is interesting and novel. - The authors go beyond theory and show how their finding can be exploited in practical architectures through a simple weight rescaling technique.

Weaknesses

- Corollary 1 and its extension for ReLU neural networks in Appendix A.1 is already well known in the literature [2] (see discuss after Remark 2) and [3] (Equation 4). These and the references therein should at least be mentioned. - The works of [1] and [2] are both very related to the phenomenon uncovered here. In particular the notion of "neuron sharing" [2] seems to exactly correspond to the "feature reuse" phenomemon described in this work. - It's not clear to me why the paper focuses on linear diagonal networks at all. These are not used in practice and are only studied for theoretical reasons. Since the theoretical contributions seem to be corollaries of previous results I think the paper would benefit from just focusing on shallow ReLU networks. - The regularizer used during finetuning in Proposition 3 isn't well motivated. Minimizing the $\ell_2$ norm of the weight changes during fine tuning is not weight decay and seems a bit ad-hoc. It would be great to provide some discussion on why you choose to analyze this regularizer or whether this is done in practice. - The notation is very loaded and makes the paper difficult to parse. For example $\vec{w}^{(1),aux}_{h}$. - The figures are also very small and hard to read. It would be better to fill in all the whitespace between the different plots and include a legend for the different settings instead of a colorbar. **Minor things** - Lines 51 and 54 reuse the word "Finally". - Line 82: Should the inputs to the model $x \in \mathbb{R}^{D}$? - Line 113 paraeter --> parameter - Line 135: Typo $\mathbb{R}^{d} \in \mathbb{R}$ - Line 251: yon --> on - In Section 3 $D$ and $d$ are both used interchangeably for the input dimension of the networks. [1] Collins, Liam, et al. "Provable multi-task representation learning by two-layer relu neural networks." arXiv preprint arXiv:2307.06887 (2023). [2] Shenouda, Joseph, et al. "Variation Spaces for Multi-Output Neural Networks: Insights on Multi-Task Learning and Network Compression." arXiv preprint arXiv:2305.16534 (2023). [3] Yang, Liu, et al. "A better way to decay: Proximal gradient training algorithms for neural nets." OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop). 2022.

Questions

- Where is the proof for Corollary 2? - When finetuning are the weights of both layers trained or only the second layer weights? - Where are the results from the experiments discussed in Section 3.3? - In Fig 2a, what do the labels on the colorbars mean "# Active dims" and "Units"? Is this the number of features in the teacher model? - Line 176 what do you mean by "overlapping dimensions"? You mean overlapping features? - On line 178 what is LP in PT+FT(LP)?

Rating

5

Confidence

3

Soundness

3

Presentation

1

Contribution

2

Limitations

Limitations have been mentioned in the weaknesses section.

Summary

Abstract The goals of the paper are very clear from the abstract, as are the results. Introduction Lines 23-42: The authors do a great job of summarizing the applications of MTL (using it here loosely to capture MTL and PF+FT) while recognizing that very little work has been done to explore exactly why it works as well as it does in data-limited settings. We have strong intuitions about its regularization effects with minimal work backing our intuitions Related Work Lines 60-68: Noting the extent of prior work analyzing regularization effects in single-task modeling is hopeful in showing the impact of the author’s work in adding to similar work for MTL. Lines 85-89: I would do a better job of defining “large initialization” vs “small initialization” for the reader without requiring a reference check by the reader. One could assume that large = overparameterized and small = underparameterized versus large vs small magnitudes for the same number of parameters at initialization. Implicit and explicit regularization penalties for MTL and PT+FT Lines 94-102: I think most readers would take the authors at their heuristic given what they note in Lines 101-102, that explicit weight decay use in practice yields a permissible heuristic. However, the references, especially to [12] enhance their argument. Lines 103-111: The notation in Corollary 1 is a bit confusing. I am not sure why beta vector aux has the aux not in subscript when defining vector beta 2. I also do not know why we introduce vector beta 2 other than just showing that we are indexing the outputs since we continue to stick with vector beta aux. I see that in 3 we now subscript the dimension but I think the notation can be more clean. Line 112: Typo: “A ReLU networks” Lines 112-115: I have not seen a ReLU network defined this way and from the authors’ text the notation is again confusing. The authors state they are defining a ReLU network with O outputs but (4) that follows defines only a single final output from the network. What follows then makes us think that the “outputs” are the units in a single hidden layer? The notation is unnecessarily confusing for a basic ReLU network and something more conventional would help especially in the context of MTL. Also, the L1,2 norm is referenced several times already but not defined. There is another typo on line 113 “paraeter”. Lines 117-124: Why are the DLN network weights initialized with a constant magnitude? Lines 128-129: We are not told how we get to (6) and not appendix section is referenced Line 131: Typo Line 134: Herein we find more confusion because now (4) is referenced as an equation for a ReLU network with a single output, whereas this was not how it was defined in the preceding text before Equation (4). Also, weights within the set of R d within the set of R? Line 135: Is vector gamma simply a uniform vector (constant magnitude) of the same scalar similar to how the DLN was defined or is the vector gamma different values at re-initialization (not uniform)? Lines 137-139: It is not until Line 139 that theta and m are well defined despite numerous equations/references to them earlier in the manuscript. Figure 2: I would define STL as single task learning for the less familiar reader Lines 144-151: The experiment is explained well and easy to follow. Lines 152-155: More explanation for the ReLU would help. By “sparse” number of (hidden?) units we mean a low dimensional hidden layer? Lines 158-165: Everything is clear save how we arrived at c. Appendix reference? Lines 158-182: I like the elegance of the intuition to simplify the experiments but just having no feature overlap in aux and main tasks in small pretrained feature case and full feature overlap in aux and main in large pretrained feature case. Lines 183-194: Again, elegance in defining the feature sharing to explicitly study what was of interest (simultaneous sparsity and feature sharing). Lines 206-215: The definitions are clear and the intermediate regime is set up well by the authors. Lines 223-250: Again, the elegance of the experiment set-ups to reinforce the nested feature selection regime. Lines 256-257: I would like to hear more about this from the authors in the Conclusion. It seems like the authors are saying that there seems to be some “critical” (used loosely here) rescaling magnitude between aux aux and main tasks in MTL over which the network is unable to enter the nested feature selection regime. I would have liked to see the authors stress this to see what they found re: network tolerance to different rescaling magnitudes. Hopefully this can be a future work. Lines 262-264: I am not sure this is the case. MTL can certainly learn correlated features in shallow layers where one task’s use of the common features are simply rescaled later in the network relative to the features learned for an auxiliary task in isolation. Unless the authors are speaking more strictly of actual soft parameter sharing here, in which case, yes, separate parameters are being learned but kept similar. I think this is confirmed in the low sample MTL results (where we know MTL works best when it can). It does not work as well as PT+FT but still beats STL. It is a marginal beat. Lines 312-315: Very interesting finding on a large dataset and more complex network. Summary: The authors explore the inductive biases behind the success of MTL and PT+FT, a very under-studies phenomenon. Many interesting discoveries were made, most interestingly in their newly names "nested feature selection regime". This is important and timely work. I would like some of the notation cleaned-up for better reach to a wide audience.

Strengths

The theoretical justifications for the set-up of the experiments are excellent. The experiments that follow are also well-designed to allow the authors to probe their implicit regularization penalties. The authors kept probing until arriving on their nested feature regime. I liked the trajectory from theory to synthetic data/systems experiments to findings to larger/classic datasets and networks.

Weaknesses

I think some of the notation needs to be cleaned-up to have this important work reach as broad an audience as possible. Much is left to the appendix and when it is the references are usually there, but I would make sure they are always there for the reader to reference from the main body of the paper.

Questions

Why are the DLN network weights initialized with a constant magnitude? The authors mention how entrance into the nested feature regime is sensitive to the correlation between aux and main task features. Some more thoughts in the Conclusions about further theoretical thresholding of this or future experiments would be of interest to me as the authors seemed to have good theoretical foundation for all other experiments in the paper. Beyond me being pedantic, is this entire paper working on multi-output learning not multi-task learning? Traditionally when we talk about MTL we have two tasks with a different loss function, different goals, different processes, different dimensionality of input features (image classification as task 1 and image object segmentation as task 2). From what I can tell, the auxiliary and main tasks in the synthetic experiments are different instantiations of the same process, with carefully crafted differences in features/coefficients/etc in order to generate different outputs, then described as different tasks. I would describe this as multi-output learning. All of the findings still hold of course, but I would make the distinction more clear.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

NA

Summary

In this study, the authors explore the inductive biases associated with multi-task learning (MTL) and the sequential process of pretraining followed by finetuning (PT+FT) in neural networks. Specifically, they analyze the implicit regularization effects in diagonal linear networks and single-hidden-layer ReLU networks under these training paradigms. The findings reveal that both MTL and PT+FT promote feature reuse across tasks and favor sparsity in the feature set, establishing a conservation law that illustrates a tradeoff between these biases. Additionally, a unique "nested feature selection" pattern in PT+FT is identified, where a sparse subset of features learned during pretraining is selectively refined. This behavior contrasts with the broader feature learning strategies seen in MTL. Empirical validation using teacher-student models and experiments on deep networks trained for image classification confirms the theoretical insights. The authors also propose a practical enhancement involving weight rescaling post-pretraining, which their results suggest can optimize finetuning by encouraging the network to engage more effectively in the nested feature selection regime.

Strengths

1. The article provides an in-depth characterization of the inductive biases associated with two common training strategies—Multi-Task Learning (MTL) and Pretraining followed by Fine-Tuning (PT+FT)—in diagonal linear and ReLU networks. This detailed analysis is crucial for understanding the impacts of different training strategies. 2. By pushing networks into the nested feature-selection regime, the article proposes simple techniques to improve PT+FT performance, which have shown promising empirical results, adding practical application value.

Weaknesses

1. Although the findings are promising, the article notes that more work is needed to test these phenomena in more complex tasks and larger models. This suggests that the research's applicability and universality might be limited. 2. The article outlines promising avenues for extending the theoretical work, such as connecting derived penalties for ReLU networks to the dynamics of gradient descent, and extending the theory to the case of cross-entropy loss. This implies that the current theoretical foundation still requires further development. 3. The expression of the article needs further improvement. It is suggested that the author should introduce the theoretical settings of analyzing MTL and PT+FT in separate sections in detail, rather than the current form scattered across various sections.

Questions

See in Weaknesses.

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

N/A

Summary

This paper studies the implicit bias of gradient descent on the linear diagonal model and two-layer ReLU networks but instead of looking at the standard single output regression / classification setting, the authors study multiple outputs. In particular, the dataset X is associated with labels y and another dataset X_aux is associated with labels y_aux, which comprise an auxiliary task. With these, they study the multitask setting and the pretraining and fine-tuning setting. First, they apply several pre-existing results on implicit regularization to derive the implicit regularizer in each of these cases. They then analyze the qualitative implications of these results and run experiments to validate those qualitative analyses.

Strengths

* I really like the style of the paper! Take a theoretical result, extend it to a new setting, understand its implications, and verify them! The implications are not only verified on simplified models but also in more realistic models. * I view the audience of your paper as a much broader community than just the theory of deep learning community, as there are concrete and possibly algorithmic implications for your findings (i.e., it may be possible to design algorithms that actively incentivize certain kinds of feature learning in different settings via your results).

Weaknesses

A good paper overall. See some questions below.

Questions

* This paper seems somewhat related to the PT + FT setting. Mainly out of curiosity, as I think the settings are a little different, do you see any connections? Evron, I., Moroshko, E., Buzaglo, G., Khriesh, M., Marjieh, B., Srebro, N. & Soudry, D.. (2023). Continual Learning in Linear Classification on Separable Data. Proceedings of the 40th International Conference on Machine Learning. * These papers seem related to your feature sparsity and sharing experiments. Can you please discuss how they relate? In particular, do you see the emergence of a similar phenomenon to these works in the intermediate regime? Lee, S., Goldt, S. & Saxe, A.. (2021). Continual Learning in the Teacher-Student Setup: Impact of Task Similarity. Proceedings of the 38th International Conference on Machine Learning. Lee, S., Mannelli, S.S., Clopath, C., Goldt, S. & Saxe, A.. (2022). Maslow’s Hammer in Catastrophic Forgetting: Node Re-Use vs. Node Activation. Proceedings of the 39th International Conference on Machine Learning * While for someone familiar with the literature it is not so necessary, for accessibility to a more general audience, it would be helpful to get a little more discussion on what the settings and results are in [Dai et al], [Woodworth et al], [Azulay et al] and in particular why you need different results for the different settings. As mentioned above, your paper is one for a broad audience, one that is certainly much broader than just the implicit regularization community, who would be the ones most familiar with the subtleties in those above results.

Rating

8

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

Work is self-contained and aims to explain; the limitations of the theory in explaining certain cases are raised and addressed.

First Response

I appreciate the authors' follow-up responses. Much of my questions were asking for clarity on theoretical implications of the authors' work. Although it is not all addressed in the work here, the authors have many ideas that they are adding into the Discussion and Conclusions sections to drive further work by them and others. I think my point re: " I am not sure this is the case. MTL can certainly learn correlated features in shallow layers where one task’s use of the common features are simply rescaled later in the network relative to the features learned for an auxiliary task in isolation" came from the authors and myself referencing two different things. I was certainly referring to a NN with more than one layer whereas the authors were making their statement in the context of a NN with a single hidden layer. I do not think we have any disagreement in the more restricted case being described by the authors.

Thank you!

Thanks for engaging with my suggestions. I hope you found the literature suggestions valuable. I appreciate the discussion you have provided here. To reiterate, I really like your paper!