Cross-stitch-Networks-for-Multi-task-Learning

This project is a TensorFlow implementation of a Multi Task Learning method described in the paper Cross-stitch Networks for Multi-task Learning.

Arguments

• --lr, learning rate
• --n_epoch, number of epoch
• --n_batch_size, mini batch size
• --reg_lambda, L2 regularization lambda
• --keep_prob, Dropout keep probability
• --cross_stitch_enabled, Use Cross Stitch or not

Dataset

Fashion-MNIST

Fashion-MNIST is a dataset of Zalando's article images, consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28x28 grayscale image, associated with a label from 10 classes:

Label	Description	Label	Description
0	T-shirt/top	5	Sandal
1	Trouser	6	Shirt
2	Pullover	7	Sneaker
3	Dress	8	Bag
4	Coat	9	Ankle boot

For multi task learning, I created another label for each image, which is based on the original labels:

Label	Original Labels	Description
0	5, 7, 9	Shoes
1	3, 6, 8	For Women
2	0, 1, 2, 4	Other

The network will train these two classifiers together.

Network

Without task sharing

As a baseline, a network without cross stitch is built, which simply concats two convolutional neural networks side by side. Each network is for one task, although their parameters are not shared. The final loss function is the sum of two loss functions of sub networks.

Here is an overview of this structure:

Both sub convolutional neural networks have the same architecture:

Layer	Output size	filter size / stride
conv1	28x28x32	3x3 / 1
pool_1	14x14x32	2x2 / 2
conv2	14x14x64	3x3 / 1
pool_2	7x7x64	2x2 / 2
fc_3	1024
output	10 or 3 depends on task

With Cross Stitch

Cross Stitch is a transformation applied between layers, it describes the relationship between different tasks with a linear combination of their activations.

The network should learn the relationship by itself, in comparison with manually tuning the shared network structure, this end-to-end approach works better.

Here is an overview of this structure:

The convolutional sub networks have the same architecture as above. As in paper suggested the cross stitch units are only added after Pool layers and Fully Connected layers.

Training

• The input images are standardized with z-score.
• L2 regularization is used for convolution layers and fully connected layers, lambda = 1e-5.
• Dropout has keep_prob = 0.8
• Batch normalization is used
• Weights of sub networks are initialized with He initialization
• Weights of Cross Stitch are initialized with identity matrix (i.e no sharing between tasks at the beginning)
• Learning rate is set to a constant value 0.001
• Trained 30 epochs with batch size = 128

Evaluation

The overall accuracy is calculated by averaging the accuracies of all sub tasks.

With cross stitch transformation it gets more than 1% improvement on test dataset.

Orange: without sharing. Blue: with cross stitch.

Result

For Fashion-MNIST new labels are created based on the original labels, so two classification tasks are highly related. I also used this technique to build a gender-age classifier with VGGFace2 dataset, which labels are more independent. In both tests cross stitch improves the accuracy. Although this project only trained with two tasks but it can be extended to more tasks easily.

I didn't pretrain the sub networks as in paper suggested and I also used a different initialization strategy. A better result might be found with more tuning.