Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse
Authors: Samuel Lippl
Jack Lindsey
Venue: NeurIPS 2024 poster (2024)
Keywords: multi-task learning
implicit regularization
finetuning
pretraining
implicit bias
Primary area: learning_theory
TL;DR: Analysis of the inductive biases associated with pre-training+finetuning and multi-task learning

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.

Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse
Samuel Lippl
Jack Lindsey* Anthropic
Center for Theoretical Neuroscience
Columbia University New York, NY samuel.1ippl.e aamuel.lippl..eolmmiii...c
San Francisco, CA
backullinddey@@mail.ccomm
Abstract
Neural networks are often trained on multiple tasks, either simultaneously (multi- task learning, MTL) 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 downstream task with fewer samples. Despite the prevalence of this approach, the inductive biases that arise from learn- ing multiple tasks poorly characterized. In this work, address this gap. We describe novel implicit regularization penalties associated with MTL and PT+FT in diagonal linear networks and angle-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 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 prior work, which biases to rely on a sparse subset of the features learned during pretraining. This regime is much narrower for MTI. 3) PT+FT (but not MTL.) 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, 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 nested feature selection. Our results suggest that nested feature selection may be important inductive bias for finetuning neural networks.
Introduction
Neural networks often trained multiple tasks, either simultaneously ("multi-task learning." henceforth MTL, see (12) or sequentially ("pretraining" and subsequent "finetuning," henceforth Empirically, models can transfer knowledge from auxiliary tasks improve performance on tasks of interest. However, theoretical understanding of how auxiliary tasks influence learning and generalization is limited.
PT+FT,
Auxiliary tasks are especially useful when there is less data available for the target task. Modern "foundation models," trained data-rich general-purpose auxiliary tasks (like next-worg word prediction or image generation) before adaptation to downstream tasks, are a timely example of this use case 5
Equal contributions.
WWrk primarily conducted while at the Center for Theoretical Neuroscience, Columbia University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 [3]. Alternatively, some others characterize the effects of explicit parameter regular- ization (e.g. an Eg penalty on weights) on the inductive bias of networks towards learning certain functions [14 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 MTI. 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 ReL.U 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 B).
We find that both MTL and PT+FT are biased towards solutions that reuse features between tasks, and that rely
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 a.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 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.
2 Related work
A variety of studies have characterized implicit regularization effects in deep learning. These include biases toward low-frequency functions [17]. stable minima the loss landscape [18). low-rank
solutions [9). 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
Fi norm, an ciinite-dimensione analogue of the E 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 Fi
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 +T 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 ( 25).
Others have characterized the learning dynamics of deep linear networks trained from nonrandom initializations, which can be applied to understand finetuning dynamics (76 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 (18 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. investigate generalization on a finetuning task), and, similarly, our tools could he 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 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 (2) that studies the inductive bias of a simplified "diagonal linear network"
model. Diagonal linear networks parameterize linear maps f R RR
(z) = B(T)) It, Bu(wi) == w (2) (1) + w
(1)
where ((T) E R... These correspond two-layer linear networks in which the first layer consists of one-to-one connections, with duplicate + and pathways to avoid saddle point dynamics around i = O. Woodworth el al. (38) showed that overparameterized diagonal linear networks trained with gradient descent on mean squared error loss find the zero-training-erve solution that minimizes ||ā||, when trained from large initial weight magnitude (the "lazy" regime, equivalent to ridge regression). When trained from small initial weight magnitude, networks instead minimize || (the "rich" regime). This bias is a linear analogue of feature learning/feature selection, as a model with an l penalty tends to learn solutions that depend on 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 4-- regularization. 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. contrast, the impact of explicit weight regularization has been characterized for both multi-output diagonal linear networks [16] and multi-output ReLU networks [44.
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 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 proxy for understanding the more theoretically complex problem of implicit regularization.
the
3 Implicit and explicit regularization penalties for MTL and PT+FT
3.1 Theoretical setup
Architectures. First, we consider diagonal linear networks with hidden weights w., T_ E RD and O E {1,2} output weights v4, ROXD, (0 is or 2 depending on the training paradigm, see below.) The resulting network function defined as
(T) = Bww, v)æ,
Ba(w, v) = J4.o o W4 J_ w_ B(u,)) v) E RXD
(2)
Second, we consider ReLU networks with H hidden neurons, hidden weights w E "xD and O E (1 readout weights E ROOH The network function is defined as
fu Em= Tn ((WT, #))+;
where (.)+ 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 a 0 and the readout weights by L. It will therefore be useful to consider a re-parameterization in terms of the
magnitude m, == ux || ll2 and direction 0, == / Under this re- ce-parameterization, an equivalent definition of the network function is given by
(4)
f D mn mnl10r. x))+ (z).
(5)
3

---

b
Auxiliary
Cor:::1
Corr:: 0.99
Corr.: 0.9
Cor::
task feature coefficient
10
Norm
Norm
Norm
0.001
10.
1000
10
10:9-
10
10
0.1
10
10
10
10
Magnitude
Magnitude
Magnitude
Figure I: Theoretically derived regularization penalties. a, Explicit regularization penalty associated with multi-task learning. b, Implicit regularization penalty associated with finetuning in diagonal linear networks. c, Explicit regularization penalty associated with finetuning in ReLU networks. This penalty also depends on the changes in feature direction over finetuning (measured by the correlation between the unit-normalized feature directions pre vs. post finetuning).
Training paradigms. We consider two datasets: an auxiliary task X xouz RRow and main task mannn
R Poaz X.D grue
IR "Rmaa XD Jmain E IR Rmmein
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 g := 31 and gmain := 32 in the diagonal linear network and yrur := T2 and Jmann := Ug in the ReLU network.
Second, 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 and y and the parameters learned after finetuning by wain and umain We define 3main gauz (for the diagonal network) and main Orain m aur faur (for the ReLU network) analogously.
3.2 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 minimizing the l parameter norm as an approximation of the implicit bias of training with gradient descent from small initialization. We argue that this is reasonable heuristic. First, the analogous result holds the single-output case for infinitesimally small initialization and two layers (though not for deeper networks, (38). ) Second, for cross-entropy loss it has been shown that gradient flow on all positively homogeneous networks (including diagonal linear networks KKT point of max
and ReLU networks) converges Finally, explicit & parameter norm regularization ("weight decay") is commonly used practice, making its inductive bias important to understand in its right well.
hatpammmmmmnparatterrreee norm objective [12).
We now derive the norms minimized by explicit regularization in MTL:
Corollary For the multi-output diagonal linear network defined Eq. 2 a solution f3* with minimal parameter norm || + ||el subject to the constraint that it fut the training data mXmas Emein ipnain X Xaux gouz guuz ) also minimizes the following:
8 arg min (2E2) (8guz)z (5mmainy2
S.f. Xmain Amain -main Xaux Jaur guur
This norm is known as the group lasso [47] and denoted as 11. For ReLU networks, by an argument analogous the one above, parameter norm minimization translates to minimizing E (mmauzy + (mgaainy (see Appendix A. and [44 46)).
The norm ||- I11.2 is plotted in Fig. 11. We analyze its impact in Section
3.3 Regularization penalties in finetuning
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 gemar from the auxiliary
task. After pretraining, the second-layer weights (v4, and v-) are reinitialized with constant magnitude
y. Further, to ensure that the network output pre-fineture feeeenntunnngg is zero (as in [3 we set the values of
corresponding positive and negative pathway weights to be equal to their sum following pretraining"
The network weights are further optimized to convergence on the main task dataset (XX main 1 Jmain))
The dynamics of the pretraining and finetuning steps can be derived as corollary of [38 40
Corollary 2. If the gradient flow solution Jaux for the diagonal linear model in Eq. I during
pretraining fis the auxiliary task training data with zero error (i.e. X Jour jaus), , and following reinitialization of the second-layer weights and finetuning, the gradient flow solution 3 after finetuning fis the main task data with zero training error (i.e. Xmoin Jmain = sJmain) ), then
8
= arg min ||
[.mmainn
lle
S.L.
xB =
Bmain
D
2,3main
Imain
l1
(13g
+
+g2)
0==)=2 2- v + azrrsssin(222)
Baur
d=1
This corollary is proven in Appendix A.2 The norm is plotted in Fig. Dp.
ReLU networks. We assume that after pretraining, the readout layer is re-initialized with arbitrary new weights 7 E "R We then characterize the solution to the finetuning task that minimizes the le norm of weight changes from this initialization. This is similar to the explicit weight regulariza- tion 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 theoretical analysis (which we then validate using experiments in Section[
We show that this finetuning solution implicitly minimizes the following penalty:
Proposition 3. Consider single-output ReLU network (see Eq. 3 with first-layer weights wh e RDD after pretraining, and second-layer weights re-initialized to 5 E RR The solution to the finetuning task that minimizes the & norm of changes in the weights, iee. minimizes E | w 1/ + (vh- - -²h²², is equivalent to the solution that minimizes
R(@main .m main gaax m maa)
Σ
(fmain maain (0h.u mm 0212)
-mh
h=1
rOmmain mmmain (0hu0 mh mp"") == (mg. //"" - 72) + (u²)² + mg aux 2u- V 1
aux (fmain Honay
where is the unique positive real root of
-(m + 7ymm main mh
aux (ōrmain Jouzy yu² 3 + u =.
(6)
We 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 Ph. == (0)main Geury 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. Tp.
4 Implications of the theory: multiple regimes of feature reuse
4.1 Sample efficiency in teacher-student models
To validate these theoretical characterizations and illustrate their consequences, now perform experiments in a seacher-student setup. In the diagonal linear network case, we consider a linear
Note that the need for this procedure an idiosynerasy of the diagonal linear network setup, and further unnecessary if = O. 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 pathway
5

---

MTL
PT.FT
Full averlsp
20/40
30/40
Non-zero dims
Gen. loss
10:9
10"1
Gen.
10
10:8
PT+FT
MTL PT+FT
STL (LP)
STL
Gen. (LP)
10
10.
32
256
32
256
32
256
32
256
32
256
32
256
Samples
Samples
Samples
b
d
f
STL
MTL
PT
Full overlap
3/6
5/6
# Units
Gen. loss
0.1
6
loss
0.1
MTL
PT.FT
Gen. loss
0.1-
STL
PI+FT (LP)
10
PT+FT (NTI)
STL (UP)
STL (NTK)
32
256
32
256
32 256
32
256 # Samples
32
256
# Samples
Samples
Figure 2: PT+FT and MTL benefit from feature sparsity and reuse. a.b, Generalization loss for a) diagonal linear networks and b) ReLU networks trained on a) a linear model with distinct active dimensions and b) a teacher network with distinct units between auxiliary and main task (STL: single- task learning). MTL and PT+FT benefit from a sparser teacher on the main task. c,d, Generalization loss for c) diagonal linear networks and d) ReLU networks trained teacher model sharing all features between the auxiliary and main task. PT+FT and MTL both generalize better than STL. e,f, Generalization loss for e) diagonal linear networks and f) ReLU networks trained on teacher model with overlapping features. Networks benefit from feature sharing and can learn new features.
regression task defined by i R 1000 with sparse set of l non-zero entries. We sample two such vectors, corresponding to "auxiliary"" and "main" tasks, varying the number of non-zero entries hauz and kmain and the number of shared features (overlapping non-zero entries). We then uniformly sample input vectors 7 E R 1000 from the unit sphere, using the ground-truw 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. -iweeennniionll hidden layer) and different kinds of overlaps (e.g. shared, correlated, orthogonal features) between the auxiliary and main task. We randomly sample input data o' E R15 from the unit sphere and use the teacher network to generate the target. We train 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 10-3 /2/H.
4.2 PT+FT and MTL benefit from sparse and shared features
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, consider the limit 18gl
1 +y² ac (capturing the case of a feature not
used during pretraining). this limit, the MTL penalty converges to 2|3a]- Similarly, for finetuning in diagonal linear networks, the penalty converges to c|ßa where c ~ O (log (1/()Sg"a + +)))) (per an analysis in (38). For new features, both networks therefore have an & norm minimization bias, suggesting that they tend to learn a sparse set new features (just like in the single-task case).
To test this insight, we consider stacher-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. 21). 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 teacher- student network with six auxiliary task features and one six uncorrelated main task features. Again, both MTL and PT+FT have lower generalization loss for fewer features (Fig. 2p).
Feature sharing. Next, we consider the opposite limit, i.e. large pretrained features:
[3al
- 0.
L H+y
22
In this limit, both the MTL and the PT+FT penalty for diagonal linear networks converges to

---

a weighted &g 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. 2 Notably, they perform similarly to a network where we only finetune the second layer (PT+FT (Linear Probing. I.P), which exactly implements the weighted & bias.
In the case of ReLU networks, we considered 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. 2
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. (b). On the one hand, both PT+FT and MTL outperformed STL, indicating that they were able benefit from the pretrained features. On the other hand, they also performed better than the FT+ (LP) model which only finetuned the second layer (and therefore did not have 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 FT+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 & norm for small auxiliary features and a weighted E2 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 difference in the limit behavior for diagonal linear networks: the relative weights of the e-- and weighted & f2-penalty are different between MTL and PT+FT. In particular, in the & penalty limit, there is an extra factor of order O (log (1/()S2uz + 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, 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)
4.3 A conservation law
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 given penalty P( ,32m0) (where graain is the main task feature coefficient and f is the auxiliary task B P feature coefficient): 1) the "É-order," Blg Blgg Sraam which measures, locally, how strongly P changes
with increasing grain and 2) the "feature dependence" (FD), Flog Squs, which measures, locally, how much P decreases for a larger auxiliary feature. In the previous section, we found that for ( gauz 0, ê-order -> and FD 0, i.e. the penalty becomes fyllike and does not depend on the magnitude of the auxiliary feature. In contrast, for gauz + 1, l-order -> 2 and FD - -1, i.e. the penalty becomes likee and depends inversely on the auxiliary task feature coefficient magnitude.
Blog P
How are torder and FD related for intermediate values of Sguz? 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, t-order FD 1.
See Appendix A.4 for proof.
Thus, there is 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).
7

---

MTL
PT+FT
PT-FT (0.1) PT+FT (0.01)
Order of penally
Gen. loss
Non-zero entries
Gen. loss
MTL PT+FT
10
10-9
10
10:45
0.001
10
32
256
32
256
32
256
32
256
Auxiliary feat.
# Samples
Samples
Order of penalty a
MTL
PT-FT
PT-FT (0.1) PT (10)
2
Con
Gen. loss
0.1
Units 6
Gen. loss
0.1
T
10-1
10
0.001
256 # Samples
258
Auxiliary feat.
Samples
Figure 3: PT+FT (much moreso than MTL) exhibits nested feature selection regime. a-c, Diagonal linear networks. a, --order/feature dependence plotted for Greein = 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, f-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 the auxiliary task and subset of those units on the main task. f, Generalization loss for PT+FT models whose weights are rescaled before finetuning.
4.4 The nested feature selection regime
Proposition 4 predicts the existence of 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 subset of the auxiliary task features, i.e. Amain I kaaz, and the number of overlapping units is equal kmain- Solving this task most efficiently involves "nested feature selection," " bias towards feature reuse and towards sparsity among the reused features. We plot É-order and FD for 3 find that for gaue z 1, both norms exhibit "lazy regime"- behavior (f-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: kimain decreases, the networks' sample efficiency does not become substantially better (Fig. 3D).
and varying the auxiliary task feature coefficient (Fig. 33)) and
However, for features with auxiliary task coefficients that are moderately smaller than their main task coefficients, Fig. Ba suggests a broad regime where the f-order closer to I (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 factor less than 1. In line with the prediction the theory, performing this manipulation enables PT+FT to leverage sparse structure within auxiliary task features (Fig. 3b), even while retaining their ability to privilege features learned during pretraining above others (see Fig, D. By contrast, this regime is much narrower for MTL (Fig. 3p.
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 (fmain (ggue) between the post-pretrainin and post-finetuning feature weights (Fig. Bl. For features that are fully aligned before and after pretraining, this penalty is again like for m xe 1. However, in most cases, these features change direction at least little bit, and we find that in that case the penalty more like while still remaining sufficiently feature- deauurede This suggests that 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 FT+ (even without any weight rescaling) was able to benefit from the nested sparsity structure, but MTL was not (Fig. 3-) Performing weight rescaling in either direction following pretraining uncovers the
We 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 postdiction rather than prediction.
8

---

initialization inneisstiiiinssiiivee (FD near 0), sparsity-bissed (é-order near 1) rich / feature-learning regime and the mitializaaiion-iase (FD near 1), no-sparsity-boe -iias (e-order near 2) lazy learning regime (Fig. . This suggests that for different architectures and different tasks, different rescaling values may be required to enter the nested feature selection regime.
4.5 PT+FT, but not MTL, in ReLU networks benefits from correlated features
The regularization associated with PT+FT yields benefits even when main/auxiliary task direc- tions are correlated but not identical (Fig. [ In contrast, MTL cannot softly share features as it encodes correlated features with entirely dis- tinct 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. 4b.. Notably, if fea- tures are identical, MTL outperforms PT+FT in the low-sample setting (Fig. Dtd. Thus, PT+FT (compared to MTL) trades off the flexibility to "softly" share features for reduced sample- efficiency when such flexibility is not needed.
b
STL
MTL PT STL
FT (Corr.: 1)
PT+FT Corrr: 0.9)
PT (Corr: 0.8)
Corr: 0.9
Mag.:
Mag.: 0.1 Mag.: 0.01
Gen. loss
Gen, loss
0.1
10
32
256
32
256
32
256
32
256
# Samples
Samples
Figure 4: 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, Gener- alization loss for main task features with varying correlation and magnitude (mag.). FT+ only outperforms STL if the features are either identical in direction or identical in magnitude.
The analysis in Fig. 34 would predict that ReLU networks can no longer benefit from correlated features if the magnitude of auxiliary task fea- tures is much higher than that of main task fea- tures. 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 perfor- mance for the case of identical but not correlated feature directions, confirming our prediction.
5 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 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+ improve sample efficiency compared to single-task learning (Fig. 5l)).
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 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 (5T] (see Fig. I0).
To facilitate comparison of the phenomenology in deep networks with our teacher-student experi- ments 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- posttfinttnning 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 shured 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: expect the dimensionality of network after finetuning to be

---

b
Scaling
Scaling: 0.125 Scaling:
Scaling:
Accuracy
0.5
Accuracy
0.5
20
20
20
0. 025
Pre PR Post
0.1
STL
0.1
0.5
0.25
10
10
ENSD
20
500
500
Samples
# Samples
Layer
Scaling
Scaling: 0.125
Scaling:
Scaling:
Accuracy
Accuracy
0.20
MTL PT STL
0.20
0.125
20
PR Pre PR Post ENSD
0.25
2.5
40
0.05
0.05
0.5
0.0
10
20
500
500
Samples
# Samples
Layer
Figure 5: Experiments deep neural networks trained on CIFAR- 100: a-c, ResNet-18, d-f, ViT. a.d, Accuracy for MTL, PT+FT, and STL in a) ResNet-18 and d) ViT. b,e Accuracy for PT+FT with weight rescaling in b) ResNet-18 and e) ViT. The participation ration of e) ResNet- and f) ViT's layers before and after finetuning (PR Pre and PR Post) as well as their ENSD.
lower than after pretraining (PR(XFT) < PR(XPT)) and for nearly all of the representational dimensions expressed by the network post-finetuning to be inherited from the network state after pretraining (EN (EPTT X PR(XFr)). We validate that this description holds in nonlinear teacher-student experiments with networks in the nested feature selection regime (Fig. [T]]
Remarkably, we find that the ResNet-18 exhibits the same phenomenology, but only for weights rescaled by a small value (Fig. 5 This supports the hypothesis that the observed benefits of rescaling indeed arise from pushing the network into the nested feature selection regime.
Finally, conduct the same experiment for Vision Transformers (ViT) (54). We confirm that PT+ improves performance over single-task learning, though MTL offers no similar benefit in this case (Fig. (1). We further find that rescaling before finetuning (both by larger and smaller values) decreases generalization performance (Fig. 5p)) Notably, our representational analysis reveals that rescaling by larger value yields a higher-dimese dineer-aimensionall subspace both before and after finetuning, whereas rescaling by smaller value yields
subspace, but pushes the effective number of shared dimensions down substantially (Fig. (5.. This indicates that a rescaling value of may already give rise to the nested feature selection regime and rescaling by 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 does not arise naturally.
Conclusion
In this work have provided 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 novel nested feature selection learning regime which encourages sparsity within features inherited from pretraining. This insight motivates 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 ReI.U networks (which assumed explicit parameterization) to the implicit regularization induced by dynamics of gradient descent, and extending our theory 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.
10

---

Acknowledgments
We are grateful to the members of the Center for Theoretical Neuroscience for helpful comments and discussions. The work was supported by NSF 1707398 (Neuronex), Gatsby Charitable Foundation GAT3708, and the NSF Al Institute for Artificial and Natural Intelligence (ARNI).
References
1. Vafaeikia, Namdar, K. & Khalvati, A Brief Review of Deep Multi-task Learning and Auxiliary Task Learming MEXiv:2007.01226 [cs, stat]. July 2020. http:://rxixs org/ abs/ 2007 01126 (2023).
2. Zhang. Y. Yang, Q. Survey on Multi- Task Learning. IEEE Transactions on Knowledge and Data Engineering 34. Conference Name: IEEE Transactions Knowledge and Data Engineering, 5586-5609, ISSN: 1558-2191 (Dee. 2022).
3. Du, Y., Liu, Z., Li, J. & Zhao, W. X. A Survey of Vision-Language Pre- Fre-Trained Models arXiv:2202. 10936 [es]. July 2022, http:://axiisi org/ abt/////..22222/////////// 10936 (2023).
4. Zhou, C. al. A Comprehensive Survey on Pretrained Foundation Models: History from BERT to CharGPT arXiv:2302009419 [es]. May 2023. http::/atxti......///// org/ abs htt///as.232200020 09419 (2023).
5. Bommasani, R. et al. On the Opportunities and Risks of Foundation Models arXiv:2108.07258 [es]. July 2022. http://arxiv org abs//////////////////aa//// 07258 (2023).
6. Jaderberg, M. al. Reinforcement Learning with Unsupervised Auxiliary Tasks 1rXiv:1611.05333 [cs]. Nov. 2016. http:://rxii org/ wwwaaa/s/ia. 161...111.... 05397 (2023).
7. Wang. A. & Russakovsky, O. Overwriting Pretrained Bias with Finetuning Data arXiv:2303.06167 [cs], Aug. 2023, /htt:://arrat
hwtgg/abs/23033. 06167 (2023).
8. Steed, R., Panda, S., Kobren, A. & Wick, M. Upstream Mitigation Is Not All You Need: Testing the Bias Transfer Hypothesis in Pre-Trained Language Models in Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (Association for Computational Linguistics, Dublin, Ireland, May 2022), 3524-3542. 202.022200.../2222222222222222 acl long 247 (2023).
9. Neyshabur, B. Implicit Regularization in Deep Learning 170.:1999900953 [es]. Sept. 2017.
httt:://arxiw.org/bbb//////
01953 (2023).
10. Gunasekar, S., Lee, J. D., Soudry, D. & Srebro, N. Implicit Bias of Gradient Descent on
Linear Convolutional Networks in Advances in Neural Information Processing Systems 31 (Curran Associates, Inc., 2018),
neur www.ineeurrippppeeeeeee /ttps//cpaapaee//////////// (2023).
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S. Srebro, N. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research 19, 2822-2222220222 2878. ISSN: 1532-4435 (Jan. 2018).
12. Lyu, K. & Li, J. Gradient Descent Maximizes the Margin of Homogeneous Neural Networks 05890 [cs, stat]. Dec. 2020. /ttt:://arxiv org/
05890 (2023). 13. Chizat, L. & Bach, F. Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss en. in Proceedings of Thirty Third Conference on Learning Theory ISSN: 2640-3498 (PMLR, July 2020), 1305-1338,
arXiv:1906. 1999900000890
https / proceedinggg
/ mlr
press cises//v155/0 http//11//chh2a2 chwzat20a....a.mamaa html (2023).
14. Savarese, P., Evron. I. Soudry, D. & Srebro, N. How do infinite width bounded norm networks
look in function space?: 32nd Conference on Learning Theory, COLT 2019. Proceedings of Machine Learning Research 99, BYYFLLL (2023) (2019).
http: / /ww.www..... sccpus cww.comiiiiiiie coom/invard/reeodd
15. Ongie, G., Willett, R. Soudry, D. & Srebro, N. A function space view of bounded norm infinite width relu nets: The multivariate case. arXiv preprint arXiv:1910.01635 (2019).
16. Dai, Z., Karzand, M. Srebro, N. Representation Costs of Linear Neurul Networks: Analysis and Design in Advances in Neural Information Processing Systems 34 (Curran Associates, Inc., 2021), 26884 26896 https: htttps::/ neurips cc/ paper wwwess/pppeeeeeee files//ppper/2211/ cctpsper..pp....pee
hash/ haalee222000022200eeeefffaaaa... Abstrrat.ss.tt...aa.......aaa... html (2023).

---

17. Rahaman, N. et al. On the Spectral Bias of Neural Networks, en. https://openreviow. net/ (2023) (Sept. 2018).
Morum?id-r1gh22cS33
18. Mulayoff, R., Michaeli, T. & Soudry, D. The Implicit Bias of Minima Stability; A View from Function Space in Advances in Neural Information Processing Systems 34 (Curran Associates, Inc., 2021), (445ae334Baadcaaa00bccc11a222223 http:/9999btttattttom
https http:://poooededeeeeeee neurips cc pttpp/ppperr/2222 20ah//ahh/hhhhh hash (2023).
19. Huh, M. et al. The Low- Rank Simplicity Bias in Deep Networks arXiv:2103.10422 [es]. Mar.
2023. htt.:/ttt://arra http://xxiioox//////////oo /ttp//bb//2111122211 10427 (2023).
20. Refinetti, M., Ingrosso, A. & Goldt, S. Neural networks trained with SGD learn distributions of [cond- mat, stat]. May 2023.
increasing complexity urXiv:2211.
org/
abs/2211. 11567 (2023).
21. Razin, N. & Cohen, N. Implicit Regularization in Deep Learning May Not Be Explainable by Norms in Advances in Neural Information Processing Systems 33 (Curran Associates, Inc., 2020), 21174-21187, https ptwppeeiice......cooo neurips cc cc/ paper 2020/ hash $2lc55f800ff555cc444444eeece Aat.-G-bbttrattttttcoo (2023).
22. Boursier, E., Pillaud Vivien, L. & Flammarion, N. Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs en. in (May 2022).
wttfforrmmid-L44400 iUxQ1I (2023).
23. Maurer, A., Pontil, M. & Romera-Paredes B. The Benefit of Multitask Representation Leannnn
Journal of Machine Learning Research 17, 1-32. ISSN: 1533-7928. http:// .rg/
papers/v17/15 242. html (2023) (2016).
24. Wu, S., Zhang, R. Ré, C. Understanding and Improving Information Transfer in Multi-Tass Learning arXiv:2005 00944 [es]. May 2020. /ttp::/ htta:://axxiv.ooo org abs/2005 00944 (2023).
25. Collins, L., Hassani, H., Soltanolkotabi. M., Mokhtari, A. & Shakkottai, S. Provable Multi-Tane Representation Learning by Twy-Layer ReLU Neural Networks en. in (June 2024). https 7/openreview. ftte/forrmiidd-MUU
(2024).
26. Braun, L., Dominé, C., Fitzgerald. & Saxe, A. Exact learning dynamics of deep linear networks with prior knowledge. en. Advances in Neural Information Processing Systems 35, 6615-66629, https / neurips cc / /eeppaprr_. files paper 2022 / hash / 22230222511911109771111002222223
Abstract Conference html (2023)
(Dec. 2022).
27. Shachaf, G., Brutzkus, A. & Globerson, A. A Theoretical Analysis of Fine-tuning with Lin- ear Teachers in Advances in Neural Information Processing Systems 34 (Curran Associates, Inc., 2021), 15382-15394 https: pttp/eecceeennnseeo neurips cc caperrrr///pppapapppppooa / 2021 / hash / 62033d16dcceaaab313bbt htt.://99930baa/a7aaaaa-------- Abstract.stt.t.c..o..m html (2023).
28. Evron. I., Moroshko. E., Ward, R., Srebro, N. Soudry. How catastrophic can catastrophic forgetting be in linear regression? in Conference on Learning Theory (PMLR, 2022), 4028- 4079.
29. Lin, S., Ju, P., Liang, Y. & Shroff, N. Theory on Forgetting and Generalization of Continual Learning en. in Proceedings of the 40th International Conference on Machine Learning ISSN: -490-3498 (PMLR, July 2023), 21078-21001
https: //pooeedd//oc mlr hal..press/v222
1in23f. html (2024).
30. Evron, I. et al. Continual learning linear classification on separable data in International Conference on Machine Learning (PMLR, 2023), 9440-9484.
31. Goldfarb, D., Evron, I., Weinberger, N., Soudry. D. & HAnd, P. The Joint Effect of Task Similarity and Overparametterzaation. on Catastrophic Forgettigg- Analytical Model The Twelfth International Conference on Learning Representations (2024).
32. Lee, S., Goldt, S. & Saxe, A. Continual Learning in the Teacher-Student Setup: Impact of Task Similarity en. in Proceedings of the 38th International Conference on Machine Learning ISSN: 2640-3498 (PMLR, July 2021), 6109-6119, 1se21e.html (2024).
https: : ttttps///proce /ttp////cceeeeg....m mlr. press v139/
33. Lee, S., Mannelli, S. S., Clopath, C., Goldt, S. & Saxe, A. Maslow's 's Hammer in Catastrophic Forgetting: Node Re- Rse vs. Node Activation en. in Proceedings of the 39th International Conference on Machine Leurning ISSN: 2640-3498 (PMLR, June 2022), 12455-12477. https //proceedings. mlr. press//162/leee2gg..ot arrroeeiinggnl....pppps/////////
(2024).
12

---

34. Dery, L. M., Michel, P., Talwalkar, A. & Neubig, G. Should We Be Pre-training? An Argument for End-task Aware Training as an Alternative en. in (Oct. 2021). / (2023).
35. Weller, O., Seppi, K. & Gardner, M. When to Use Multi-Task Learning Intermediate Fine- Tuning for Pre-Trained Encoder Transfer Learning in Proceedings the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers) (Association for Computational Linguistics, Dublin, Ireland, May 2022), 272-282.
https: /atps://aalaanie aalannholny
org/2022. acl short 30 (2023).
36. Kumar, A., Raghunathan, A., Jones, R., Ma, T. & Liang, P. Fine- Tuming can Distort Pretrained Features and Underperform
urXiv:2202.1005: [cs]. Feb. 2022. http
7/arxiv. org/ abs./202/ 10054 (2023).
37. Kornblith, S., Shlens, J. & Le, Q. V. Do belter ImageNet models transfer better? in (2019).
https://arxiv. www.:/ogggppp.....e0 08974. pdf (2023).
38. Woodworth, B. et al. Kernel and Rich Regimes in Overparametrized Models en. in Proceedings of Thirty Third Conference on Learning Theory ISSN: 2640-3498 (PMLR, July 2020), 3635- 3673. https://proceed/::
mlr. halx.pres//vee /vaps//11220 www.ppas//////ooooorrrereree wwwp///owwwothaaa..ooo html (2023).
39. Pesme, S., Pillaud- L. & Flammarion, N. Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity en. in (Nov. 2021).
https: bectps://openrvviii net/
Forum?id=wviiRqqile
(2023).
40. Azulay, S. et al. On the Implicil Bias of Initialization Shupe: Beyond Infinitesimal Mirror Descent en. Proceedings of the 38th International Conference on Machine Learning (PMLR,
July 2021), 468-477,
mml wtw...press/vv3/// azulll////lllllaa.amm 1ttml .ttl (2023).
41. HaoChen, J. Z., Wei, C., Lee. & Ma. T. Shape Matters: Understanding the Implicit Bias of the Noise Covariance en. in Proceedings of Thirty Fourth Conference on Learning Theory (PMLR, July 2021), 2315-2351 (2023).
https: // proceedings. mlr. press vv34/ http//hachhhn1.. html
42. Moroshko, E. et al. Implicit Bias in Deep Linear Classification: Initialization Scale vs Train- ing Accuracy in Advances in Neural Information Processing Systems 33 (Curran Associates, Inc., 2020), 22182-22199 https: : http:/prooceeeee proceedings neur
cc/ wttppcppppprp..// / 2020/ / /hash/ (2023).
fc2022c69b61cc76befff78888 (ttp:///bbbbbb1.7111711166-6aaa- wwwt:oftffbbbtttttttoooc
43. Ji, Z., Dudík, M., Schapire, R. E. & Telgarsky, M. Grudient descent follows the regularization
path for general losses in Conference on Learning Theory (PMLR, 2020), 21 2109-2136.
44. Yang, L. et al. A Better Way to Decay: Proximal Gradient Training Algorithmss for Neural Nets en. (Nov. 2022). httpt::t/////t http::///pperevieww net/forunt 4xhh 1xhh (2024).
45. Parhi, R. & Nowak, R. D. Near-Minimal Optimal Estimation With Shallow ReLU Neural Networks. IEEF Transactions on Information Theory 69. Conference Name: IEEE Transactions
on Information Theory, 1 1125-11 140. ISSN: 1557-9654., hattps://eeexpllrr
ieee. org/
(2024) (Feb. 2023).
46. Shenouda, J., Parhi, R., Lee, K. & Nowak, R. D. Variation Spaces for Multi-Output Neural Networks: Insights on Multi-Task Learning and Network Compression. Journal of Machine
Learning Reseurch 25, 1-40. ISSN: 1533-7928.
http: : // org papers / v25 /3
0677 htmlmmlllp (2024) (2024).
47. Yuan, M. & Lin, Y. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society: Series B (Statistical Methodology) 68. _eprint: wtwpee/eiwiiaiiilllleoo/o//o///a
1461-1167777888111 000.068.00000005555 00532.x, 49 67. ISSN: 1467- 9868. https: http///nnnnnnnnrrrrrrroaa wiley. com/ /tsi/b////////e/e////////eaae/ba 1111/j 1467 9868 2005 00532. x (2023) (2006).
48. Lubana, H. S., Trivedi, P., Koutra, D. & Dick, R. How do quadratic regularizers prevent catastrophic forgetting: The role of interpolation in Conference on Lifelong Learning Agents (PMLR, 2022), 819-837.
49. He, K., Zhang, X., Ren. S. & Sun. J. Deep Residual Learning for Image Recognition in (2016). 770-778. https: / thecvf com/ content_cvpp__0000 2016/ http//tmmll////// He _Depp_ Residual Ciw.ar__inn_____2266666666 _apper_pprr... html (2024).
50. Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. Publisher: Toronto, ON, Canada (2009).
51. Deng, J. et al. ImageNet: large-scale hierarchical image database in 2009 IEEE Conference Computer Vision and Pattern Recognition (2009), 248-255.
13

---

52. Gao, P. et al. A theory of multineuronal dimensionality, dynamics and measurement en. Nov. 2017. https://w: biorxiv org/ content/tottoooooooooooooooo 111p///11110111111122222122222 (2023).
53. Giaffar, H., Buxó, C. R. Aoi, M. The effective number of shared dimensions: simple method
for revealing shared structure between datasets July 2023. https://wwat www.:..biixiv..ommmmo org/
content/10. 11111101222111 07 27 550815v1 (2023).
54. Dosovitskiy, A. et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1192: (2020).
55. Paszke, A. el al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems 32 (2019).
56. Korbblith, S., Norouzi, M., Lee, H. & Hinton, G. Similarity of Neural Network Representations Revisited en. in Proceedings of the 36th International Conference on Machine Learning (PMLR, May 2019), 3519-3529, https: http://pooceedigg..oe mlr. ppwas/v97 http://9///ooooblttt11..ee html (2023).
httpss//proceedngaa
14

---

A Derivation of regularization penalties
A.1 Multi-task learning in ReLU networks
Multi-task ReLU networks with a shared feature layer and outputs can be written as
H
(x)
=
Σ
Th (1w 17)+ = Σ
mn ( z))
(7)
h=1
mih = Uh || 2|,,, On - wn/ |
(8)
where Wh my IR. We interested in how minimizing the parameter norm Σk=1 | | ū7 l
maps minimizing norm over the solution weights m. For given my R the norm of the
input weight ||x2 2 is a free parameter Σn > 0, we can set Th == m. 2. As any value for 24, leads the same function, we choose 24 so to minimize
1 11
1 11 = +
11 12/22/
Setting the derivative to zero implies
2 = n 1|a-
(10)
As result,
H
H Σ
Σ
12
|0 11
Imalla
(11)
h=1
where the right-han hand side of the equation is the م12 norm.
A.2 Finetuning in diagonal linear networks
Corollary 2.
If the gradient flow solution 3aua for the diagonal linear model in Eq. 1 during
pretraining fits the auxiliary task training data with zero error (i.e. Xaux Jaus gonz), ), and following reinitialization of the second- layer weights and finetuning, the gradient flow solution 3* after finetuning fits the main task data with zero training error (i.e. X Xmain Gmain gmain ), then
3
arg min 1.maainll
s.t. X =
maain
2,3main
(12 Σ
+ +
g(2) = V arcsinh(z/2))
|Baua
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 zero and one pathway encodes the effective linear predictor gaur, Thus, the first hidden layer has the weights / /auz which yields the new weight magnitude along the positive and negative pathway (due to the summing operation we perform following finetuming, which ensures that the weights of both pathways are set /Bauz). Having set the readout initialization to 7, we then apply Theorem 4.1 in 440].
A.3 Finetuning in ReLU networks
Proposition 3. Consider single-output ReLU network (see Eq. with first-lyrr weights up RD after pretruining, and second-layer weights re-initialized to 7 RR The solution to the finetuming task that minimizes the lg norm of changes in the weights, i.e. minimizes Eh |! 1 112 + (Un - > equivalent to the solution that minimizes
M( - main yaur aaa)
rOmmain
main (0gaz .mh (21))
,thh
h=1
r@@main mm main (g mn m..") := (m momin /u² 72 + (2²)² + aur 2u² mm /m (
@rmain Haury
where u² is the unique positive real root of
- -(m Ammainga + mmmm mmmin mnh aux (fpmain σouz) u² =0. 0.
(6)
15

---

Proof. We define the norm of the deviation from the parameters at initialization as
R(w, 1 t 7)
U/ 7)² ||07H
wh wouuee
(12)
Note that w and parameterized in redundant manner, as multiply tp by constant 0 and obtain the same function as long as we divide üi by the same constant. To parameterize function in unique fashion, we describe it in terms of the normalized hidden weight 0% := / ll2 and the (signed) magnitude my := |, 1/. We wish compute the norm of the deviation from the initialization in terms of this parameterization. Specifically, we compute
R(0, m m ?7 == min R(w, aw 7) s.t. Vy, ui, w/ll2 my 1 1120 w.e
(13)
Note 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 hidden units, we can solve it for each hidden unit individually:
R(Ø m aa :7)
Σ
f(Oh, mn. mm/11y 1 ,7),
(14)
1=1
f(6h, my 22th
:7) = min (un, - 7)² + 112ih - 15% 12
(15)
1h.
s.t. 0n w/./ /1/n/11 ma = 11
(16)
We can express this as
F(O,, mA wh ouu
.7) min p(u), p(u) = (ma/u 1 -)) +) 3/7++1100 u>0
+11100 |uth Whh
(17)
Note that this function is smooth and diverges to do both for u - oo and u - O. As long as we have unique positive stationary point u", this stationary point must therefore be minimum. We simplify
puu) = (ma/u 7) + u² + ||
2u(O,. (7)".))
(18)
and, defining mi maur := || ll2, and the cosine similarity between corresponding pretraining and (ar) .ang
finetuning task weights Ph := | ( Ila 10g
write this as
puu)= = (ma/u 7)² + 2 + (marr)) - 2wm Ph-
(19)
We then determine the stationary points by computing the derivative
p'(u)= = 2(mma/u - 7)mm// /² + 2u 2m Pho
(20)
and setting it to zero. Simplifying this equation (and multiplying by 22/2) yields
-m + 7m - m, ph u + + =
(21)
We can compute the solution to this quartic equation, e.g. using np. roots and select the (unique, in all empirical cases explored) positive real root, u². We can then compute the original norm by plugging in u":
F(O,, my w .7) (ma/u" /u" -7)² + x² ( 2u m, 2u"m," ouz Ph-
(22)
A.4 Proof of Proposition
Proposition 4. For the MTL and FT+ penalties derived in both the diagonal linear and ReLU cases, t-order FD = 1.
16

---

Proof. Below, for simplicity's 's sake, we always denote gmaii
3, m main and gmarn = A.
MTL penalty. We first consider the MTL penalty, which is identical for diagonal linear networks and ReL. networks. We assume (without loss of generalization) mane > 0. In that case,
a bog
V(maue) )a m maur
fodder
log
(23)
77
m
m
(maur)a y + m² mauz V(maurya + m²
(maux)² m2² mar (maur)² m²
and
i log V(mauzja + m m 982
FD
① log maur
(maux)² m (mouz) )2 mm² (maur)) + m2 maar v (manay + m²
172
772 ammaa j2
1
V (maaz) )2 + m mou
m
(24)
Thus, --rddrr + FD 1.
Finetuning in diagonal network. We assume (without loss of generalization) $2
and derive
@ log (13g
28g
torrerr
Hour
2,3d
29 [
(25)
a logg 3d
(3aur))2
203 32
and
i log
1.9gue 1
28a
TA
FD
i Olog (3guxr)
(26)
2,3
28d (
q'
23a [,50*|
d
28
28d [
(3qu2) )2g
29
[18900)
Thus, f-order FD 1.
Finetuning in ReLU networks. We define
F(O,, m., u/th m, mm (ma/u (ma/u-7)² 7) + u² + mh guz 2 / m
(fn, (
(27)
where, without loss of generalization, we assume that mn aaz and note that
r(O,, m. 10 ,, ) F(On, m. u" 1 m auz ) where u is the critical point, depending 0n, m4, hauz m aue By the chain rule,
(28)
Ör (0,, my mm 18.u.
auzz m
OF mh, u Øaur mp Or(O,, mp, u 08 .ma. du"
(29)
Omn
Om
Ou
and the analogous statement is true for m que
However, by definition,
aF mh, 16 mg
(30)
Ou
Thus,
my OF
2m r(u))2
E-order
0mp
(31)
m que
OF
aux
Ph
1 auz
FD
amau aue
m
auz
mm
Ph
m

---

To prove the statement, we therefore must prove that
(u*) )2
2m 2 aua
+ m
PV
h
u
=
mi (u))²
+ (u²)²² + (m² - 2u m aaz Ph.
(32)
To do so, we leverage that 30 implies
m
+ +(u))² u² "m/
/mguu Ph. 0,
(33)
(u*) )²
and therefore
m
(u*)² u² m Ph. O. (u²)2
(34)
Plugging this into
32 yields
2m 2 u )2
mi
+
m
aue
PV m ous
u'
+ (2*)² maw
2p m u =r
(35)
(u
and completes the proof.
B Detailed methods
We trained all networks with PyTorch [55
B.1 Diagonal linear networks
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. a rough heuristic for determining the learning rate, we begin at learning rate of 10 and divide by ten whenever the loss exceeds 100 (indicating divergence).
B.2 ReLU networks
We train ReLU networks with 1000 hidden units using the same learning rate heuristic and train until the mean squared error has reached threshold of 10-8
B.3 ResNets
We consider ResNet-18 in the standard Py Torch implementation and train it with gradient descent with learning rate of 10-3 and momentum of 0.9.
B.4 Vision Transformers
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 learning rate 10-3 and weight decay 0.00005.
B.5 Reproducibility
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 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-boume estimate, the experiments therefore required 50,000 hours of CPU time and 60,000 hours of GPU time.
18

---

b
# Teacher units
0/40
20/40
35/40
40/40
1000
40
0.1
PT+FT
MTL
Loss
Loss
RT+FT (LP)
10"7-
10
PT4FT (NTK)
Scale of init.
STL (LP)
STL (NTK)
64
4096
0.001
32
4096
32
4096 32 ammples
4096 32
4096
amples
Figure 6: Larger-scale teacher-student experiments. a, Generalization loss of shallow ReLU networks trained data from a ReLU teacher network. b. Generalization loss for different numbers of overlapping features (out of 40 total) between main and auxiliary tasks. NTK indicates the (lazy) tangent kernel solution. This is comparable to Fig. 2H. except with more teacher units and more data.
C Robustness of main results to choice of number of auxiliary task samples and input dimension
To increase confidence that main results robust 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. 2 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. (h). Further, MTL and PT+FT tend to outperform STL if the features are overlapping and MTL tends to outperform PT+FT (Fig. 6p). In particular, the finetuned networks still benefit from feature learning, especially if some features are novel.
D Analysis of learned solutions in linear and nonlinear networks
Our theory predicts inductive biases towards solutions that minimize norms, often either fi-like (incentivizing sparsity) or lz-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.
D.1 Diagonal linear networks
To check whether the implicit regularization theory is good explanation for these performance results, we directly measured the f1 and Q norms the solutions Iearned by networks, compared to the corresponding penalties of the ground truth weights. In Fig. 77 we see that 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 -tuut 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. 7p by analyzing the overlap between the auxiliary task features and the learned linear predictor for the main task.
In Fig. 7 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, Fig. 7 we show that the LI norm of the solution increases with the training sample size, consistent with a bias towards minimization. There an interesting discrepancy between behavior of the sparsity of the solutions (nonmonotonic, see Fig. 7 and the L1 norm (largely monotonic, see Fig. 7t). This attributable to the discrepancy between the L. norm (which diagonal linear networks the rich regime are biased to minimize) and sparsity (for which LI norm is only a proxy).
19

---

b
L12
Overlap law tasks
1.0
1.00
Rel. penalty
10.0
3.0
MTL PT+FT STL
0.01
0.9
overlap
0.8
0.0
no
256
256
32
256
yes
Samples
ammple
STL
STL (LP)
1.0
Unexplained Variance
6
# Samples
256
32
2
0.0
200 0 # Weights
200
32
256
# Samples
0/40
20/40
30/40
40/40
1.0
0.5
Unexplained
Variance
#Samples
0.0
256 32
1.0
0.5
MTLL
0.0
0
Weights
0/40
20/40
30/40
40/40
Overlap b/w tasks
1.0
STL
0.5
STL (LP)
MTL
0.0
32
256
32
256
32
256
32
256
# #Samples
Figure 7: a, f,2 norm and Q penalty for MTL, STL, and PT+FT networks from Fig. (40/40 overlapping features case). b, Proportion of the weight norm in the 40 dimensions relevant for the auxiliary task, for the networks in Fig.Bp.c. Weight rescaling decreases this overlap. c, Proportion of variance concentrated in the top k weights, as function of k. for training on 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, LI norm for STL as a function of the number of samples, Proportion of variance across different overlaps and for different learning setups (see also Fig. 22)). 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. 2)
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 which uses 40 features). We find that learned solutions across range of overlaps sparse using single-task learning (see Fig. 7b)
20

---

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. 7 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.
D.2 ReLU networks
We adopt a clustering- based approach to analyzing the effective sparse structure of learned task solu- tions. Specifically, for a given trained network, we perform k-means clustering with predetermined value of clusters on the normalized input weights to each hidden-layer neuron in the network" We measure the extent to which K cluster centers are able 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 zero, we can say that (to good approximation) the network effectively makes use of at most nonlinear features.
D.2.1
Single-task learning: rich inductive bias yields clusters of similarly tuned neurons that approximate sparse ground-truti 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. Ba). For tasks with many (1000) units 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 the number of samples increased (Fig. (a)), at which point the network matches the teacher function very well. This bias towards sparser-thant ground-trutt 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 likee inductive bias is 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 li-like inductive bias is not exactly the same as 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 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 cluster 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.
D.2.2
finds sparse solutions and improves alignment of feature clusters learned during pretraining
Pretraining+finetue
We find that wewrainiingiinetnning 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+ 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. (8)). Moreover, we find that MTL also learns sparse solutions (Fig. [d). 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 magle-task-trainee networks, especially when main task samples
'weighting the importance of each unit to the k-means objective by the weight of its contribution to the network's 's onput-output function, specifically the magnitude of the product of its associated input and output weights
21

---

a
Units:
Units: 1000
3/6
5/6
Inertia
1.0
0.5
Samples
0.5
0.0
512
0
20
20
Inertia
0.0
# C lusters
6/6
Corr
Mixture
Alignment
Units:
Units: 1000
Scale of init 0.001
1.0
0.5
0.5
0.0
512
Samples
# C Clusters
d
0/6
3/6
5/6
6/6
Corr.
Mixture
Inertia
0.5
0.0
0
6
20
0
6
20
0
20
20
6
20
6
20
# C lusters
0/6
3/6
5/6
6/6
Corr.
Mixture
Alignment
1.0
MTL PT+FT
0.5
256
32
256
32
256
32
256
32
256
256
S Samples
PT+FT 0.1
PT+FT
PT+FT 10
MTL
PT+FT
MTL
Inertia
Weight scaling 10.0
Allgnment
1.0
Same Units no
0.5
yes
0.0
0.1
32 256
32 256
32
256
32
256
32
256
32
256
# #amples
# Samples
Figure 8: Analysis of effective sparsity of learned ReLU network solutions. Inertia (k-means reconstruction error for clustering of hidden-unit input weights) as function of the number of clusters used for k-means, for different numbers of main task samples and ground-truth teacher network units, in single-task learning. b Alignment score average alignment (across teacher units) of the best-aligned student network cluster uncovered via k-means. c, Inertia for networks trained using PT+FT for the tasks of Fig. Lie and Fig. T.. d, Same as panel but for networks trained with MTL. e, Alignment score for networks trained with MTL, PT+ and STL the same tasks as in panels and d. f Inertia (using & = 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 the auxiliary task (as in Fig. fe,,), using MTI. or PT+FT with various rescaling factors applied to the weights prior to finetuning. g Alignment score for the networks and task panel
limited (Fig. (e). This provides mechanistic underpinning for the relationship between the inductive bias of PT+ that we describe in the main text and its performance benefits.
D.2.3 Nested feature selection regime allows network to prioritize sparse subset of feature clusters learned during pretraining
In the main text we describe the "nested feature selection" regime, which occurs at intermediate values 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 subset of the features used in the auxiliary task. Importantly, they still maintain a bias towards reusing features from the auxiliary task, show for diagonal linear networks and ReLU networks in Fig. Here
22

---

MTL
PT+FT 1
PT+FT 0.1
PT+FT 0.01
overlap yes
LOSS
1006
10
256
32
256
32
256
# amples
MTL
PT+FT 10
PT4FT
PT4FT
Same Units
0.1
10
no
256
256
32
256
32
256
yes
# 5 Samples
Figure 9: Comparison between tasks with sparse main task features that are either subsets of auxiliary task features or new features. Networks are trained with MTL or with PT+FT, potentially with rescaling (as indicated by the number). a, Diagonal linear networks trained on five main task features. b, ReLU networks trained on teacher network with one feature. We see that MTL (to some extent) and PT+FT can benefit from such an overlap, but for small rescaling values, this benefit becomes smaller.
ResNet18, STL10
VGG11, CIFAR100
ReaNet18, CIFAR100
Accuracy (%)
100
100
100
PT+FT
PT+rescaling-FT
0
50
500
5000
500
5000
50000
500
5000
50000
Main task samples
Main task samples
Main task samples
ResNet18, STL10 (binary)
VGG11, CIFAR100 (binary)
ResNet18, CIFAR100 (binary)
Error (%)
Error (%)
Error (%)
2
2
5
no rescaling
rescaling
no rescaling
rescaling
no rescaling
rescaling
Figure 10: Results for finetuning deep convolutional networks trained on ImageNet, with/without weight rescaling (factor of 0.5) prior to finetuning.
show that networks in this regime (obtained most clearly in the shallow ReLU network when networks are rescaled by a value of 1.0 after pretraining) indeed learn very sparse (effectively I-feature) solutions when the ground-truth main task consists of 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 bias in their internal representation as they still need to maintain the features needed for the main task (Fig. (8, 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 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. (g.. This alignment benefit mostly lost when networks are rescaled by factor of 0.1 following pretrainning (a signature of rich-regime-like like behavior).
23

---

E Further evaluations of the rescaling method for finetuning
To evaluate the robustness/ general-purpose utliity of our suggested approach of rescaling network weights following pretraining, experimented with finetuning convolutional networks pretrained on ImageNet on downstream classification tasks: pretrained ResNet- 18 finetuned on CIFAR100, pretrained VGG11 I finetuned CIFAR 100, 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 found exposes performance differences more strongly). Due to computational constraints, we did not sweep over the choice of the rescaling factor, but simply used 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 networks
We 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. 5p, 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 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 most-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: [2] as measure of dimensionality. For an n x p matrix representing n mean-centered samples p-dimensional network responses, with a p x p covariance matrix Cx _X X, the participation ratio is defined as
(E/
2₂)
trace (Cx)²
PR(X)
Σ 21=1 X trace (C²)
brace (X [x))
(36)
trace (X XXXXXX
where \; are the eigenvalues of the covariance matrix Cx. The PR scales from 1 to 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 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. Tha..
Our measure of shared dimensionality between two representations is the effective number of shared dimensions (ENSD) introduced by Giaffar el al. (53). The ENSD for n X p matrix of responses X from one network and an n x matrix of responses Y from another network, denoted ENSD(X,Y) given by
trace (Y XXXXXX Y) trace (XTX) trace (YTY) trace (X XX X) trace (Y YY YY
(37)
This measure 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. 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.
24

---

b
Participation Ratio
Participation Ratio
Scale of init.
Units
0.001
1
512
32
256
32
256
32
256
Samples
Samples
d
PT PR
FT PR
PT/FT ENSD
PT
FT PT/FT ENSD
Weight scaling
Weight scaling
10.0
1.0
0.1
0.125
2
32
256
32
256
256
012345
012345
013455
Samples
Layer
Figure 11: Dimensionality of the network representations before and after finetuning. a, Participation ratio of the ReLU networks' internal representation after training task with six teacher units. b, Participation ratio the network representation after finetuning on the nested sparsity task with different weight rescalings. Participation ratio before (left panel) and after finetuning (middle panel) and the effective number of shared dimensions between the two representations (right panel). Small weight scaling decreases the participation ratio after training. d, The same quantities for ResNet18 before and after finetuning (see also Fig. 5.).
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 phe- nomena the nested feature selection regime. case in which the main task uses subset of the fea- tures of the auxiliary task, if the network truly extracts this sparse subset of features, expect the di- mensionality of network after finetuning to be lower than after pretraining (PR(X < P PR(X (PT)) and for nearly all of the representational dimensions expressed by the network post-finetuning to be inherited from the network state after pretraining (EN SENSS(XPT PT XFT) ) = PR(XFT)), By contrast, networks not in the nested feature selection regime should exhibit fz-like rather than -ikkee 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 pre- training), and also networks in the rich regime, exhibit decreased PR following finetuning (Fig. lTb). By contrast, lazy networks (weights rescaled by 10.0 following pretraining) exhibit dimensionality decrease during finetuning. Additionally (see Fig. ITb), the ENSD between pretrained (PT) and finetumed (FT) networks is almost identical to the dimensionality of the finetuned representation (PR FT).
Strikingly, we observe very similar behavior in our ResNet-18 18 model pretrained on 98 CIFAR-100 classes and finetuned on the 2 remaining classes (Fig. ITH), 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.
25

---

NeurIPS Paper Checklist
1. Claims
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 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 and find that they relevant for large-scal scale neural network architectures 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 long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in the latter half of Section
Guidelines:
The answer NA means that the paper has 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
26

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We derive 1) the penalty associated with MTL in ReLU networks in Ap- pendix A..] 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 include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proof' 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 short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that 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 reproduce all figure in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, 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 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 novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, 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 good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
27

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have attached codebase that enables reproduction of all figures in the main text.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ //.pps.cc/publi//ggug guwdes/ CoodeSummis ssi casssnonnccyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide these details in Appendix
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 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.
7. 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.
28

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We describe these resources Appendix B.5
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httt:p/://errrps..cc/p/b.///cece
Answer: [Yes]
This manuscript presents foundational research without direct path to negative applications. Further, the manuscript does present any experiments involving human subjects.
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This manuscript presents foundational research without direct societal impact.
29

---

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), fairesss considerations (e.g., dep.oymmly 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 tied to particular applications, let alone deployments. However, if there is 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, 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 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 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 MI).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 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, were unable to find associated license. For Imagenet, the data is available for free to researchers for con-ommerccaal use. Both websites hosting these datasets request a citation of the relevant paper which 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 dataset.
30

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperssi thoode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's ' creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
31

---

The answer NA means that the paper does not involve crowdsoureing 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 expect authors adhere 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.
32


Reviews:
[{"review_id": "Y15Nzi7V9a", "reviewer": "Reviewer_969G", "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.", "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? \n\nEvron, 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.\n* 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? \n\nLee, 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.\n\nLee, 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\n\n* 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": "Work is self-contained and aims to explain; the limitations of the theory in explaining certain cases are raised and addressed.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* I really like the style of the paper! Take a theoretical result, extend it to a new setting, understand its implications, and verify them! The implications are not only verified on simplified models but also in more realistic models.\n* I view the audience of your paper as a much broader community than just the theory of deep learning community, as there are concrete and possibly algorithmic implications for your findings (i.e., it may be possible to design algorithms that actively incentivize certain kinds of feature learning in different settings via your results).", "weaknesses": "A good paper overall. See some questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive feedback and your helpful comments. Below we respond to your questions.\n\n> 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?\n\nWe agree that this paper is related. This paper (which focuses on classification, though see Evron et al. (2022), which focuses on regression) also analyzes sequential learning, focusing on how learning of subsequent tasks affects forgetting of previous tasks (i.e. a continual learning perspective). In contrast, we focus (in our PT+FT analysis) on how learning of a previous task impacts generalization on the subsequent task. Further, the authors analyze a shallow network using only a linear readout, whereas we analyze networks with a hidden layer, enabling us to characterize the impact of feature learning (e.g. giving rise to feature sparsity and nested sparsity biases). Thus, these papers use complementary methods to study a related, but different problem. Future work could both analyze generalization on sequential tasks using the methods established in Evron et al., or use our analysis to analyze the continual learning setup. We find both of these directions quite promising and will add the paper to the Related Work section.\n\n> 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?\n\nThere are indeed a number of similarities between these papers and our own investigation. In particular, they also investigate a teacher-student setup (albeit with a different nonlinearity), using two different teachers with various overlaps. Notably, their focus (just like that of Evron et al.) is on forgetting, but they also investigate forward transfer, studying how well the networks can learn novel tasks with various overlaps. In particular, they find that these networks do not always reach optimal training error, which may be due to the online learning setup, the challenging learning landscape caused by the initialization induced by pretraining, and the fact that the student networks have about as many parameters as the teacher. In contrast, we study overparameterized student networks which are able to reach arbitrarily low training error and then investigate their generalization error. Nevertheless, these phenomena may be affected by similar mechanisms: in particular, Lee et al. (2021; 2022) also manipulate the alignment between the two teachers by varying the correlations of their hidden features and find maximal forgetting for intermediate correlations. While we focus on the generalization error on the finetuning task (rather than measuring forgetting on the pretraining task), we expect that in our experiments, the error on the original pretraining dataset after finetuning (i.e. a measure of forgetting) might exhibit a similar non-monotonicity with respect to correlations.\n\nNotably, Lee et al. (2022) consider interleaved replay (similar to MTL in our setup) and find that it performs worse than finetuning (or regularized finetuning) for intermediate task similarity. This is also the case in our studies, where correlated teacher features (corresponding to intermediate task similarity) generally yield worse performance on MTL compared to PT+FT. Indeed, the intuition of Maslow's hammer may transfer to our setting: MTL networks try to re-use the same features on both tasks, which harms their generalization. In contrast, finetuned networks change their pretrained features (which yields forgetting in the setup of Lee et al.), yielding better generalization on the finetuning task.\n\nAll in all, these prior works on continual learning provide useful contextualization of our own work. We will add a paragraph on this to the related work section and will further note the similarity discussed in the previous paragraph in section 4.4. Thank you for bringing these papers to our attention!\n\n> 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.\n\nWe agree with this point and will add a paragraph to the related work section on the contrast between implicit and explicit regularization, explaining why we are using a mixture of both. Roughly, our summary is: Woodworth et al. and Azulay et al. are able to characterize the implicit regularization induced by gradient descent for diagonal linear networks trained from arbitrary initialization. However, it is technically much more challenging to derive a similar result for multi-output diagonal linear networks, or for ReLU networks of any kind. It is generally easier to characterize the impact of explicit weight regularization, which Dai et al. do for multi-output diagonal linear networks. Our own contributions to the theoretical landscape are (1) spelling out the implications of existing results on implicit regularization in diagonal linear networks, when applied to the PT+FT setting, and (2) characterization of the penalty conferred by explicit regularization on finetuning ReLU networks from arbitrary initialization. Again, we take an explicit regularization perspective in the ReLU case (and validate experimentally that our results are a good description of implicit regularization dynamics), because this is technically much more tractable than characterizing the implicit regularization of ReLU networks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P14oZi8fME", "reviewer": "Reviewer_Ry6c", "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.", "questions": "See in Weaknesses.", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The article provides an in-depth characterization of the inductive biases associated with two common training strategies—Multi-Task Learning (MTL) and Pretraining followed by Fine-Tuning (PT+FT)—in diagonal linear and ReLU networks. This detailed analysis is crucial for understanding the impacts of different training strategies.\n\n2. By pushing networks into the nested feature-selection regime, the article proposes simple techniques to improve PT+FT performance, which have shown promising empirical results, adding practical application value.", "weaknesses": "1. Although the findings are promising, the article notes that more work is needed to test these phenomena in more complex tasks and larger models. This suggests that the research's applicability and universality might be limited.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your helpful review. Below we respond in detail to your comments.\n\n> 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.\n\n\nWe are indeed quite excited about investigating these phenomena in such more complex settings — we expect that such an investigation would reveal an even more nuanced picture in terms of when and how an inductive bias towards nested feature sparsity, correlated features, and shared features aids generalization. We believe that such an investigation is best conducted in a new manuscript, whereas the main focus of this paper was to carefully derive and explain how these inductive biases arise in simpler networks. We’d also like to emphasize that we’re already providing an investigation of these empirical phenomena that considers several deep neural network architectures (ResNet, VGG, Vision Transformers) and standard benchmark datasets (Imagenet and CIFAR-10). As such, we believe that we are already providing useful evidence for the applicability of our insights, but we acknowledge that there is much more work to be done on this in future papers.\n\n> 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.\n\nWe agree that these are exciting avenues. Understanding how the penalties arising from explicit regularization differ from and connect to the penalties arising from the implicit regularization of gradient descent is technically challenging but an important direction. (We note that characterizing the impact of explicit regularization is also useful and important in its own right, and is a common strategy for understanding the inductive biases of networks, see e.g. Savarese et al. (2019); Evron et al. (2022); see also l.97-102 in our manuscript.) Similarly, we would be interested in extending our theory to the cross-entropy loss (as well as other losses) and expect the same principles of feature sparsity, sharing, and correlations to be important in that setting as well. Doing so would require introducing new teacher-student setups and introduce some additional practical considerations (e.g. weights of cross-entropy networks grow without bound rather than converge) which we believe are best left for another paper. We felt it was better to focus this paper on analyzing the regression setting in depth, to provide a solid foundation for future extensions to our work.\n\n> 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.\n\nThank you for this suggestion. We agree that describing the theoretical setup in a separate section could be a helpful change. As a result, we will make the following change to the camera-ready version if accepted: we will add a new section 3.1 (“Theoretical setup”), which defines diagonal linear networks and ReLU networks and defines multi-task learning and pretraining+finetuning. We will then introduce our theoretical results in dedicated sections.  Currently, our sections 3.1 and 3.2 both introduce the network setups and describe theoretical results – these changes will have the effect of refactoring these sections to separate our description of the setup from our theoretical results. Could you clarify if this is the kind of change you had in mind? We are open to other ways of organizing the content.\n\n**References**\n\nSavarese, Pedro, et al. \"How do infinite width bounded norm networks look in function space?.\" Conference on Learning Theory. PMLR, 2019.\n\nEvron, Itay, et al. \"How catastrophic can catastrophic forgetting be in linear regression?.\" Conference on Learning Theory. PMLR, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pfhPIFVLDA", "reviewer": "Reviewer_FiYt", "summary": "Abstract\n\nThe goals of the paper are very clear from the abstract, as are the results.\n\nIntroduction\n\nLines 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\n\nRelated Work\nLines 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.\n\nLines 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.\n\nImplicit and explicit regularization penalties for MTL and PT+FT\n\nLines 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.\n\nLines 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. \n\nLine 112: Typo: “A ReLU networks”\n\nLines 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”. \n\nLines 117-124: Why are the DLN network weights initialized with a constant magnitude?\n\nLines 128-129: We are not told how we get to (6) and not appendix section is referenced\n\nLine 131: Typo\n\nLine 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?\n\nLine 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)?\n\nLines 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.\n\nFigure 2: I would define STL as single task learning for the less familiar reader\n\nLines 144-151: The experiment is explained well and easy to follow. \n\nLines 152-155: More explanation for the ReLU would help. By “sparse” number of (hidden?) units we mean a low dimensional hidden layer?\n\nLines 158-165: Everything is clear save how we arrived at c. Appendix reference?\n\nLines 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. \n\nLines 183-194: Again, elegance in defining the feature sharing to explicitly study what was of interest (simultaneous sparsity and feature sharing).\n\nLines 206-215: The definitions are clear and the intermediate regime is set up well by the authors.\n\nLines 223-250: Again, the elegance of the experiment set-ups to reinforce the nested feature selection regime.\n\nLines 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. \n\nLines 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.\n\nLines 312-315: Very interesting finding on a large dataset and more complex network. \n\nSummary: 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.", "questions": "Why are the DLN network weights initialized with a constant magnitude?\n\nThe 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.\n\nBeyond 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": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The theoretical justifications for the set-up of the experiments are excellent. The experiments that follow are also well-designed to allow the authors to probe their implicit regularization penalties. The authors kept probing until arriving on their nested feature regime. I liked the trajectory from theory to synthetic data/systems experiments to findings to larger/classic datasets and networks.", "weaknesses": "I think some of the notation needs to be cleaned-up to have this important work reach as broad an audience as possible. Much is left to the appendix and when it is the references are usually there, but I would make sure they are always there for the reader to reference from the main body of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive feedback and your helpful comments. We will make sure to make the notation more accessible and really appreciate your detailed notes on this. Below we respond to the questions and comments you raised in your review. (Due to the character limit, we focus on the most important responses in the rebuttal and have relegated additional responses to a comment.)\n\n> Why are the DLN network weights initialized with a constant magnitude?\n\nWe wanted to remove the random variation induced by having randomly sampled readout weights, and  to be consistent with Woodworth et al. who also considered constant-magnitude initialization. However, we would expect qualitatively similar results for randomly sampled weights from a Gaussian. Unlike for ReLU networks, random weights are not really necessary for symmetry breaking in diagonal linear networks, as each unit already receives input from a different dimension.\n\n> 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.\n\nWe agree that the task dependence of where the nested feature selection regime falls is an interesting phenomenon.  As the reviewer points out, we show a theoretical basis for such sensitivity in Fig. 3d.  However, some of our empirical results are not completely described by the theory – empirically, even when auxiliary and main task features are identical, ReLU networks exhibit behavior that our theory would expect to arise in the case where they are highly correlated but not identical. We conjecture that this discrepancy arises due to noise in stochastic gradient descent dynamics (i.e. the inferred features on the two tasks may be highly correlated but not identical).  Furthermore, we observe differences between ResNets and Vision Transformers in where the nested feature selection regime occurs.  Our shallow network theory is not able to speak to these architectural differences. We think that clarifying this picture is an important direction for future work, and will add a sentence on this to the conclusion.\n\n> Beyond me being pedantic, is this entire paper working on multi-output learning not multi-task learning? (...) 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.\n\nWe agree that in practice multi-task learning is much more versatile than the setup we consider here. We opted for such similar main and auxiliary tasks to remove additional factors of variation and to keep the teacher-student setup parsimonious, but will update the manuscript to acknowledge the much broader range of multi-task training used in practice. We expect that our findings are relevant to other kinds of multi-task learning, but have not yet established this, and think it is an important direction for future work.\n\n> 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.\n\nWe agree that this is a really fascinating phenomenon. So far, we have found a (not fully rigorous) theoretical basis for differences in the rescaling magnitude required to enter the nested feature selection regime between diagonal linear networks and ReLU networks (see lines 245-250 for our explanation). We have also empirically found that rescaling by a lower value is required for ResNets to enter the nested feature selection regime, whereas this is not required for Vision Transformers (and indeed such rescaling appears to cause them to leave this regime). We don’t yet understand the reason for this difference, and we agree that a more comprehensive investigation of this phenomenon (across networks and tasks) would be very interesting future work. We will add a sentence on this topic to the conclusion.\n\n> 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.\n\nTo clarify, while MTL can certainly learn correlated features using different hidden units, it does not benefit from such correlated units in terms of its overall norm (at least in a neural network with one hidden layer). Note that in the low-sample regime, MTL simply uses the same unit for both tasks, which is not entirely accurate but works reasonably well (see Fig. 8e, which shows that the student units have a correlation of 0.9 with the main task teacher, suggesting that they are identical with the units the student uses for the auxiliary task). We'll change this sentence to make that point clearer.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "75yjfXLniF", "reviewer": "Reviewer_4YZg", "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:\n1. 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).\n2. The authors prove that during finetuning networks interpolate between a \"lazy\" and \"rich\" regime. (Proposition 4).\n3. For PT+FT it is observed empirically that the choosen features are a sparse subset of the features learned during pretraining. \n4. The paper provides empirical evidence that for ReLU networks only PT+FT benefits from features that are correlated between the auxilary and primary task. \n5. Finally, they present a practical technique to improve PT+FT by weight rescaling after pretraining.", "questions": "- Where is the proof for Corollary 2?\n- When finetuning are the weights of both layers trained or only the second layer weights?\n- Where are the results from the experiments discussed in Section 3.3?\n- 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?\n- Line 176 what do you mean by \"overlapping dimensions\"? You mean overlapping features?\n- On line 178 what is LP in PT+FT(LP)?", "limitations": "Limitations have been mentioned in the weaknesses section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "I think the setting is interesting and very relevant to modern deep learning. This is also one of the few papers studying the inductive bias of multi-task networks. In particular:\n- The paper offers a number of results with many experiments on toy models.\n- The finding that MTL and PT+FT induces a regularizer that interpolates between $\\ell_1$ and $\\ell_2$ regularization is interesting and novel.\n- The authors go beyond theory and show how their finding can be exploited in practical architectures through a simple weight rescaling technique.", "weaknesses": "- Corollary 1 and its extension for ReLU neural networks in Appendix A.1 is already well known in the literature [2] (see discuss after Remark 2) and [3] (Equation 4). These and the references therein should at least be mentioned.\n- 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.\n- 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.\n- 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. \n- The notation is very loaded and makes the paper difficult to parse. For example $\\vec{w}^{(1),aux}_{h}$. \n- 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.\n\n\n**Minor things**\n- Lines 51 and 54 reuse the word \"Finally\".\n- Line 82: Should the inputs to the model $x \\in \\mathbb{R}^{D}$?\n- Line 113 paraeter --> parameter\n- Line 135: Typo $\\mathbb{R}^{d} \\in \\mathbb{R}$\n- Line 251: yon --> on\n- In Section 3 $D$ and $d$ are both used interchangeably for the input dimension of the networks.\n\n[1] Collins, Liam, et al. \"Provable multi-task representation learning by two-layer relu neural networks.\" arXiv preprint arXiv:2307.06887 (2023).\n\n[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).\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your helpful review. Below we respond to your comments and questions. (Due to the character limit on the rebuttal, we focus here on the most important responses and have relegated additional responses to a comment.)\n\nFirst, in response to the reviewer’s summary of the paper, we wanted to clarify a few of our key contributions. If you think we could highlight these contributions more clearly in the revised version of the paper, we would appreciate any suggestions.\n\n> The authors prove that during finetuning networks interpolate between a \"lazy\" and \"rich\" regime. (Proposition 4).\n\nSpecifically, we show that this interpolation can be described in terms of a conserved quantity expressing a tradeoff between initialization dependence and sparsity, and that a network’s position on this tradeoff is tied to the scale of the weights following pretraining.\n\n> For PT+FT it is observed empirically that the choosen features are a sparse subset of the features learned during pretraining.\n\nThe insight above allowed us to predict the existence of a “nested feature selection” regime, in which finetuning extracts a sparse subset of the features learned during pretraining (due to intermediate levels of both sparsity and initialization dependence). This is a qualitatively different regime from the lazy regime (which is initialization-dependent, but not sparse) and rich regime (which is sparse, but not initialization-dependent) that were known previously. Networks do not always exhibit this behavior, but do in some cases, and can be induced to exhibit it via the weight rescaling technique we introduce.\n\n**Responses to comments**\n\n> - 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.\n\n> - 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.\n\nThank you for pointing us to these highly relevant papers; they were interesting to read and indeed highly relevant. We’ll make sure to cite these papers when describing the MTL penalty and when describing the feature reuse bias (in particular, Theorem 9 in [2]). We will also discuss this line of work in the Related Work section. We note that in our original draft we did not claim that Corollary 1 was an original result (we cited reference [16], Dai et al.) but we appreciate the relevant references for the ReLU network case and agree they are important to cite.\n\n> 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.\n\nWe see our results on diagonal linear networks as complementary to our results on ReLU networks, as we are able to prove our theoretical claims with fewer assumptions in the diagonal linear case.  In particular, the diagonal linear setting allows us to analytically derive the effects of implicit regularization due to gradient descent for PT+FT (in the ReLU case we had to assume an explicit regularizer, due to the difficulty of analyzing gradient descent dynamics in nonlinear networks). This helps clarify the mechanism underlying our key theoretical results -- the nested feature-selection regime and our “conservation law” -- and motivates our subsequent ReLU network analysis.\n\n> The regularizer used during finetuning in Proposition 3 isn't well motivated. (...) It would be great to provide some discussion on why you choose to analyze this regularizer or whether this is done in practice. \n\nWe agree that this deserves discussion separately from the explicit regularization considered in the multi-task learning setup. We chose to consider this regularization penalty for two reasons. First, infinitesimal explicit regularization from initialization is equivalent to the implicit regularization induced by gradient descent in the case of shallow linear models [4 below]. While this is not true more generally, it is a useful heuristic to motivate theoretical analysis (which of course must then be checked against experiments). Assuming the explicit regularization heuristic allowed us to derive a number of nontrivial predictions that we would not otherwise have come up with, and which we confirmed hold true in empirical simulations of unregularized gradient descent.  We also note that this explicit regularization penalty is sometimes studied in the context of continual learning [e.g. 5,6 below].  We will make sure to clarify the motivation for basing our ReLU theory on this regularization penalty in the revised manuscript. \n\n> Where is the proof for Corollary 2?\n\nThis 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.\n\n> Where are the results from the experiments discussed in Section 3.3?\n\nSection 3.3 describes the setup for all the teacher-student experiments presented in section 4. We will clarify this in the revised manuscript.\n\n4. Gunasekar et al. \"Characterizing implicit bias in terms of optimization geometry.\" ICML (2018).\n5. Lubana et al. \"How do quadratic regularizers prevent catastrophic forgetting: The role of interpolation.\" CoLLAs (2022).\n6. Evron et al. \"Continual learning in linear classification on separable data.\" ICML (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their helpful feedback and comments. In particular, the reviewers have suggested a number of relevant papers that we will add to the related work section. In addition, they have provided valuable feedback on clarity, which we will take into account in revising the manuscript (see specific responses to reviewers). We respond to the individual reviews below.

Meta-review: 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.
Decision: Accept (poster)