Improving Multi-task Learning via Seeking Task-based Flat Regions

• concerns about efficiency:
  • SAM is computationally expensive, doubling the computation cost compared with ERM/SGD. In Algorithm 1, each task requires computing the SAM gradient for shared/non-shared parameters. In total, the algorithm needs at least $2m$ gradient calculations, where $m$ is the number of tasks. Hence, the algorithm is computationally inefficient.
  • In experiments, there are no results (like training time) for comparing efficiency or training curve (performance w.r.t. training time).
  • this problem will be very serious when there are many tasks, e.g., the QM9 data set has 11 tasks.
  • some suggestions for mitigating this issue: use efficient variants of SAM, e.g.,
    • AE-SAM, An Adaptive Policy to Employ Sharpness-Aware Minimization, ICLR 2023
    • ESAM, Efficient Sharpness-aware Minimization for Improved Training of Neural Networks, ICLR 2022
• Eq(4) in Theorem 1, $[...]\_{i=1}^m \leq \max [...]\_{i=1}^m$ means ?
• Theorem 1 can be directly obtained from Theorem 1 of Foret et al. (2021): decomposing the parameters into two parts and using different $\rho$'s.
• in "Update the shared part" (P5), "However, a direct gradient aggregation ... can be negatively affected by the gradient cancelation or conflict because it aims to combine many individual elements with different objectives", a direct gradient aggregation means? not clear
• Why the proposed aggregation in Section 4.4 is better than the above "direct gradient aggregation"?
• In the Conclusion Section, "proving that they can help enhance previous works both theoretically," which theorem(s)?
• how to calculate the entropy in Figure 2, note that the entropy in the figure has negative values.
• Figure 1, the "2-task problem", where is the definition?

The statement of Theorem 1 is too intuitive, which does not require rigorous proof on it. At the right side of Theorem 1, maximum of maximum is utilized for deriving the upper bound. It is intuitive based on my knowledge.

The analytical decomposition of SAM gradient into 1) loss and 2) flatness parts are not novel at all. It is well known analysis based on existing methods (SAM, GSAM, GAM). Rather, The new modeling parts of SAM-MTL is gradient decomposition on each task and gradient aggregation based on whole tasks. However, i do not get convinced why these gradient decomposition and re-organization are required in the context of multi-task learning. This is not empirically validated by additional ablation studies.

In Figure 4, the author claim that suggested algorithms significantly improves the task-wise flatness than ERM algorithm. What if we conduct simple SAM on MTL, not based on your gradient decomposition and re-organization? I conjecture that the flatness would be similar to SAM-MTL, your method. The extensive comparison with SAM variants (SAM,GSAM,GAM) is required.

Please empirically provide the computation cost increments by applying SAM-MTL. SAM is well known for increasing the computation cost about 2 times than ERM. is there any other increments during the adaptation of SAM-MTL?

Please see questions below.

1. Second point in strengths.

2. The proofs and theorems—which the authors claim to be a major contribution of the present work and on which the proposed algorithm is supposedly based—are poorly presented. In turn, without which, the proposed approach is largely heuristic and lack theoretical support (excluding results that have been established in prior work, i.e., the constituent component of SAM and gradient-based MTL methods).

   • The "mild assumptions" are not clearly stated nor justified. E.g., theorem 2 used the assumption that the loss function is bounded by $L$, which is not mentioned anywhere except in the proof. Also, please justify and elaborate on the assumption that "that adding Gaussian perturbation will raise the test error": is it required for all $\theta$, or local minima? It would be best if the assumptions are listed explicitly.

   • The conclusion of theorem 3 looks wrong. First of all, in the proof, the induction is incorrectly applied—the $\xi$ cannot alter between cases. The $\log1/\delta$ term in $f^i$ should be $\log m/\delta$. And, does the conclusion not follow theorem 2 directly via a simple union bound?

   • The outer $\max _ {\\|\epsilon_\textrm{sh}\\|<\rho_\textrm{sh}}$ in the statement of Theorem 1 and 3 does not make sense to me. The max is taken over a vector of $m$ dimensions. Is the max coordinate-wise? If so, it should go inside the square bracket. If not, is the max well-defined? Or, how is the total order of the vector space defined?

3. Regardless of the above potential issue with the theorem statement, I fail to see the connection between Theorem 1 (or its complete version 3) and the approach in section 4.3, i.e., the idea that "the worst-case shared perturbation $\epsilon_\mathrm{sh}$ is commonly learned for all tasks". Specifically, how is computing the worst-case perturbation on each task separately and then aggregate the gradients $\\{g^{i,\textrm{SAM}}_\textrm{sh}\\} _ {i\in m}$ related to the idea above?

4. As mentioend in point 1 of strengths, there are some grammatical issues and weird word choice that may lead to confusions. E.g., what is the "ultimate gradient descent direction" (in the abstract)? Also, "is the compliment set" --> "is the complement set".