Optimal Transport for Deep Joint Transfer Learning

TL;DR

This paper introduces a joint transfer learning method using optimal transport loss to improve deep neural network training with limited target data, explicitly modeling relatedness between source and target tasks.

Contribution

The paper proposes a novel joint fine-tuning approach with OT loss that explicitly incorporates task relatedness, enhancing transfer learning efficiency.

Findings

JTLN outperforms standard fine-tuning on image classification tasks.

Using different cost matrices allows encoding prior knowledge about task relatedness.

The method is general and applicable to various neural network architectures.

Abstract

Training a Deep Neural Network (DNN) from scratch requires a large amount of labeled data. For a classification task where only small amount of training data is available, a common solution is to perform fine-tuning on a DNN which is pre-trained with related source data. This consecutive training process is time consuming and does not consider explicitly the relatedness between different source and target tasks. In this paper, we propose a novel method to jointly fine-tune a Deep Neural Network with source data and target data. By adding an Optimal Transport loss (OT loss) between source and target classifier predictions as a constraint on the source classifier, the proposed Joint Transfer Learning Network (JTLN) can effectively learn useful knowledge for target classification from source data. Furthermore, by using different kind of metric as cost matrix for the OT loss, JTLN can incorporate different prior knowledge about the relatedness between target categories and source categories. We carried out experiments with JTLN based on Alexnet on image classification datasets and the results verify the effectiveness of the proposed JTLN in comparison with standard consecutive fine-tuning. This Joint Transfer Learning with OT loss is general and can also be applied to other kind of Neural Networks.