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

Authors: samuel  lippl
jack  lindsey
anthropic  san francisco
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:19:55

Abstract: Neural networks are often trained on multiple tasks, either simultaneously (multi-task learning, MTL) or sequentially (pretraining and subsequent finetuning, PT+FT). In particular, it is common practice to pretrain neural networks on a large auxiliary task before finetuning on a downstream task with fewer samples. Despite the prevalence of this approach, the inductive biases that arise from learning multiple tasks are poorly characterized. In this work, we address this gap. We describe novel implicit regularization penalties associated with MTL and PT+FT in diagonal linear networks and single-hidden-layer ReLU networks. These penalties indicate that MTL and PT+FT induce the network to reuse features in different ways. 1) Both MTL and PT+FT exhibit biases towards feature reuse between tasks, and towards sparsity in the set of learned features. We show a "conservation law" that implies a direct tradeoff between these two biases. 2) PT+FT exhibits a novel "nested feature selection" regime, not described by either the "lazy" or "rich" regimes identified in prior work, which biases it to *rely on a sparse subset* of the features learned during pretraining. This regime is much narrower for MTL. 3) PT+FT (but not MTL) in ReLU networks benefits from features that are correlated between the auxiliary and main task. We confirm these findings empirically with teacher-student models, and introduce a technique -- weight rescaling following pretraining -- that can elicit the nested feature selection regime. Finally, we validate our theory in deep neural networks trained on image classification. We find that weight rescaling improves performance when it causes models to display signatures of nested feature selection. Our results suggest that nested feature selection may be an important inductive bias for finetuning neural networks.

Main paper: 
 introduction: 
 Neural networks are often trained on multiple tasks, either simultaneously ("multi-task learning," henceforth MTL, see [1, 2] ) or sequentially ("pretraining" and subsequent "finetuning," henceforth PT+FT, see [3, 4] ). Empirically, models can transfer knowledge from auxiliary tasks to improve performance on tasks of interest. However, theoretical understanding of how auxiliary tasks influence learning and generalization is limited.Auxiliary tasks are especially useful when there is less data available for the target task. Modern "foundation models," trained on data-rich general-purpose auxiliary tasks (like next-word prediction or image generation) before adaptation to downstream tasks, are a timely example of this use case [5] .Auxiliary tasks are also commonly used in reinforcement learning, where performance feedback can be scarce [6] . Intuitively, auxiliary task learning biases the target task solution to use representations shaped by the auxiliary task. When the tasks share common structure, this influence may enable generalization from relatively few training samples for the task of interest. However, it can also have downsides, causing a model to inherit undesirable biases from auxiliary task learning [7, 8] .An influential strategy in the literature on the theory of single-task learning has been to characterize the implicit regularization conferred by the combination of network architecture and optimization algorithm [9] [10] [11] [12] [13] . Alternatively, some others characterize the effects of explicit parameter regularization (e.g. an ℓ 2 -penalty on weights) on the inductive bias of networks towards learning certain functions [14] [15] [16] . Compared to the single-task case, the inductive bias of MTL and PT+FT, whether obtained via explicit regularization or implicit regularization induced by optimization dynamics, is less well understood. Here we make progress on this question by studying inductive biases of MTL and PT+FT in two network architectures that have been extensively theoretically studied: diagonal linear networks, and densely connected networks with one hidden layer and a ReLU nonlinearity. We then demonstrate that our insights transfer to practically relevant scenarios by studying deep neural networks trained on image classification tasks.Our specific contributions are as follows:• We characterize regularization penalties associated with MTL or PT+FT in both diagonal linear networks and single hidden-layer ReLU networks (Section 3). • We find that both MTL and PT+FT are biased towards solutions that reuse features between tasks, and that rely on a sparse set of features (Section 4.2). • We then find that under suitable scalings, PT+FT exhibits a "nested feature selection" regime, distinct from previously characterized "rich" and "lazy" regimes, which biases finetuning to extract sparse subsets of the features learned during pretraining (Sections 4.3 and 4.4). • We find that PT+FT in ReLU networks can benefit from correlated (not just identical) features between the main and auxiliary task, but only when coefficients of features in the main task weights are of comparable magnitude (Section 4.5). • Finally, we study deep neural networks trained on natural image data (CIFAR-100 and ImageNet) (Section 5). Remarkably, we find that rescaling weights before finetuning improves accuracy in ResNets. Our analysis of the network representations suggests that this weight rescaling also results in the network relying on a low-dimensional subspace of its pretrained representation (i.e. exhibiting nested feature selection behavior). Intriguingly, Vision Transformers already exhibit signatures of nested feature selection without weight rescaling, and do not benefit from weight rescaling, suggesting that the nested feature selection regime is beneficial for finetuning performance.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We characterize the inductive bias associated with MTL and PT+FT by characterizing existing penalties and deriving a new penalty for PT+FT in ReLU networks, which we do in Section 3. We then test these inductive biases to understand how they manifest in practice in Section 4 and find that they are relevant for large-scale neural network architectures in Section 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 A variety of studies have characterized implicit regularization effects in deep learning. These include biases toward low-frequency functions [17] , stable minima in the loss landscape [18] , low-rank solutions [19] , and lower-order moments of the data distribution [20] . Chizat & Bach [13] show that when using cross-entropy loss, shallow (single hidden layer) networks are biased to minimize the F 1 norm, an infinite-dimensional analogue of the ℓ 1 norm over the space of possible hidden-layer features [see also 12, 14] . Other work has shown that implicit regularization for mean squared error loss in nonlinear networks cannot be exactly characterized as norm minimization [21] , though F 1 norm minimization is a precise description under certain assumptions on the inputs [22] .Compared to the body of work on inductive biases of single-task learning, theoretical treatments of MTL and PT+FT are more scarce. Some prior studies have characterized benefits of multi-task learning with a shared representational layer in terms of bounds on sample efficiency [23] [24] [25] .Others have characterized the learning dynamics of deep linear networks trained from nonrandom initializations, which can be applied to understand finetuning dynamics [26, 27] . Similarly, insights on the implicit regularization of gradient descent in linear models has been applied to better understand forgetting and generalization in a continual learning setup [28] [29] [30] [31] . However, while these works demonstrate an effect of pretrained initializations on learned solutions, the linear models they study do not capture the notion of feature learning we are interested in. Finally, teacher-student setups have been used to study the impact of task similarity on continual learning in deep neural networks [32, 33] . This methodology could also be applied to our setup (i.e. to investigate generalization on a finetuning task), and, similarly, our tools could be applied to continual learning setups.A few empirical studies have compared the performance of MTL vs. PT+FT in language tasks, with mixed results depending on the task studied [34, 35] . Others have observed that PT+FT outperforms PT + "linear probing" (training only the readout layer and keeping the previous layers frozen after pretraining), implying that finetuning benefits from the ability to learn task-specific features [36, 37] .Inductive biases of diagonal linear networks. The theoretical component of our study relies heavily on a line of work [38] [39] [40] [41] [42] that studies the inductive bias of a simplified "diagonal linear network" model. Diagonal linear networks parameterize linear maps f : R D → R asEQUATION+,d w+,d -w-,d wEQUATIONwhere ⃗ β( ⃗ w) ∈ R D . These correspond to two-layer linear networks in which the first layer consists of one-to-one connections, with duplicate + andpathways to avoid saddle point dynamics around ⃗ w = 0. Woodworth et al. [38] showed that overparameterized diagonal linear networks trained with gradient descent on mean squared error loss find the zero-training-error solution that minimizes ∥ ⃗ β∥ 2 , when trained from large initial weight magnitude (the "lazy" regime, equivalent to ridge regression). When trained from small initial weight magnitude, networks instead minimize ∥ ⃗ β∥ 1 (the "rich" regime). This bias is a linear analogue of feature learning/feature selection, as a model with an ℓ 1 penalty tends to learn solutions that depend on a sparse set of input dimensions.Implicit vs. explicit regularization. Theoretical work on the inductive biases conferred by different architectures has studied both the implicit regularization induced by gradient descent and explicit ℓ 2regularization. Notably, in homogeneous networks trained with crossentropy loss, implicit and explicit regularization yield identical inductive biases in the limit of infinite training time and infinitesimal regularization [12, 43] , but this does not hold in general [38] . While [38, 40] are able to characterize the implicit regularization of gradient descent for diagonal linear networks, it is technically much more challenging to derive a similar result for multi-output diagonal linear networks, or for ReLU networks of any kind. In contrast, the impact of explicit weight regularization has been characterized for both multi-output diagonal linear networks [16] and multi-output ReLU networks [44] [45] [46] .Our main technical contributions to this theoretical landscape are (1) spelling out the implications of existing results on implicit regularization in diagonal linear networks and (2) providing a novel characterization of the inductive bias induced by applying explicit parameter regularization to the finetuning of ReLU networks from arbitrary initialization. Our choice to study explicit parameter regularization for ReLU networks is made primarily for theoretical tractability; we view these results as a proxy for understanding the more theoretically complex problem of implicit regularization.3 Implicit and explicit regularization penalties for MTL and PT+FT EQUATIONSecond, we consider ReLU networks with H hidden neurons, hidden weights w ∈ R H×D , and O ∈ {1, 2} readout weights v ∈ R O×H . The network function is defined asEQUATIONwhere (•) + is the ReLU nonlinearity. Importantly, the ReLU nonlinearity is homogeneous and, as a result, the network function is invariant to rescaling the hidden weights by α > 0 and the readout weights by 1 α . It will therefore be useful to consider a re-parameterization in terms of the magnitudem h := v h ∥ ⃗ w h ∥ 2 and direction ⃗ θ h := ⃗ w h /∥ ⃗ w h ∥ 2 .(4) Under this re-parameterization, an equivalent definition of the network function is given by Training paradigms. We consider two datasets: an auxiliary task X aux ∈ R naux×D , ⃗ y aux ∈ R naux and a main task X main ∈ R nmain×D , ⃗ y main ∈ R nmain . First, we consider multi-task learning (MTL): simultaneous training on both the auxiliary and main task. In this case, we consider networks with O = 2 outputs. The first output corresponds to the auxiliary task and the second output to the main task. Accordingly, we denote β aux := β 1 and β main := β 2 in the diagonal linear network and ⃗ v aux := ⃗ v 1 and ⃗ v main := ⃗ v 2 in the ReLU network.EQUATIONSecond, we consider pretraining on the auxiliary task and subsequent finetuning on the main task (PT+FT). In this case, we consider networks with a single output (O = 1), but re-initialize the readout weights before finetuning. Accordingly, we denote the parameters learned after pretraining by w aux and v aux , and the parameters learned after finetuning by w main and v main . We define β main , β aux (for the diagonal network) and m main , θ main , m aux , θ aux (for the ReLU network) analogously.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations in the latter half of Section 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
explicit regularization penalties in multi-task learning: 
 To theoretically understand the inductive bias of diagonal linear and ReLU networks trained with MTL, we consider the effect of minimizing the ℓ 2 parameter norm as an approximation of the implicit bias of training with gradient descent from small initialization. We argue that this is a reasonable heuristic. First, the analogous result holds in the single-output case for infinitesimally small initialization and two layers (though not for deeper networks, see [38] ). Second, for cross-entropy loss it has been shown that gradient flow on all positively homogeneous networks (including diagonal linear networks and ReLU networks) converges to a KKT point of a max-margin/min-parameter-norm objective [12] . Finally, explicit ℓ 2 parameter norm regularization ("weight decay") is commonly used in practice, making its inductive bias important to understand in its own right as well.We now derive the norms minimized by explicit regularization in MTL: Corollary 1. For the multi-output diagonal linear network defined in Eq. 2, a solution β * with minimal parameter norm ∥w∥ 2 2 + ∥v∥ 2 2 subject to the constraint that it fits the training data (X main ⃗ β main = ⃗ y main , X aux ⃗ β aux = ⃗ y aux ) also minimizes the following:β * = arg min β 2 D d=1 (β aux d ) 2 + (β main d ) 2 s.t. X main ⃗ β main = ⃗ y main , X aux ⃗ β aux = ⃗ y aux .This norm is known as the group lasso [47] and denoted as ∥ • ∥ 1,2 . For ReLU networks, by an argument analogous to the one above, parameter norm minimization translates to minimizing [44, 46] ).H h=1 (m aux h ) 2 + (m main h ) 2 (see Appendix A.1 andThe norm ∥ • ∥ 1,2 is plotted in Fig. 1a . We analyze its impact in Section 4.Diagonal linear networks. We now consider the behavior of PT+FT in overparameterized diagonal linear networks trained to minimize mean-squared error using gradient flow. We assume that the network is initialized prior to pre-training with infinitesimal weights, and that during pretraining, network weights are optimized to convergence on the training dataset (X aux , ⃗ y aux ) from the auxiliary task. After pretraining, the second-layer weights (v + and v -) are reinitialized with constant magnitude γ. Further, to ensure that the network output pre-finetuning is zero (as in [38] ), we set the values of corresponding positive and negative pathway weights to be equal to their sum following pretraining 1 . The network weights are further optimized to convergence on the main task dataset (X main , ⃗ y main ).The dynamics of the pretraining and finetuning steps can be derived as a corollary of [38, 40] :Corollary 2.If the gradient flow solution ⃗ β aux for the diagonal linear model in Eq. 1 during pretraining fits the auxiliary task training data with zero error (i.e. X aux ⃗ β aux = ⃗ y aux ), and following reinitialization of the second-layer weights and finetuning, the gradient flow solution ⃗ β * after finetuning fits the main task data with zero training error (i.e. X main ⃗ β main = ⃗ y main ), then⃗ β * = arg min ⃗ β main ∥ ⃗ β main ∥ Q s.t. X ⃗ β = ⃗ y, ∥ ⃗ β main ∥ Q := D d=1 |β aux d | + γ 2 q 2β main d |β aux d | + γ 2 , q(z) = 2 -4 + z 2 + z • arcsinh(z/2).This corollary is proven in Appendix A.2. The norm is plotted in Fig. 1b .ReLU networks. We assume that after pretraining, the readout layer is re-initialized with arbitrary new weights ⃗ γ ∈ R H . We then characterize the solution to the finetuning task that minimizes the ℓ 2 norm of weight changes from this initialization. This is similar to the explicit weight regularization considered in the previous section, except we now penalize weight changes from a particular initialization rather than the origin. We chose to consider this regularization penalty for two reasons. First, it is sometimes studied in the continual learning setting [30, 48] . Second, infinitesimal explicit regularization is equivalent to the implicit regularization induced by gradient descent in the case of shallow linear models [10] . While this is not true more generally, it is a useful heuristic to motivate our theoretical analysis (which we then validate using our experiments in Section 4).We show that this finetuning solution implicitly minimizes the following penalty: Proposition 3. Consider a single-output ReLU network (see Eq. 3) with first-layer weights ⃗ w aux h ∈ R D after pretraining, and second-layer weights re-initialized to ⃗ γ ∈ R H . The solution to the finetuning task that minimizes the ℓ 2 norm of changes in the weights, i.e. minimizesH h=1 ∥ ⃗ w h - ⃗ w aux h ∥ 2 2 + (v h -γ h ) 2, is equivalent to the solution that minimizesEQUATIONWe prove the proposition in Appendix A.3. It implies that the regularization penalty associated with finetuning in the ReLU network only depends on the correlation ρ h := ⟨ ⃗ θ main h , ⃗ θ aux h ⟩ between the first-layer feature weights before and after finetuning, and the magnitudes of the weights of these features. We plot this penalty in Fig. 1c .4 Implications of the theory: multiple regimes of feature reuseQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: We derive 1) the penalty associated with MTL in ReLU networks in Appendix A.1, 2) the penalty associated with PT+FT in ReLU networks in Appendix A.3, and 3) the conservation law (Prop. 4) in Appendix A.4.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
sample efficiency in teacher-student models: 
 To validate these theoretical characterizations and illustrate their consequences, we now perform experiments in a teacher-student setup. In the diagonal linear network case, we consider a linear regression task defined by ⃗ w ∈ R 1000 with a sparse set of k non-zero entries. We sample two such vectors, corresponding to "auxiliary" and "main" tasks, varying the number of non-zero entries k aux and k main , and the number of shared features (overlapping non-zero entries). We then uniformly sample input vectors ⃗ x ∈ R 1000 from the unit sphere, using the ground-truth weights to generate the target. We train on 1024 auxiliary samples and vary the number of main task samples.In the ReLU network case, we consider a "teacher" ReLU network with a sparse number of units (i.e. a low-dimensional hidden layer) and different kinds of overlaps (e.g. shared, correlated, or orthogonal features) between the auxiliary and main task. We randomly sample input data ⃗ x ∈ R 15 from the unit sphere and use the teacher network to generate the target. We train on 1024 auxiliary samples and vary the number of main task samples. During finetuning, we randomly re-initialize the readout weights using a normal distribution with a variance of 10 -3 2/H.Feature sparsity. To understand the impact of the derived penalties in the case of features either used or not used during pretraining, we consider their limit behavior for very large or very small pretrained features. First, we consider the limit|β d | |β aux d |+γ 2 → ∞ (capturingthe case of a feature not used during pretraining). In this limit, the MTL penalty converges to 2|β d |. Similarly, for finetuning in diagonal linear networks, the penalty converges to c|β d | where c ∼ O log 1/(|β aux d | + γ 2 ) (per an analysis in [38] ). For new features, both networks therefore have an ℓ 1 norm minimization bias, suggesting that they tend to learn a sparse set of new features (just like in the single-task case).To test this insight, we consider a teacher-student setup without any shared features between the auxiliary and main task, and vary the number of main task features. Indeed, we found that both MTL and PT+FT have a more rapidly decreasing generalization loss for fewer features (just like single-task learning (STL)) (Fig. 2a ). We further confirmed that they learned a sparse set of weights (Fig. 7 ).The same phenomenon holds true for the ReLU network penalties as well. We considered a teacherstudent network with six auxiliary task features and one to six uncorrelated main task features. Again, both MTL and PT+FT have a lower generalization loss for fewer features (Fig. 2b ).Feature sharing. Next, we consider the opposite limit, i.e. large pretrained features:|β d | |β aux d |+γ 2 → 0.In this limit, both the MTL and the PT+FT penalty for diagonal linear networks converges toβ 2 d |β aux d | ,a weighted ℓ 2 bias. This implies that the networks preferentially use large pretrained features. To test this, we consider a teacher-student setup with fully overlapping dimensions. Indeed, both MTL and PT+FT outperform STL in this case (Fig. 2c ). Notably, they perform similarly to a network where we only finetune the second layer (PT+FT (Linear Probing, LP)), which exactly implements the weighted ℓ 2 bias.In the case of ReLU networks, we considered a teacher network with the same six features for the auxiliary and main task. Again, we found that MTL and PT+FT outperformed STL, though training a linear readout from a model with fixed features (either by using the hidden layer (PT+FT (LP)) or the neural tangent kernel (PT+FT (NTK))) generalized even better (Fig. 2d ).Simultaneous sparsity and feature sharing. Finally, our analysis above considered the limit behavior of each feature separately. This suggests that models should be able to 1) preferentially rely on pretrained features and 2) when necessary, learn a sparse set of new features. To test this insight, we consider partially overlapping teacher models. In diagonal linear networks, we consider teacher models with forty auxiliary and main task features, with twenty or thirty of those features overlapping (Fig. 2e ). On the one hand, both PT+FT and MTL outperformed STL, indicating that they were able to benefit from the pretrained features. On the other hand, they also performed better than the PT+FT (LP) model which only finetuned the second layer (and therefore did not have a sparse inductive bias), indicating that they tended to learn a sparse set of new features. In ReLU networks, we consider teacher models with six auxiliary and main task features, varying their overlap. Again, we find that MTL and PT+FT outperformed both STL and PT+FT (LP), indicating that they benefitted from feature learning by implementing an inductive bias towards sparse and shared features.Differences between the MTL and PT+FT norms. So far, we have highlighted several similarities between the MTL and PT+FT norms in diagonal linear networks: they tend towards the ℓ 1 norm for small auxiliary features and a weighted ℓ 2 norm for large auxiliary features. In the next section, we will highlight an important difference that arises in the intermediate regime. Here we briefly highlight a difference in the limit behavior for diagonal linear networks: the relative weights of the ℓ 1 -and weighted ℓ 2 -penalty are different between MTL and PT+FT. In particular, in the ℓ 1 penalty limit, there is an extra factor of order O log 1/(|β aux d | + γ 2 ) in the PT+FT penalty. Assuming small initializations, this factor tends to be larger than 2, the corresponding coefficient in the MTL penalty. Thus, PT+FT is more strongly biased toward reusing features from the auxiliary task compared to MTL. We are careful to note, however, that in the case of ReLU networks this effect is complicated by a qualitatively different phenomenon with effects in the reverse direction (see Section 4.5).We now turn our attention to the intermediate regime where the coefficients of a feature are of similar magnitude in the auxiliary and main tasks. We define two functions for a given penalty P (β main , which measures, locally, how much P decreases for a larger auxiliary feature. In the previous section, we found that for β aux d → 0, ℓ-order → 1 and FD → 0, i.e. the penalty becomes ℓ 1 -like and does not depend on the magnitude of the auxiliary feature. In contrast, for β aux d → 1, ℓ-order → 2 and FD → -1, i.e. the penalty becomes ℓ 2 -like and depends inversely on the auxiliary task feature coefficient magnitude.How are ℓ-order and FD related for intermediate values of β aux d ? Remarkably, we find an exact analytical relationship between these measures that holds for all penalties considered here: Proposition 4. For the MTL and PT+FT penalties derived in both the diagonal linear and ReLU cases, ℓ-order + F D = 1.Thus, there is an exact tradeoff between the sparsity and feature dependence even for intermediate auxiliary task feature coefficient values: an increase in sparsity (i.e. a smaller ℓ-order) yields a corresponding decrease in feature dependence (i.e. FD becomes closer to zero). = 1 and varying the auxiliary task feature coefficient. b, Generalization loss for models trained on a teacher with 40 active units during the auxiliary task and a subset of those units active during the main task. c, Generalization loss for PT+FT models whose weights are rescaled by the factor in the parentheses before finetuning. d-f, ReLU networks. d, ℓ-order/feature dependence plotted for the explicit finetuning and MTL penalties, for m = 1 and varying the auxiliary task feature coefficient. e, Generalization loss for models trained on a teacher network with six active units on the auxiliary task and a subset of those units on the main task. f, Generalization loss for PT+FT models whose weights are rescaled before finetuning.Proposition 4 predicts the existence of a novel "nested feature selection" regime: for intermediate magnitudes of the auxiliary features, the penalties should encourage both sparsity and feature dependence. To test this prediction in diagonal linear networks, we use a teacher-student setting in which all of the main task features are a subset of the auxiliary task features, i.e. k main ≤ k aux , and the number of overlapping units is equal to k main . Solving this task most efficiently involves "nested feature selection," a bias towards feature reuse and towards sparsity among the reused features. We plot ℓ-order and FD for β main d = 1 and varying the auxiliary task feature coefficient (Fig. 3a ) and find that for β aux d ≈ 1, both norms exhibit "lazy regime"-like behavior (ℓ-order of around 2, and FD near 0). This predicts that neither MTL nor PT+FT networks should be able to benefit from nested sparsity task structure in this regime, which we confirm empirically: as k main decreases, the networks' sample efficiency does not become substantially better (Fig. 3b ).However, for features with auxiliary task coefficients that are moderately smaller than their main task coefficients, Fig. 3a suggests a broad regime where the ℓ-order is closer to 1 (incentivizing sparsity), but feature dependence remains high. We can produce this behavior in these tasks by rescaling the weights of the network following pretraining by a factor less than 1. In line with the prediction of the theory, performing this manipulation enables PT+FT to leverage sparse structure within auxiliary task features (Fig. 3c ), even while retaining their ability to privilege features learned during pretraining above others (see Fig. 9 ). By contrast, this regime is much narrower for MTL (Fig. 3a ).For ReLU networks, the MTL penalty is the same as in diagonal linear networks. We plot the regularization penalty derived for PT+FT, conditioned on various correlations ⟨θ main h , θ aux h ⟩ between the post-pretraining and post-finetuning feature weights (Fig. 3d ). For features that are fully aligned before and after pretraining, this penalty is again ℓ 2 -like for m aux ≈ 1. However, in most cases, these features change direction at least a little bit, and we find that in that case the penalty is more ℓ 1 -like while still remaining sufficiently feature-dependent. This suggests that a nested feature selection regime may arise in ReLU networks even when auxiliary task feature coefficients have comparable magnitude to main task feature coefficients. To test this insight, we considered a set of tasks in which the main task solution relies exclusively on a subset of auxiliary task features. We found that PT+FT (even without any weight rescaling) was able to benefit from the nested sparsity structure, but MTL was not (Fig. 3e ). 2 Performing weight rescaling in either direction following pretraining uncovers the initialization-insensitive (FD near 0), sparsity-biased (ℓ-order near 1) rich / feature-learning regime and the initialization-biased (FD near -1), no-sparsity-bias (ℓ-order near 2) lazy learning regime (Fig. 3f ). This suggests that for different architectures and different tasks, different rescaling values may be required to enter the nested feature selection regime. : PT+FT, but not MTL, in ReLU networks benefits from correlated features. a, Generalization loss for main task features that are correlated (0.9 cosine similarity) with the auxiliary task features. PT+FT outperforms both MTL and STL. b, Generalization loss for main task features with varying correlation and magnitude (mag.). PT+FT only outperforms STL if the features are either identical in direction or identical in magnitude.The regularization associated with PT+FT yields benefits even when main/auxiliary task directions are correlated but not identical (Fig. 1c ). In contrast, MTL cannot softly share features as it encodes correlated features with entirely distinct units. To test this hypothesis, we conduct experiments in which the ground-truth auxiliary and main tasks rely on correlated but distinct features. Indeed, we find PT+FT outperforms STL in this case, whereas MTL only does so in the low-sample setting (Fig. 4a ). Notably, if features are identical, MTL outperforms PT+FT in the low-sample setting (Fig. 2d ). Thus, PT+FT (compared to MTL) trades off the flexibility to "softly" share features for reduced sampleefficiency when such flexibility is not needed.The analysis in Fig. 3d would predict that ReLU networks can no longer benefit from correlated features if the magnitude of auxiliary task features is much higher than that of main task features. To test this, we varied the magnitude of the main task features and their correlation with auxiliary task features (Fig. 4b ). We found that for lower magnitudes, PT+FT still improved performance for the case of identical but not correlated feature directions, confirming our prediction.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We describe the way in which we generate the teacher-student experiments in Section 3.3 and further discuss our methods in detail in Appendix B. Finally we are providing a codebase to reproduce all figure in the main text. 
weight rescaling in deep networks gives rise to nested feature selection: 
 Our analysis has focused on shallow networks trained on synthetic tasks. To test the applicability of our insights, we conduct experiments with convolutional networks (ResNet-18, [49] ) on a vision task (CIFAR-100, [50] ), using classification of two image categories (randomly sampled for each training run) as the primary task and classification of the other 98 as the auxiliary task. As in our experiments above, MTL and PT+FT improve sample efficiency compared to single-task learning (Fig. 5a ).Our findings in Section 4.4 indicate that the nested feature selection bias of PT+FT can be exposed or masked by rescaling the network weights following pretraining. Such a bias may be beneficial when the main task depends on a small subset of features learned during pretraining, as may often be the case in practice. We experiment with rescaling in our CIFAR setup. We find that rescaling values less than 1 improve finetuning performance (Fig. 5b ). These results suggest that rescaling network weights before finetuning may be practically useful. We corroborate this hypothesis with additional experiments using networks pre-trained on ImageNet [51] (see Fig. 10 ).To facilitate comparison of the phenomenology in deep networks with our teacher-student experiments above, we propose a signature of nested feature selection that can be characterized without knowledge of the underlying feature space (since the correct feature basis to analyze is less clear in multi-hidden-layer entworks). Specifically, we propose to measure (1) the dimensionality of the network representation pre-and post-finetuning, and (2) the extent to which the representational structure post-finetuning is shared with / inherited from that of the network following pretraining prior to finetuning. We employ the commonly used participation ratio (PR, [52] ) as a measure of dimensionality, and the effective number of shared dimensions (ENSD, [53] ) as a soft measure of the number of aligned principal components between two representations. Intuitively, the PR and ENSD of network representations pre-and post-finetuning capture the key phenomena of the nested feature selection regime: we expect the dimensionality of network after finetuning to be lower than after pretraining (P R(X F T ) < P R(X P T )), and for nearly all of the representational dimensions expressed by the network post-finetuning to be inherited from the network state after pretraining (EN SD(X P T , X F T ) ≈ P R(X F T )). We validate that this description holds in our nonlinear teacher-student experiments with networks in the nested feature selection regime (Fig. 11 ).Remarkably, we find that the ResNet-18 exhibits the same phenomenology, but only for weights rescaled by a small value (Fig. 5c ). This supports the hypothesis that the observed benefits of rescaling indeed arise from pushing the network into the nested feature selection regime.Finally, we conduct the same experiment for Vision Transformers (ViT) [54] . We confirm that PT+FT improves performance over single-task learning, though MTL offers no similar benefit in this case (Fig. 5d ). We further find that rescaling before finetuning (both by larger and smaller values) decreases generalization performance (Fig. 5e ). Notably, our representational analysis reveals that rescaling by a larger value yields a higher-dimensional subspace both before and after finetuning, whereas rescaling by a smaller value yields a lower-dimensional subspace, but pushes the effective number of shared dimensions down substantially (Fig. 5f ). This indicates that a rescaling value of 1 may already give rise to the nested feature selection regime and rescaling by a smaller value pushes the ViT towards the pure feature learning regime. Taken together, our results suggest finetuning performance is best when networks operate in the nested feature selection regime, and weight rescaling can push networks into this regime when it does not arise naturally.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We have attached a codebase that enables reproduction of all figures in the main text.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 In this work we have provided a detailed characterization of the inductive biases associated with two common training strategies, MTL and PT+FT, in diagonal linear and ReLU networks. These biases incentivize both feature sharing and sparse task-specific feature learning. In the case of PT+FT, we characterized a novel nested feature selection learning regime which encourages sparsity within features inherited from pretraining. This insight motivates a simple technique for improving PT+FT performance by pushing networks into this regime, which shows promising empirical results.There are several avenues for extending our theoretical work: for example, connecting our derived penalties for ReLU networks (which assumed explicit parameterization) to the implicit regularization induced by dynamics of gradient descent, and extending our theory to the case of cross-entropy loss.In addition, more work is needed to extend our theory to more complex tasks and larger models. For instance, we are interested in investigating how the weight magnitudes required to enter the nested feature regime depend on architecture and properties of the tasks -we observed a difference between the rescaling values required for ResNets and Vision Transformers, which our shallow network theory is unable to speak to. Better understanding the conditions needed for nested feature selection could also inspire more sophisticated interventions than our weight rescaling trick. Finally, we considered a particularly simple multi-task setup, with identical formats between auxiliary and main tasks -our theory could be extended to cases where the different tasks use different objectives. Nevertheless, our work already provides new and practical insights into multi-task learning and finetuning.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide these details in Appendix B Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: For all experiments, we report the mean ± two standard deviations.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This manuscript presents foundational research and there is no clear potential risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We use two standard vision datasets: CIFAR-100 and Imagenet. For CIFAR-100, we were unable to find an associated license. For Imagenet, the data is available for free to researchers for non-commercial use. Both websites hosting these datasets request a citation of the relevant paper which we provided.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Multi-task ReLU networks with a shared feature layer and O outputs can be written asEQUATIONEQUATIONwhere ⃗ w h , m h ∈ R O . We are interested in how minimizing the parameter normH h=1 ∥ ⃗ w h ∥ 2 2 +∥⃗ v h ∥ 2 2maps to minimizing a norm over the solution weights ⃗ m. For a given ⃗ m h ∈ R O , the norm of the input weight ∥ ⃗ w h ∥ 2 is a free parameter z h > 0, as we can set ⃗ v h := ⃗ m h /z. As any value for z h leads to the same function, we choose z h so as to minimizeEQUATIONSetting the derivative to zero impliesEQUATIONEQUATIONwhere the right-hand side of the equation is the ℓ 1,2 norm.A.2 Finetuning in diagonal linear networks Corollary 2. If the gradient flow solution ⃗ β aux for the diagonal linear model in Eq. 1 during pretraining fits the auxiliary task training data with zero error (i.e. X aux ⃗ β aux = ⃗ y aux ), and following reinitialization of the second-layer weights and finetuning, the gradient flow solution ⃗ β * after finetuning fits the main task data with zero training error (i.e. X main ⃗ β main = ⃗ y main ), then⃗ β * = arg min ⃗ β main ∥ ⃗ β main ∥ Q s.t. X ⃗ β = ⃗ y, ∥ ⃗ β main ∥ Q := D d=1 |β aux d | + γ 2 q 2β main d |β aux d | + γ 2 , q(z) = 2 -4 + z 2 + z • arcsinh(z/2).Proof. We derive this result in two steps: first, we note that after pretraining (because of the infinitesimal initial weight magnitudes), one of the pathways remains at zero and one pathway encodes the effective linear predictor β aux . Thus, the first hidden layer has the weights √ β aux , which yields the new weight magnitude along the positive and negative pathway (due to the summing operation we perform following finetuning, which ensures that the weights of both pathways are set to √ β aux ). Having set the readout initialization to γ, we then apply Theorem 4.1 in [40] .Proposition 3. Consider a single-output ReLU network (see Eq. 3) with first-layer weights ⃗ w aux h ∈ R D after pretraining, and second-layer weights re-initialized to ⃗ γ ∈ R H . The solution to the finetuning task that minimizes the ℓ 2 norm of changes in the weights, i.e. minimizesH h=1 ∥ ⃗ w h - ⃗ w aux h ∥ 2 2 + (v h -γ h ) 2, is equivalent to the solution that minimizesEQUATIONProof. We define the norm of the deviation from the parameters at initialization asEQUATIONNote that w and v are parameterized in a redundant manner, as we can multiply v h by a constant a > 0 and obtain the same function as long as we divide ⃗ w h by the same constant. To parameterize a function in a unique fashion, we describe it in terms of the normalized hidden weight ⃗ θ h := ⃗ w h /∥ ⃗ w h ∥ 2 and the (signed) magnitude m h := ∥ ⃗ wh ∥ 2 v h . We wish to compute the norm of the deviation from the initialization in terms of this parameterization. Specifically, we computeEQUATIONNote that Woodworth et al. [38] solve a version of this optimization problem for diagonal linear networks (though only for the case in which corresponding first and second layer weights are initialized with equal magnitude).As this optimization problem decomposes over hidden units, we can solve it for each hidden unit individually:EQUATIONEQUATIONWe can express this asEQUATIONNote that this function is smooth and diverges to ∞ both for u → ∞ and u → 0. As long as we have a unique positive stationary point u * , this stationary point must therefore be a minimum. We simplifyEQUATIONand, defining m aux h := ∥ ⃗ w aux h ∥ 2 , and the cosine similarity between corresponding pretraining and finetuning task weightsρ h := ⟨ ⃗ w (1) h , ⃗ w aux h ⟩ ∥ ⃗ w (1) h ∥2∥ ⃗ w aux h ∥2, write this asEQUATIONWe then determine the stationary points by computing the derivativeEQUATIONand setting it to zero. Simplifying this equation (and multiplying by u 3 /2) yieldsEQUATION)We can compute the solution to this quartic equation, e.g. using np.roots and select the (unique, in all empirical cases we explored) positive real root, u * . We can then compute the original norm by plugging in u * :EQUATIONA.4 Proof of Proposition 4 Proposition 4. For the MTL and PT+FT penalties derived in both the diagonal linear and ReLU cases, ℓ-order + F D = 1.Proof. Below, for simplicity's sake, we always denote main β, main m, and θ = MTL We first consider the penalty, which identical for diagonal linear networks ReLU (without loss of generalization) aux > 0. In that case,EQUATIONandEQUATION)Thus, ℓ-order + FD = 1.Finetuning in diagonal network. We assume (without loss of generalization) β aux d > 0 and deriveEQUATIONandEQUATIONThus, ℓ-order + FD = 1.Finetuning in ReLU networks. We defineEQUATIONwhere, without loss of generalization, we assume that m aux h > 0 and note thatEQUATIONwhere u * is the critical point, depending on ⃗ θ h , m h , ⃗ θ aux h , m aux h . By the chain rule,EQUATIONand the analogous statement is true for m aux h . However, by definition,EQUATIONThus,EQUATIONTo prove the statement, we therefore must prove thatEQUATIONTo do that (30) impliesEQUATIONand thereforeEQUATIONPlugging this into (32) yieldsEQUATIONand completes the proof.We trained all networks with PyTorch [55] .We train the diagonal linear networks until they have reached a mean squared error below 10 -10 , for pretraining, finetuning, multi-task learning and single-task learning. As a rough heuristic for determining the learning rate, we begin at a learning rate of 10 6 and divide by ten whenever the loss exceeds 100 (indicating divergence).We train ReLU networks with 1000 hidden units using the same learning rate heuristic and train until the mean squared error has reached a threshold of 10 -8 .We consider a ResNet-18 in the standard PyTorch implementation and train it with gradient descent with a learning rate of 10 -3 and momentum of 0.9.We consider a Vision Transformer with seven layers, eight heads, 384 hidden dimensions and an MLP with 1,536 hidden dimensions. We train the transformer using the Adam optimizer with a learning rate of 10 -3 and weight decay 0.00005.We provide a code database that enables reproducing all data used to create the main figures. We used one CPU for each diagonal network (total number of experiments: 5,544) and ReLU network experiment (total number of experiments: 6,740) and all of them took up to four hours. We used on GPU for the CIFAR and ViT experiments and all of them took up to twelve hours (total number of experiments: 5,000). As a rough (upper-bound) estimate, the experiments therefore required 50,000 hours of CPU time and 60,000 hours of GPU time. To increase confidence that our main results are robust to the number of data samples used (1024 auxiliary task samples and up to 1024 main task samples in most of our experiments), and the number of ground-truth units in the teacher network (6), we repeated the experiments of Fig. 2d with 8192 auxiliary task samples and 40 ground-truth features. Indeed, in this setting the rich regime also helps with generalization if and only if the teacher units are sparse (Fig. 6a ). Further, MTL and PT+FT tend to outperform STL if the features are overlapping and MTL tends to outperform PT+FT (Fig. 6b ). In particular, the finetuned networks still benefit from feature learning, especially if some features are novel.Our theory predicts inductive biases towards solutions that minimize norms, often either ℓ 1 -like (incentivizing sparsity) or ℓ 2 -like. Our experiments in the main text corroborate these description by analyzing how sample complexity depends on the feature sparsity of the ground-truth task solution, and how the sparse feature structures of the main and auxliary tasks relate. However, this evidence for sparsity biases (or lack thereof) is indirect; here we present more direct analyses of the learned solutions in linear and nonlinear networks that support the account we provide in the main text.To check whether the implicit regularization theory is a good explanation for these performance results, we directly measured the ℓ 1,2 and Q norms of the solutions learned by networks, compared to the corresponding penalties of the ground truth weights. In Fig. 7a we see that as the amount of training data increases, the norms all converge to that of the ground truth solution, but in the low-sample regime, MTL and PT+FT find solutions with lower values of their corresponding norm than the ground-truth function, consistent with the implicit regularization picture (by contrast, STL does not consistently find solutions with lower values of these norms than the ground truth).Our theory predicts that weight rescaling by a factor less than 1.0 following pretraining reduces the propensity of the network to share features between auxiliary and main tasks during finetuning. We confirm that this is the case in Fig. 7b by analyzing the overlap between the auxiliary task features and the learned linear predictor for the main task.In Fig. 7c we show that our measure of effective sparsity of learned solutions in diagonal linear networks effectively distinguishes between networks trained in the feature selection regime and networks trained with linear probing (only training second-layer weights). Moreover, in Fig. 7d we show that the L1 norm of the solution increases with the training sample size, consistent with a bias towards L1 minimization. There is an interesting discrepancy between the behavior of the sparsity of the solutions (nonmonotonic, see Fig. 7c ) and the L1 norm (largely monotonic, see Fig. 7d ). This is attributable to the discrepancy between the L1 norm (which diagonal linear networks in the rich regime are biased to minimize) and sparsity (for which L1 norm is only a proxy). 3b, c . Weight rescaling decreases this overlap. c, Proportion of variance concentrated in the top k weights, as a function of k, for training on a single-task. When both layers are trained from small initialization (STL), this variance decreases much more rapidly than for pure linear readout training (STL (LP)), demonstrating the sparsity of the learned solution. d, L1 norm for STL as a function of the number of samples. e, Proportion of variance across different overlaps and for different learning setups (see also Fig. 2b ). The rapid decrease in variance demonstrates the sparsity of the learned solutions both for PT+FT and MTL. f, Proportion of weight norm in the 40 dimensions relevant for the auxiliary task (see also Fig. 2b ).In Fig. 7e we show the sparsity of networks trained on teacher with different values of the number of overlapping features between main and auxiliary tasks (each of which uses 40 features). We find that learned solutions across a range of overlaps are as sparse as using single-task learning (see Fig. 7c ) when task features do not overlap (0/40 case) and more sparse otherwise (on account of the bias toward reuse of the sparse features learned during the auxiliary task, see next paragraph).In Fig. 7f we show the overlap between features active on either task, finding that learned main task solutions are biased to share auxiliary task features when few samples are available.We adopt a clustering-based approach to analyzing the effective sparse structure of learned task solutions. Specifically, for a given trained network, we perform k-means clustering with a predetermined value of K clusters on the normalized input weights to each hidden-layer neuron in the network 3 . We measure the extent to which K cluster centers are able to explain the variance in input weights across hidden units; the fraction of variance left unexplained is commonly referred to as the "inertia." For values of K at which the inertia is close to zero, we can say that (to a good approximation) the network effectively makes use of at most K nonlinear features.In the single-task learning case, we measure the inertia of trained networks as a function of K. We find that for networks in the rich regime (small initialization scale), for tasks with sparse ground-truth (six units in the ReLU teacher network), the networks do indeed learn solutions that make use of approximately six nonlinear features (Fig. 8a ). For tasks with many (1000) units in the teacher network, the network finds solutions that use a small number of feature clusters when main task samples are limited, but gradually uses more clusters as the number of samples is increased (Fig. 8a ), at which point the network matches the teacher function very well. This bias towards sparser-thanground-truth solutions given insufficient data corroborates our claim of an inductive bias towards sparse solutions. By contrast, networks in the lazy learning regime (large initialization scale) display no such bias, corroborating our claim that the sparse ℓ 1 -like inductive bias is a property of the rich regime but not the lazy regime. Interestingly, in the sparse ground-truth case learned solutions are relatively less sparse for an intermediate number of training examples. This may arise because an ℓ 1 -like inductive bias is not exactly the same as a bias toward sparse solutions over nonlinear features, particularly when training data is limited. We leave an in-depth investigation of this phenomenon to future work.Our clustering analysis allows us to measure the extent to which the effective features employed by the network (cluster centers) are aligned with the ground-truth task features. Specifically, for each teacher unit, we compute an "alignment score" between teacher and student networks by taking each teacher unit, measuring its cosine similarity with all the centers, choosing the maximum value, and averaging this quantity across all teacher units. We find that the learned feature clusters are indeed highly aligned with the ground-truth teacher features in the sparse ground-truth case, and moreso as the number of main task samples (and consequently task performance) increases (Fig. 8b ).We find that pretraining+finetuning improves performance over single-task learning when main and auxiliary task features are shared (or correlated), and maintains an apparent bias toward sparsity in new task-specific features. To corroborate these claims, we performed our clustering analysis on the solutions learned through PT+FT. We find that the solutions learned are indeed quite sparse (comparable to the sparsity of solutions learned by single-task learning), even when the auxiliary task and main task features are disjoint (Fig. 8c ). Moreover, we find that MTL also learns sparse solutions (Fig. 8d ). In particular, as expected, the effective features on tasks with overlapping features is equal to the number of total unique features. Moreover, we observe that when main task and auxiliary task features are shared, PT+FT and MTL networks exhibit higher alignment between learned features and ground-truth features than single-task-trained networks, especially when main task samples are MTL. e, Alignment score for networks trained with MTL, PT+FT, and STL on the same tasks as in panels c and d. f Inertia (using k = 1 clusters) for networks trained on an auxiliary task that relies on only one ground-truth feature, which is one of the six ground-truth features used in the auxiliary task (as in Fig. 3e, f ), using MTL or PT+FT with various rescaling factors applied to the weights prior to finetuning. g Alignment score for the networks and task in panel f. limited (Fig. 8e ). This provides a mechanistic underpinning for the relationship between the inductive bias of PT+FT that we describe in the main text and its performance benefits.In the main text we describe the "nested feature selection" regime, which occurs at intermediate values of the ratio between ground-truth main task feature coefficients and pretrained network weight scale. In this regime, networks can more efficiently learn main tasks that make use of a subset of the features used in the auxiliary task. Importantly, they still maintain a bias towards reusing features from the auxiliary task, as we show for diagonal linear networks and ReLU networks in Fig. 9 . Here we show that networks in this regime (obtained most clearly in the shallow ReLU network case when networks are rescaled by a value of 1.0 after pretraining) indeed learn very sparse (effectively 1-feature) solutions when the ground-truth main task consists of a single auxiliary task features (Fig. 8e , right). By contrast, networks with weights rescaled by a factor of 10.0 following pretraining exhibit no such nested sparsity bias (consistent with lazy-regime behavior). Similarly, multi-task networks cannot exhibit such a bias in their internal representation as they still need to maintain the features needed for the main task (Fig. 8e , left). Additionally, supporting the idea that the nested feature selection regime allows networks to benefit from feature reuse, we find that networks in this regime exhibit a higher alignment score with the ground-truth teacher network when the main task features are a subset of the auxiliary task features compared to when they are disjoint from the auxiliary task features (Fig. 8g ). This alignment benefit is mostly lost when networks are rescaled by a factor of 0.1 following pretrainning (a signature of rich-regime-like behavior).To evaluate the robustness / general-purpose utility of our suggested approach of rescaling network weights following pretraining, we experimented with finetuning convolutional networks pretrained on ImageNet on downstream classification tasks: pretrained ResNet-18 finetuned on CIFAR100, pretrained VGG11 finetuned on CIFAR100, and pretrained ResNet-18 finetuned on STL-10. We experimented both with finetuning on the full multi-way classification task, and also on binary classification tasks obtained by subsampling pairs of classes from the main task dataset (which we found exposes performance differences more strongly). Due to computational constraints, we did not sweep over the choice of the rescaling factor, but simply used a factor of 0.5 in all cases. We find that rescaling improves finetuning performance, to varying degrees, in all of our experiments (Fig. 10 ).F Analysis of representations learned in the nested feature selection regime: bridging the gap from shallow to deep networksWe were interested in testing whether our theoretical understanding of shallow networks is truly responsible for the behavior of deeper networks (with more direct evidence than performance results / sample complexity). Specifically, we sought to understand whether the observed benefit of rescaling network weights following pretraining (Fig. 5b , Appendix E) relates to the nested feature selection regime we characterized in shallow networks. Doing so is challenging, as the space of "features" learnable by a deep network is difficult to characterize explicitly (making the feature clustering analysis employed in Appendix D inapplicable). To circumvent this issue, we propose a signature of nested feature selection that can be characterized without knowledge of the underlying feature space. Specifically, we propose to measure (1) the dimensionality of the network representation pre-and post-finetuning, and (2) the extent to which the representational structure post-finetuning is shared with / inherited from that of the network following pretraining prior to finetuning.We employ the commonly used participation ratio [PR; 52] as a measure of dimensionality. For an n × p matrix X representing n mean-centered samples of p-dimensional network responses, with a p × p covariance matrix C X = 1 n X ⊤ X, the participation ratio is defined asEQUATIONwhere λ i are the eigenvalues of the covariance matrix C X . The PR scales from 1 to p and measures the extent to which the covariance structure of responses X is dominated by a few principal components or is spread across many. We argue that low-dimensional representations are a signature of networks that use a sparse set of features. We confirm that this is the case in our teacher-student setting: networks in the rich regime, which are biased towards sparse solutions, learn representations with lower PR than networks in the lazy regime, which are not biased toward sparse solutions (Fig. 11a ).Our measure of shared dimensionality between two representations is the effective number of shared dimensions (ENSD) introduced by Giaffar et al. [53] . The ENSD for an n × p matrix of responses X from one network and an n × p matrix of responses Y from another network, denoted EN SD(X, Y ), is given byEQUATIONThis measure is equal to the centered kernel alignment (CKA), a measure of similarity of two network representations [56] , multiplied by the geometric mean of the participation ratios of the two representations. It measures an intuitive notion of "shared dimensions" -for example, if X consists of 10 uncorrelated units, if Y is taken from a subset of five of those units, the ENSD(X, Y) will be 5. If Y is taken to be five uncorrelated units that are themselves uncorrelated with all those in X, the ENSD(X, Y) will be zero.Intuitively, the PR and ENSD of network representations pre-and post-finetuning capture the key phenomena of the nested feature selection regime. In a case in which the main task uses a subset of the features of the auxiliary task, if the network truly extracts this sparse subset of features, we expect the dimensionality of network after finetuning to be lower than after pretraining (P R(X F T ) < P R(X P T )), and for nearly all of the representational dimensions expressed by the network post-finetuning to be inherited from the network state after pretraining (EN SD(X P T , X F T ) ≈ P R(X F T )). By contrast, networks not in the nested feature selection regime should exhibit an ℓ 2 -like rather than ℓ 1 -like bias with respect to features inherited from pretraining and thus not exhibit a substantial decrease in dimensionality during finetuning.We show that this description holds in our nonlinear teacher-student experiments. Networks that we identified as being in the "nested feature selection" regime (weights rescaled by 1.0 following pretraining), and also networks in the rich regime, exhibit decreased PR following finetuning (Fig. 11b ). By contrast, lazy networks (weights rescaled by 10.0 following pretraining) exhibit no dimensionality decrease during finetuning. Additionally (see Fig. 11c ), the ENSD between pretrained (PT) and finetuned (FT) networks is almost identical to the dimensionality of the finetuned representation (PR FT).Strikingly, we observe very similar behavior in our ResNet-18 model pretrained on 98 CIFAR-100 classes and finetuned on the 2 remaining classes (Fig. 11d ), when we apply our method of rescaling weights post-finetuning. Analyzing the PR and ENSD of the outputs of different stages of the network following pretraining and following finetuning, we see that dimensionality decreases with finetuning, and ENSD between the pretrained and finetuned networks is very close to the PR of the finetuned network. Moreover, this phenomenology is only observed when we apply the weight rescaling method; finetuning the raw pretrained network yields no dimensionality decrease. These results suggest that the success of our rescaling method may indeed be attributable to pushing the network into the nested feature selection regime.NeurIPS Paper Checklist• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Note that the need for this procedure is an idiosyncrasy of the diagonal linear network setup, and further is unnecessary if γ = 0. Following pretraining, for each input dimension, either the positive or negative pathway weights will be zero, so setting both pathway parameters to equal the sum across pathways has the effect of copying the nonzero value over to the zeroed-out pathwayWe note that we had observed this behavior before deriving the PT+FT regularization penalty for ReLU networks. Unlike all other described experiments (for which we derived the described predictions from inspecting the norms before running any simulations), this is therefore a postdiction rather than a prediction.weighting the importance of each unit to the k-means objective by the weight of its contribution to the network's input-output function, specifically the magnitude of the product of its associated input and output weights 


Decision: The main contribution of this paper is a theoretical study of the implicit regularization in multi-task learning, comparing it to a pre-training and fine-tuning schedule.  The authors study a simple toy model of diagonal linear network, showing that MTL and PT+FT display different implicit regularization, where PT+FT benefits from features correlated between the main and auxiliary task.

While there were some concerns regarding the applicability of the theoretical results derived for the simple toy model studied in the paper, reviewers agree that the results in the paper are novel and that the setting studied in the paper and the implication of the results are very interesting. I therefore recommend that this paper is accepted.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: A good paper overall. See some questions below.
questions: 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.
limitations: limitations: Work is self-contained and aims to explain; the limitations of the theory in explaining certain cases are raised and addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See in Weaknesses.
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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)?
limitations: limitations: Limitations have been mentioned in the weaknesses section.

Comments: title: Seconding the opinion to keep diagonal network in main paper not Appendix
Thank you!
A Small Contrasting Opinion Re Linear Diagonal Networks
Thank you for the response
First Response
Additional responses
Additional responses
comment: We are glad that our rebuttal addressed the reviewer's concerns and thank them for increasing their score.

We also appreciate their note on the diagonal linear network section (suggesting that it should be moved to the appendix), as well as the input on this question by reviewers 969G and FiYT (who suggest it is a useful part of the main article, for the reason that we provide in the rebuttal). We really appreciate everyone's input on this. We have decided to keep the diagonal linear network section in the main text. However, we will clarify in more detail the role of these networks in connecting the theory in our paper to existing lines of work in deep learning theory. Note that in response to a suggestion by reviewer Ry6c, we will move our definitions of the diagonal linear and ReLU networks to a dedicated section ("Theoretical setup") before describing our theoretical results. We will add to this section a detailed explanation of the role of diagonal linear networks in our paper.

We believe that this will make our paper more broadly accessible: for readers who are interested in these networks as a theoretical stepping stone, we present diagonal linear networks in an integrated manner with our other findings; for empirically minded readers, we clarify early on why we focus on these networks and make it as easy as possible for them to focus their attention on ReLU networks and large-scale neural networks (if they wish to do so). We will also make sure to provide more justification on the assumptions used in Prop. 3 and make our figures bigger (using the extra page allotted to the camera-ready version).

We want to thank all the reviewers for providing their input on this question, as this was a very helpful discussion.

**Response to typos**

> Line 87 and 88: What is $\ell_1$ and $\ell_2$ norm of $f$ are these typos? It seems like you're taking an $\ell_2$ norm of a function which doesn't really make sense... 

Thank you for pointing this out. Our intention in these lines was to define the $\ell_1$/$\ell_2$-norm of these functions in terms of the $\ell_1$/$\ell_2$-norm of the linear coefficients (as we consider linear functions). However, we agree that this is confusing and will replace $\|f\|_{\ell_1}$/$\|f\|_{ell_2}$ by $\|\beta\|_1$/$\|\beta\|_2$.

> Line 85 has a switch in the citation style.

Thanks for pointing this out.

Once again, thank you very much for your response as well as your original review.
Thank you for your response, and thank you again for your helpful review. We agree that MTL could encode correlated features in a network with more than one hidden layer, which could be an interesting direction for future work.
Thank you again for your helpful review, your positive assessment of our work, and your literature suggestions!
I will second the opinion of Reviewer 969G here. My initial impression was the same as Reviewer 4YZg on the use of linear diagonal networks as the initial proof of concept. My main issue is that these are very contrived networks with specific properties that are not used at all in contemporary deep learning. However, the more I read about them the more I saw that they still act as a stepping stone in theoretical studies and are an excellent starting point for explainable deep learning re: feature reuse. After reading the entire paper several times I think the LDN section has a place in the main paper not appendix.
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!
I want to note that I differ slightly from Reviewer 4YZg on the recommendation about moving linear diagonal network results to Appendix -- I think as the authors say, including the results gives an insight starting from strong theoretical backing with limited assumptions that they then show in increasingly realistic settings, albeit with less formalism. As a theorist, I value the range of models in which this insight is presented; it not only substantiates the authors' results from multiple epistemic perspectives but also is encouraging to theorists to see that the simple models we study corroborate insights derived via different tools (importantly, in different settings than in which the model was initially developed)! Including these results in the main body only strengthens the narrative, in my opinion. 

I could imagine that for an empiricist reading this paper, the results in the extremely stylized model may not not provide as much value. I encourage the authors to think about their target audience(s) in making the decision of whether or not to move the results to the appendix.
I thank the authors for their rebuttal which have clarified many of my questions and addressed the concerns raised. While I still have some reservations about the presentation, I think the results are interesting and I trust the authors will make the necessary changes in the camera-ready, therefore I will raise my score to 5 (Borderline Accept).

Aside: I still think that while the results on diagonal linear networks are neat and clean they don't really seem necessary and they distract from some of the more interesting points in the paper. I would really suggest moving all of the diagonal linear net stuff like Corollary 1/2 and diagonal linear experiments to a section in the Appendix. I would use the extra space to make these figures bigger, provide more justification on some of the assumptions being used in Prop. 3 and move some of the experiments from the Appendix to the main body.

**Typos**
- Line 87 and 88: What is $\ell_1$ and $\ell_2$ norm of $f$ are these typos? It seems like you're taking an $\ell_2$ norm of a function which doesn't really make sense...
- Line 85 has a switch in the citation style.
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.
Below are additional responses to your questions and suggestions that we did not have sufficient space for in the main rebuttal.

> - The notation is very loaded and makes the paper difficult to parse. For example $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

Thank you for highlighting this. We’ll change the notation to denote hidden weights (i.e. $w^{(1)}$) by $w$ and readout weights (i.e. $w^{(2)}$) by $v$, in addition to a few other changes that aim to clarify notation a bit. We also appreciate the suggestion on the figures — we’ll make them bigger and agree that replacing some of the color bars by legends would be helpful.

> When finetuning are the weights of both layers trained or only the second layer weights?
> On line 178 what is LP in PT+FT(LP)?

Finetuning means that both layers are trained whereas LP stands for linear probing, meaning that only the readout weights are trained. Importantly, this means that finetuning is able to perform feature learning whereas linear probing cannot. We will add a sentence clarifying this to the manuscript.

> 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?

Yes, exactly. Would “# Non-zero dims.” and “# Units” be clearer?

> Line 176 what do you mean by "overlapping dimensions"? You mean overlapping features?

Yes. We will change this to “overlapping non-zero dimensions (i.e. features)” (as we want to make clear that in this case, the different features are simply different dimensions of the input).

Again, thank you very much for your review!
Below, we are providing some additional responses that we did not have space for in our rebuttal:

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

Thank you for pointing this out! We will change "large/small initialization" to "large/small initial weight magnitude".

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

We agree that the notation is overly complicated and will clean this up. Specifically, we will now introduce the two output dimensions by the superscripts “aux” and “main” from the beginning.

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

Thank you for highlighting this! We’ll make the notation here more conventional and will add a separate remark early on explaining the alternative parameterization by the magnitude $m$ and unit direction $\theta$ (as this parameterization is important for characterizing the penalties we derive). We’ll also make sure to properly define the L1,2 norm.

> Lines 128-129: We are not told how we get to (6) and not appendix section is referenced

This was an oversight on our part. We will add an explanation to the appendix detailing how this follows from Azulay et al. Specifically, we derive this results in two steps: first, we note that if after pretraining the network has the effective linear predictor $\beta^{aux}$, the first hidden layer has the weights $\sqrt{\beta^{aux}}$, where the square root is applied element-wise. Having set the readout initialization to $\gamma$, we then apply Theorem 4.1 in Azulay et al. Note that we use a slightly different functional form of $q$ to enhance readability.

> 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?

To clarify, the single-output ReLU network is obtained by setting $O=1$ in Eq. 4. However, as noted above, we will change our explanation of this setup to make it clearer. The "$\in\mathbb{R}$" is a typo; thank you for pointing it out.

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

Our theorem applies to arbitrary vectors $\gamma$. In practice we use a randomly sampled readout with a variance of $10^{-3}\sqrt{2/H}$. We will add a sentence clarifying this point; thank you for pointing this out.

> Figure 2: I would define STL as single task learning for the less familiar reader

We agree.

> Lines 152-155: More explanation for the ReLU would help. By “sparse” number of (hidden?) units we mean a low dimensional hidden layer?

Yes, i.e. a low number of hidden units. We will clarify this --- thank you for pointing this out!

> Lines 158-165: Everything is clear save how we arrived at c. Appendix reference?

We should have clarified that this is a direct application of an analysis in Woodworth et al., and will add such a clarification.

> 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 223-250: Again, the elegance of the experiment set-ups to reinforce the nested feature selection regime.

Thank you very much for these (and the other) positive comments! It’s useful for us to get this kind of feedback, as it helps us know when our explanations/experiments are clear.

Again, thank you very much for your helpful review!