Multitask learning NN only trains on a few tasks

Question

I am fairly new to the Machine Learning field, but I am working on building a multi-task network for a project I am working on. Unfortunately, despite all my best googling and digging through documentation, blogs and research papers, I am so far unable to get my network training on more than a few of the tasks at a time.

For context, the problem entails a set of images of metallic test samples, all taken under standard lighting and at a standard distance. For confidentiality reasons, I cannot share the dataset itself. The goal of the network is to be able to score 3 different surface defects in 3 distinct regions each, a total of 9 outputs, each with an integer rating score from 0 to 5 (6 possible outcomes for each output).

I have been writing the NN in Python using Keras and visualising the metrics and graph with TensorBoard.

The problem is that, as I attempt to train the network, it only (at best) seems to learn for up to 4 of the 9 tasks, even though each image has a rating for each task. To illustrate my point, here are two screenshots of the TensorBoard metrics for one of the more successful training runs along with the graph showing the network structure:

As you can see, the tasks WCP_3T and WCP_5T are training well, but nothing else really is (the high accuracy on the WCP_F task is due to heavily weighted data, and not much else I'm afraid). The plateau in accuracy for each of the other tasks seems to me to be the network quickly finding the rating that is most common for that task, and remaining there. Thus, the varying accuracies for each task is an artifact of the data for that category.

I have tried deepening (and shallowing) the other two main branches (BS and BD), I have experimented with altering learning rate, batch size, optimizer function (currently using Adam) and activation functions (currently using relu for the WCP branch and LeakyReLU for the other two, but I have also tried ordinary relu for them with no benefit). I have altered the number of filters in the various convolutional layers, and I have increased and decreased the dropout on various layers also. Still, nothing I seem to try can bring the network to train on all the tasks.

It is possible that the size of the dataset is a problem, but I only have a small number of images available and I don't think it is possible to get many more. I am currently using a training set of 506 images and a batch size of 32, which only leaves about 170 images for each testing and validation. Data augmentation to provide more images may also be out of the question as that could conceivably alter the scores placed on the different defects, and I would not know how (they were graded by experts, and the data was then passed to me for this task).

If anyone has had a similar problem before and could suggest something else I could try, or could at least point me towards somewhere that might help to solve my problem, I would be extremely grateful!

Answer 1

So, it turns out that my problem had a lot more to do with network structure than anything else. I was trying to share too many lower levels, and then still not leaving enough room for the later branches to learn individual features/locations. In the end, I re-structured my network so that there was only 1-2 shared convolutional layers, then each primary branch shared 1 further convolutional layer before again each branching into 3 sub-branches (one for each output), each with their own convolutional, dense, flatten and then softmax layers. I also possibly had too many filters on each convolutional layer.

This enabled 'training' across each of the 9 outputs, but has left me with some serious overfitting to try to overcome, due to the depth of the network and the small data set. I am currently trialing methods of correcting for this.