Organizations: Nanyang Technological University · Hong Kong University of Science and Technology (Guangzhou) · School of Artificial Intelligence, Shanghai Jiao Tong University
Multi-Task Learning (MTL) is a foundational machine learning problem that has seen extensive development over the past decade. Recently, various optimization-based MTL approaches have been proposed to learn multiple tasks simultaneously by altering the optimization trajectory. Although these methods strive to de-conflict and re-balance tasks, we empirically identify that their effectiveness is often undermined by an overlooked factor when employing advanced optimizers: the instant-derived gradients play only a marginal role in the actual parameter updates. This discrepancy prevents MTL frameworks from fully releasing its power on learning dynamics. Furthermore, we observe that Muon-a recently emerged advanced optimizer-inherently functions as a multi-task learner, which underscores the critical importance of the gradients used for its orthogonalization. To address these issues, we propose APT (Applicability of advanced oPTimizers), a framework featuring a simple adaptive momentum mechanism designed to balance the strengths between advanced optimizers and MTL. Additionally, we introduce a light direction preservation method to facilitate Muon's orthogonalization. Extensive experiments across four mainstream MTL datasets demonstrate that APT consistently augments existing MTL approaches, yielding substantial performance improvements.
Figures & tables
Figure 1: Illustration of the impact of EMA in advanced optimizer for MTL. Here we adopt gupdate=0.9∗gmom+0.1∗gmtl as the example.
β1
MGDA
CAGrad
Adam
AdamW
Muon
Adam
AdamW
Muon
default
7.29 ± 4.61
6.41 ± 5.71
3.37 ± 4.02
12.38 ± 1.49
10.57 ± 1.80
-2.37 ± 0.54
0.6
9.21 ± 4.12
5.82 ± 3.37
4.14 ± 1.10
9.78 ± 1.14
10.88 ± 1.54
1.95 ± 0.90
0.2
9.25 ± 3.75
5.17 ± 3.24
19.39 ± 6.64
10.78 ± 3.03
10.07 ± 0.96
81.12 ± 9.13
0.0
3.52 ± 1.50
10.23 ± 4.68
33.35 ± 6.61
9.11 ± 0.97
10.08 ± 2.58
120.78 ± 36.70
β1
Nash-MTL
FairGrad
Table 1: Analysis on β1 across various MTL approaches and advanced optimizers on CityScapes . AdamW is set with weight decay as 1.e-4 to distinguish with Adam. All results are average on 3 random seeds. The best scores are provided in gray .
Figure 2: Similarity and project norm comparison.
Figure 3: Effective rank comparison.
Figure 4: Effect of momentum on toy example.
Method
Segmentation ↑
Depth ↓
Surface Normal
MR ↓
Δ m% ↓
mIoU
Pix. Acc.
Abs. Err.
Rel. Err.
Angle Distance ↓
Within t∘↑
Mean
Median
11.25
22.5
30
Independent
38.30
63.76
0.68
0.28
25.01
19.21
30.14
57.20
69.15
-
-
LS
39.29
65.33
0.55
0.23
28.15
23.96
22.09
47.50
61.08
12.56
5.46
RLW
37.17
63.77
0.58
0.24
28.27
24.18
22.26
47.05
60.62
15.78
7.67
DWA
39.11
65.31
0.55
0.23
27.61
23.18
24.17
50.18
62.39
11.78
3.49
Table 2: Scene understanding ( NYUv2 , 3 tasks). We report MTAN model performance averaged over 3 random seeds.
Method
Segmentation ↑
Depth ↓
MR ↓
Δ m% ↓
mIoU
Pix.
Abs.
Rel.
Independent
74.01
93.16
0.0125
27.77
-
-
LS
75.18
93.49
0.0155
46.77
11.25
22.60
RLW
74.57
93.41
0.0158
47.79
14.50
24.37
DWA
75.24
93.52
0.0160
44.37
11.25
21.43
Uncertainty
72.02
92.85
0.0140
30.13
10.75
5.88
Table 3: Scene understanding ( CityScapes , 2 tasks). We report MTAN model performance averaged over 3 random seeds. The best scores are provided in gray , and the second scores are underlined .
Figure 5: Toy example comparison.
Method
CelebA
QM9
MR ↓
Δ m% ↓
MR ↓
Δ m% ↓
LS
9.35
4.15
12.45
177.6
SI
11.63
7.20
7.36
77.8
RLW
8.13
1.46
14.27
203.8
DWA
10.40
2.40
12.27
175.3
UW
8.70
3.23
9.09
108.0
Table 4: Results on CelebA and QM9 datasets with MR and Δ m%.
Institute of Artificial Intelligence, Beihang University · Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing, Beihang University · Beijing Academy of Artificial Intelligence (BAAI) +1