Multi-Task Cooperative Learning via Searching for Flat Minima

TL;DR

This paper introduces a novel multi-task learning framework that promotes cooperative feature learning by searching for flat minima, leading to improved performance in medical image analysis tasks.

Contribution

It formulates multi-task learning as a multi/bi-level optimization problem to enable cooperative learning and searches for flat minima to reduce negative transfer.

Findings

Outperforms state-of-the-art MTL methods on three datasets

Demonstrates the effectiveness of cooperative learning in MTL

Shows improved generalizability of features in medical image analysis

Abstract

Multi-task learning (MTL) has shown great potential in medical image analysis, improving the generalizability of the learned features and the performance in individual tasks. However, most of the work on MTL focuses on either architecture design or gradient manipulation, while in both scenarios, features are learned in a competitive manner. In this work, we propose to formulate MTL as a multi/bi-level optimization problem, and therefore force features to learn from each task in a cooperative approach. Specifically, we update the sub-model for each task alternatively taking advantage of the learned sub-models of the other tasks. To alleviate the negative transfer problem during the optimization, we search for flat minima for the current objective function with regard to features from other tasks. To demonstrate the effectiveness of the proposed approach, we validate our method on three publicly available datasets. The proposed method shows the advantage of cooperative learning, and yields promising results when compared with the state-of-the-art MTL approaches. The code will be available online.