Revisiting Scalarization in Multi-Task Learning
Han ZhaoInternational Chinese Statistical Association 2024 China Conference, Wuhan
06/28/2024(Joint work with Yuzheng Hu, Ruicheng Xian, Qilong Wu, Qiuling Fan and Lang Yin)hanzhao@illinois.edu
Amazon Visiting Academic
Assistant Professor
Department of Computer Science
University of Illinois Urbana-Champaign

What is Multitask Learning?
Multitask learning:

Suppose we need to model k di erent tasks

2ff-𝒟

Multitask Representation Learning
Multitask learning:(i) (i) nitr={(x,y)}ijjj=1-A set of k tasks-Each task can have a separate loss function ℓi( ⋅ , ⋅ )-Goal: better generalization by simultaneous training all tasks than single
task learning (train a separate model for each task)

Conceptual Framework

Model Architecture

3

Multitask Representation Learning
Multitask learning:-

Train a joint model with multiple heads, where one head ~ one task
• Pros:➡ Easy to scale to many tasks: adding one more task-speci chead➡ Can exploit potential synergies among tasks• Cons:➡ Hard to design tailored structures➡ “Negative transfer” could happen due to con icting tasks

4fiflFigure credit: “Multitask Learning”, PhD Thesis, Rich Caruana, 1997

Multitask Representation Learning
Multitask learning:
MTL Loss: min<latexit sha1_base64="9JE8x/wlnjKmBG3W567LDxGJPYo=">AAAEwHicfZNta9swEMeVNmu77KndXu6NWCkkUEpcyro3hdKtkL3ruj5BnRpZlhM1suRIcvMg9DnHPsxgcpIticJ2YDjf7393tk4X54wq3Wz+rKytV59tbG49r714+er1m+2dtzdKFBKTayyYkHcxUoRRTq411Yzc5ZKgLGbkNu59LvntE5GKCn6lRzlpZ6jDaUox0i4UbffDjPLIdPZhN6I27BcoCVOJsAms6dlQFVlk6ElgH3rwb5yXygl5LMn0nTAW0XrdVQkxlRh2GvXhg6nTho0eG/twFD1O3xrR9m7zoDkxuOoEM2cXzOwi2lk/DhOBi4xwjRlS6j5o5rptkNQUM2JrYaFIjnAPdci9cznKiGqbydlYuOciCUyFdA/XcBJdzDAoU2qUxU6ZId1VPiuD/2K6m1k/VnZRFk5siTmxFoKp5Yyy7mr4vtDpp7ahPC804Xj6F2nBoBawnCJMqCRYs5FzEJbUHQTEXeTmo92sa3vQb5GjsbC1pRYZ6hEpROZOCnFMmPsATgZYZBniiQkpT6wJy9w4Nl+tRweEdrr6jyA1A18wHM7h0Iej0RyOfDgez+F4pe/gPz1VzmcUI2a+r2ItCVlQXPmKVmuBtnxKpHQXyoRPSJJcUSa4J0jcui4U+OIXSEhK+mUBXN5kdxH5id+i7zTWnCxIPMX5MJ+P5dzvcHk5h5fWulUL/MVadW4OD4KPB0ffjnZPz2ZLtwXegw+gDgJwDE5BC1yAa4DBD/CrslHZrJ5Vu1VR7U+la5VZzjuwZNXxb8V+rJI=</latexit>

g,hinikX
X
11k i=1 ni j=1“task-speci c header”`i ((hi(i)(i)g)(xj ), yj )“shared multitask feature learning”

Multi-objective vector loss: ℓ({hi}ki=1, g) :=tr
1[ℓ1((h1 ∘ g)(X), Y )]⋮∈ ℝk
tr[ℓk((hk ∘ g)(X), Y )]k𝒟Note “≤” is not a total ordering in ℝk for k > 1, so in general wepick a preference vectorfi𝔼

𝒟𝔼w ∈ Δk−1 :=kv∈ℝ :vi = 1,vi ≥ 0
∑
{}
i=1k5

Multitask Representation Learning
Multitask learning:
MTL Loss: min<latexit sha1_base64="9JE8x/wlnjKmBG3W567LDxGJPYo=">AAAEwHicfZNta9swEMeVNmu77KndXu6NWCkkUEpcyro3hdKtkL3ruj5BnRpZlhM1suRIcvMg9DnHPsxgcpIticJ2YDjf7393tk4X54wq3Wz+rKytV59tbG49r714+er1m+2dtzdKFBKTayyYkHcxUoRRTq411Yzc5ZKgLGbkNu59LvntE5GKCn6lRzlpZ6jDaUox0i4UbffDjPLIdPZhN6I27BcoCVOJsAms6dlQFVlk6ElgH3rwb5yXygl5LMn0nTAW0XrdVQkxlRh2GvXhg6nTho0eG/twFD1O3xrR9m7zoDkxuOoEM2cXzOwi2lk/DhOBi4xwjRlS6j5o5rptkNQUM2JrYaFIjnAPdci9cznKiGqbydlYuOciCUyFdA/XcBJdzDAoU2qUxU6ZId1VPiuD/2K6m1k/VnZRFk5siTmxFoKp5Yyy7mr4vtDpp7ahPC804Xj6F2nBoBawnCJMqCRYs5FzEJbUHQTEXeTmo92sa3vQb5GjsbC1pRYZ6hEpROZOCnFMmPsATgZYZBniiQkpT6wJy9w4Nl+tRweEdrr6jyA1A18wHM7h0Iej0RyOfDgez+F4pe/gPz1VzmcUI2a+r2ItCVlQXPmKVmuBtnxKpHQXyoRPSJJcUSa4J0jcui4U+OIXSEhK+mUBXN5kdxH5id+i7zTWnCxIPMX5MJ+P5dzvcHk5h5fWulUL/MVadW4OD4KPB0ffjnZPz2ZLtwXegw+gDgJwDE5BC1yAa4DBD/CrslHZrJ5Vu1VR7U+la5VZzjuwZNXxb8V+rJI=</latexit>

g,hinikXX
11k i=1 ni j=1`i ((hitr1

Multi-objective vector loss: ℓ({hi}ki=1, g) :=(i)(i)g)(xj ), yj )[ℓ1((h1 ∘ g)(X), Y )]⋮∈ ℝk
tr[ℓk((hk ∘ g)(X), Y )]k

Note “≤” is not a total ordering in ℝk for k > 1, so in general wepick a preference vectorw ∈ Δk−1 :=𝔼𝒟Linear scalarizationkv ∈ ℝk :vi = 1,vi ≥ 0
∑
{}
i=1
min

𝔼𝒟g,hikni

wiℓi((hi ∘ g)(xj(i)), yj(i))∑ ni ∑
i=1
j=1
6

Multitask Representation Learning
Linear scalarization:
MTL Loss: ming,hiGeometrically:Task 2nikwiℓi((hi ∘ g)(xj(i)), yj(i))∑ ni ∑
i=1
j=1

4

3

2

1
0.5

1.0

1.5

w = (w1, …, wj) ∈ Δk−1

2.0

2.5

3.0

Task 17

Linear Scalarization for MTL
Linear scalarization:
MTL Loss: ming,hikni

wiℓi((hi ∘ g)(xj(i)), yj(i))∑ ni ∑
i=1
j=1

What’s the potential problem of simple scalarization?
Not necessarily possible to strike the right balance among multiple tasks

8Linear Scalarization for MTL
Linear scalarization:
Not necessarily possible to strike the right balance among multiple tasks

Empirically, this means:-For Pareto optimal solutions, the losses for di erent tasks cannot changesmoothly-

While it’s possible to optimize for one of the tasks by assigning a larger weight,
9
this is at the cost of signi cant performance drop for the otherffCannot explore all the Pareto optimal solutionsfi-Not necessarily possible to strike the right balance among multiple tasks(ℱ) =(conv(ℱ))(ℱ) ≺Feasible region of di erent models:ℱ := {−ℓ(θ) ∈ ℝk : θ = ({hi}ki=1, g), ℓi(θ) =Pareto front of ℱ:(conv(ℱ))[ℓi((hi ∘ g)(X), Y )]}

tri

ff￼

𝔼

￼

￼

𝒫

𝒟

(ℱ) := {v ∈ ℱ : ∃v′/ ∈ ℱ, ( ∀i ∈ [k], v′i ≤ vi) ∧ ( ∃i ∈ [k], v′i < vi}
𝒫𝒫𝒫𝒫

Linear Scalarization for MTL

10Multi-Objective Optimization
Methods that can potentially overcome the limitation:

Key idea: using adaptive and dynamic weight preference vector at each iteration ofthe optimization algorithms!

fl“Multiple-Gradient Descent Algorithm (MGDA) for Multi-Objective Optimization”, Comptes Rendus Mathématique,
Désidéri, 2012.
“Multi-Task Learning as Multi-Objective Optimization”, NeurIPS’18, Sener and Kolton.
“Gradient Surgery for Multi-Task Learning”, NeurIPS’20, Yu et al.
11“Con ict-Averse Gradient Descent for Multi-task Learning”, NeurIPS’21, Liu et al.

Multi-Objective Optimization
Multiple-Gradient Descent Algorithm
Let gi, i ∈ [k] be the gradient for the k tasks at a certain iteration

First order improvement along direction d: gi⊤dPrimal problem:⊤max min gi d∥d∥2≤1 i∈[k]Interpretation: nding a common direction d that maximizes theworst-task improvement

Dual problem:min ∥α∈Δk−1k∑i=1αigi∥22

Interpretation: nding a convex combination α of themultiple gradients with minimum ℓ2 norm

fifi“Multiple-Gradient Descent Algorithm (MGDA) for Multi-Objective Optimization”, Comptes Rendus Mathématique,
12Désidéri, 2012.

Linear Scalarization vs Multi-Objective Optimization
Empirical comparisons between linear scalarization vs multiobjective optimization in NNs:

13

Linear Scalarization vs Multi-Objective Optimization
Empirical comparisons between linear scalarization vs multiobjective optimization in NNs:

Observations:-For large-scale NNs, the closure of the Pareto front seems to be convex-Deeper and larger models help to enlarge the feasible region (henceachieving Pareto-dominating solutions)“Do Current Multi-Task Optimization Methods in Deep Learning Even Help?”, NeurIPS’ 22, Xin et al.

14Research Question:“For NNs, for every Pareto optimal v ∈ (ℱ), does
there exists a w ∈ Δk−1 such that the optimal solution ofthe linear scalarization problem corresponds to v?”(ℱ) =(conv(ℱ))-What is the impact of model size?
𝒫

When MOO may help?
𝒫𝒫

𝒫𝒫

Linear Scalarization vs Multi-Objective Optimization(ℱ) ≺(conv(ℱ))
15Linear Scalarization vs Multi-Objective Optimization
When all the loss functions are convex:
Theorem (convex, Boyd and Vandenberghe’ 04): If all the loss functionsℓi(θ), i ∈ [k] is convex with respect to θ, then linear scalarization withproper choice of convex combination weights w ∈ Δk−1 can explore every
Pareto optimal solution.

Does not apply to NNs due to non-convexity16Linear Scalarization vs Multi-Objective Optimization
Multi-task linear network for regression:-For each task i ∈ [k], the prediction is given by⊤fi(x, W, ai) = x Wai-

Shared input X ∈ ℝn×p, target vector yi ∈ ℝn, ∀i ∈ [k], the training
loss for task i:ℓi(W, ai) = ∥XWai − yi∥22-Parameter W ∈ ℝp×q, i.e., network width = q, ai ∈ ℝq are task-speci cparameters

Note:

The overall model parameters θ = ({ai}ki=1, W ). The optimization problem is nonconvex for each task i

17fi-For linear networks without loss of generality it su ces to consider a two layer
networkffi-

Phase-transition between over-parametrized and under-parametrizednetworks:

Over-parametrized regime (q ≥ k):

Theorem (informal): The network has su cient capacity to t all thetasks optimally, and the Pareto front reduces to a singleton(ℱ) = {c}, c ∈ ℝk and hence can be attained via an arbitrary choiceof convex coe cient w ∈ Δk−1.

fiffi“Understanding and Improving Information Transfer in Multi-Task Learning”, Zhang et al., ICLR’ 20“Revisiting Scalarization in Multi-Task Learning: A Theoretical Perspective”, Hu et al., NeurIPS’ 23
ffi

𝒫Linear Scalarization vs Multi-Objective Optimization18

Phase-transition between over-parametrized and under-parametrizednetworks:

Over-parametrized regime (q ≥ k):

Theorem (informal): The network has su cient capacity to t all thetasks optimally, and the Pareto front reduces to a singleton(ℱ) = {c}, c ∈ ℝk and hence can be attained via an arbitrary choiceof convex coe cient w ∈ Δk−1.
Intuition: let yî := X(X ⊤X)†X ⊤yi be the optimal linear predictor for task iand let Y ̂ = [y1̂ , …, yk̂ ] ∈ ℝn×k be a stack of column vectors. Then, due
to the ℓ2 loss, each task-speci c loss can be decomposed asℓi(W, ai) = ∥XWai − y∥22 = ∥XWai − yî ∥22 + ∥yî − yi∥22
Fitting loss: can choose aiapproximation error: projection lossfi19

ffi

fi

separately for a given shared W
ffi

𝒫

Linear Scalarization vs Multi-Objective Optimization

Linear Scalarization vs Multi-Objective Optimization
Over-parametrized regime (q ≥ k):

Intuition: let yî := X(X ⊤X)†X ⊤yi be the optimal linear predictor for task iand let Y ̂ = [y1̂ , …, yk̂ ] ∈ ℝn×k be a stack of column vectors. Then, due
to the ℓ2 loss, each task-speci c loss can be decomposed asℓi(W, ai) = ∥XWai − y∥22 = ∥XWai − yî ∥22 + ∥yî − yi∥22Note that yî ∈ Col(X) is the projection of yi into the column spacespanned by X ∈ ℝn×p, i.e., the optimal linear prediction. Then if thenetwork is wide enough, i.e., q ≥ k, we can:-Optimize the network parameter W ∈ ℝp×q and {ai}ki=1 ⊆ ℝq by allocating
one neuron for each task tting loss yî

-Pick W ∈ ℝp×q such that W is full-rank

For every i ∈ [k], XWai = yî has a solution in terms of ai ∈ ℝq (becauseq ≥ k)

Putting all together, we have ∥XWai − yî ∥22 = 0 and ci = ∥yî − yi∥22,∀i ∈ [k]fi20
fi-Linear Scalarization vs Multi-Objective Optimization
Under-parametrized regime (q < k):Theorem (informal): We focus on two extremal cases:-Extremely under-parametrized (q = 1): Linear scalarization su ces, i.e.,̂ Y ̂ is doubly
full-exploration of the Pareto front, if and only if G := Y⊤
non-negative, i.e., the inner products for all task pairs yî and yĵ are nonnegative, up to negating the directions of some yî , i ∈ [k]-Mildly under-parametrized (q = k − 1): Linear scalarization su ces if
and only if Q = G −1 is doubly non-negative, up to negating thedirections of some yî , i ∈ [k]Remark:Su cient and necessary conditions for the general case of 1 < q < k − 1 still open
ffi

ffi

-

G and Q could be understood as a notion of “task-similarity” — task similarity ismodel-dependent!, i.e., Gij = ⟨yî , yĵ ⟩ffi-This means that in general under the under-parametrized regime, linear scalarization
is not su cient of full exploration

ffi-21

Linear Scalarization vs Multi-Objective Optimization
Geometric intuition of the under-parametrized regime:⊤ † ⊤̂
Notation: let yi := X(X X) X yi be the optimal linear predictor for task iand let Y ̂ = [y1̂ , …, yk̂ ] ∈ ℝn×k be a stack of column vectors.

For every xed W ∈ ℝp×q, Z = XW ∈ ℝn×q are the linear representations learned⊤ † ⊤=(ZZ) Z yî .
by NNs. Each task-speci c head admits an optimal solution a*iHence, each task-speci c loss ℓi could be simpli ed tomin ∥Z(Z ⊤Z)†Z ⊤yî − yî ∥22

Z=XWLet PZ := Z(Z ⊤Z)†Z ⊤ be the projection matrix under a xed linear
representation Z = XW, then the MOO optimization problem becomesmax

Observations:

fi

fi

22
fi

fi

fi

The feasible region ℱ will be the union of 2k ellipsoidff-

Under negation, there are 2k di erent con gurations, each con guration corresponds to an(potentially degenerated) ellipsoidfi-

The objective is invariant under negation of yî , i ∈ [k]fi-PZ(y⊤1̂ PZ y1̂ , …, y⊤k̂ PZ yk̂ )Linear Scalarization vs Multi-Objective Optimization
Geometric intuition of the under-parametrized regime:⊤† ⊤Let PZ := Z(Z Z) Z be the projection matrix under a xed linear
representation Z = XW, then the MOO optimization problem becomes
PZ

Observations:

The feasible region ℱ will be the union of 2k ellipsoidfi

23fi-Under negation, there are 2k di erent con gurations, each con guration corresponds to an(potentially degenerated) ellipsoidfi-The objective is invariant under negation of yî , i ∈ [k]ff-(y⊤1̂ PZ y1̂ , …, y⊤k̂ PZ yk̂ )maxLinear Scalarization vs Multi-Objective Optimization
Under-parametrized regime (q < k): how to rescue?
Randomization!

Given two under-parametrized networks f0 and f1, we can construct a
randomized network as follows:

f0(x) if S ≤ tf(x) ={f1(x) o . w .
where t ∈ [0,1] and S ∼ U(0,1) is a uniform RV over (0,1). Then
S,X,Y[ℓ( f(X), Y )] = t X,Y[ℓ( f0(X), Y )] + (1 − t) X,Y[ℓ( f1(X), Y )]By choosing di erent t ∈ (0,1) we can interpolate and hence convexity
any given feasible region ℱ.

(ℱS) =(conv(ℱ))
𝔼fiff𝒫𝔼ℱS = conv(ℱ)
𝔼
𝒫Randomization ≈ Convexi cation24

Linear Scalarization vs Multi-Objective Optimization
Under-parametrized regime (q < k): how to rescue?
Chebyshev scalarization: w ∈ Δ :=k−1

Chebyshev scalarization{v ∈ ℝk :nikv = 1,vi ≥ 0∑ i}
i=1wimin maxℓi((hi ∘ g)(xj(i)), yj(i))
g,hi i∈[k] ni ∑j=1

Theorem (Choo & Atkins, 1983): Any feasible solution that is weakly
Pareto optimal if and only if it is a solution for a weighted Chebyshev
problem under some preference vector w ∈ Δk−1.

ff

ffi

“Proper E ciency in Nonconvex Multicriterion Programming”, Choo and Atkins,
Mathematics of Operation Research, 1983
“Robust Multi-Task Learning with Excess Risks”, He et al., ICML’ 24“Smooth Tchebyche Scalarization for Multi-Objective Optimization”, Lin et al., ICML’2425

Linear Scalarization vs Multi-Objective Optimization
Multi-task non-linear network for regression:
Theorem (informal): If q ≥ nk, then there exists a network that hassu cient capacity to t all the tasks optimally, and the Pareto front
reduces to a singleton (ℱ) = {c}, c ∈ ℝk and hence can be attained
via an arbitrary choice of convex coe cient w ∈ Δk−1.
Remark:-This upper bound is potentially very loose; ideally we would like anupper bound on the width q that only depends on the number of tasks
k not the number of data points n-No lower bound known, i.e., is there a trade-o problem in general forunder-parametrized nonlinear NNs?
NN-WidthLinear NNs

Nonlinear NNs

Upper bound (su cient)

knk

Lower bound (necessary)

k?
ffffi

fi

ffi

ffi

𝒫26

Linear Scalarization vs Multi-Objective Optimization
Multi-task non-linear network for regression:
Theorem (informal): If q ≥ nk, then there exists a network that hassu cient capacity to t all the tasks optimally, and the Pareto front
reduces to a singleton (ℱ) = {c}, c ∈ ℝk and hence can be attained
via an arbitrary choice of convex coe cient w ∈ Δk−1.
Empirical evidence:

ffi

fi

ffi

𝒫27

SummaryInsights and Implications for us:-The problem of linear scalarization vs MOO for multitask learning
is model-dependent-

Good news: with su cient capacity of the networks, linear
scalarization is su cient-

Model-dependent notion of task similarity-

For linear MTL models under regression tasks we identify aprecise phase-transition q = k-

For nonlinear MTL models we have a (loose) upper bound q = nk

Open questions:
Is there a similar phase transition phenomenon for nonlinear MTLmodels?

-Can we identify a notion of task similarity for nonlinear MTL models?
fi-ffiHow about linear MTL for classi cation?

ffi

-28

Thanks!

Q&A(Joint work with Yuzheng Hu, Ruicheng Xian, Qilong Wu, Qiuling Fan and Lang Yin)
Paper: https://arxiv.org/pdf/2308.13985.pdf (NeurIPS’ 23)29