Learning Compact Neural Networks with Deep Overparameterised Multitask
  Learning
Compact neural network offers many benefits for real-world applications.
However, it is usually challenging to train the compact neural networks with
small parameter sizes and low computational costs to achieve the same or better
model performance compared to more complex and powerful architecture. This is
particularly true for multitask learning, with different tasks competing for
resources. We present a simple, efficient and effective multitask learning
overparameterisation neural network design by overparameterising the model
architecture in training and sharing the overparameterised model parameters
more effectively across tasks, for better optimisation and generalisation.
Experiments on two challenging multitask datasets (NYUv2 and COCO) demonstrate
the effectiveness of the proposed method across various convolutional networks
and parameter sizes.
Shen Ren, Haosen Shi
# Learning Compact Neural Networks with Deep Overparameterised Multitask Learning

###### Abstract

Compact neural network offers many benefits for real-world applications. However, it is usually challenging to train the compact neural networks with small parameter sizes and low computational costs to achieve the same or better model performance compared to more complex and powerful architecture. This is particularly true for multi-task learning, with different tasks competing for resources. We present a simple, efficient and effective multitask learning overparameterisation neural network design by overparameterising the model architecture in training and sharing the overparameterised model parameters more effectively across tasks, for better optimisation and generalisation. Experiments on two challenging multitask datasets (NYUv2 and COCO) demonstrate the effectiveness of the proposed method across various convolutional networks and parameter sizes.

## 1 Introduction

Deep Multi-task Learning (MTL) techniques are widely applied to real-world embedded computer vision applications. A deep MTL model explores and exploits the synergies among multiple tasks to be learned simultaneously to improve the joint performance, as well as to reduce inference time and computational costs. However, designing efficient MTL models on budgeted devices poses two major challenges. First, the model design needs to be efficient to stay compact for meeting the computational budget constraints. Second, the model needs to be effective on resource sharing among multiple tasks learned simultaneously to avoid resource competition.

Motivated by previous research on deep linear network [14] that even the input-output map can be rewritten as a shallow network, it nevertheless demonstrates highly nonlinear training dynamics and can help to accelerate optimisation [1] and improve generalisation [1]. To tackle the aforementioned challenges, we propose an overparameterised MTL method by initialising the parameters of each shared neural network layer as the product of multiple matrices following the spatial Singular Vector Decomposition (SVD)[1]. The left and right singular vectors are trained with all task losses, and the diagonal matrices are trained using task-specific losses. Our design is mainly inspired by analytical studies on overparameterised networks for MTL [1] that the training/test error dynamics depends on the time-evolving alignment of the network parameters to the singular vectors of the training data, and a quantifiable task alignment describing the transfer benefits among multiple tasks depends on the singular values and input feature subspace similarity matrix of the training data.

In this work, we follow the definition of overparameterisation [1], referring to the replacement of neural network layers by operations of the compositions of multiple layers with more learnable parameters, but without adding additional expressiveness of the network. The contribution of this work can be summarised as follows.

* We propose an MTL neural network design with overparameterised training components and a compact inference architecture, applicable for embedded applications with limited computational budgets.
* We implement an iterative training strategy for the proposed design that is effective and efficient for the multitask computer vision dense prediction tasks, compared to the state-of-the-art.

## 2 Methodology

We replace the fully-connected layers and/or convolutional layers of modern neural networks with overparameterisation, and share the overparameterised parameters among different tasks, to achieve higher performance for reduced inference parameter size and computational cost.

Specifically, tensor decomposition is used for model expansion instead of model compression during training. The full-rank diagonal tensor is further expanded to be trained separately for each task, while the other tensors are shared among all tasks. During inference, the decomposed tensors are contracted back into a compact MTL architecture.

### Overparameterisation Mechanism

**Fully-connected Layers**

For any shared layer of a deep MTL model, given a weight matrix \(W\) that is shared among \(t\) tasks, we directly factorise the weight matrix \(W\) of the size \(m\times n\) using SVD, similar to [13] and [1], so that
\[W:=UMV \tag{1}\]

\(U\) is of size \(m\times r\) and V is of size \(r\times n\), and \(M\) is a diagonal matrix of size \(r\times r\) where the matrix is full rank \(r\geq\min(m,n)\).

#### 2.1.2 Convolutional Layers

A shared convolutional layer of a deep MTL model is parameterised as \(W\in\mathbb{R}^{c_{o}\times c_{i}\times k\times k}\) to mathematically denote a 2-dimensional convolutional layer of kernel size \(k\times k\) with \(c_{o}\) output channels and \(c_{i}\) input channels.

Using tensor factorisation on the convolutional layer following the spatial SVD format, we first change the dimensions of the weight tensor into \(c_{o}\times k\times k\times c_{i}\), and replace the original weight tensor into three tensors \(U\), \(M\) and \(V\) using Eq. (1) with the size \(c_{o}\times k\times r\), \(r\times r\) and \(r\times k\times c_{i}\) respectively, where \(r\geq\min(c_{o}\times k,c_{i}\times k)\).

### Parameter Sharing Mechanism

Considering each shared fully-connected or convolutional layer of an MTL model with the objective of learning tasks \(a\) and \(b\) together, after overparameterisation, the parameters \(U\) and \(V\) are shared across all tasks, and \(M^{(a)},M^{(b)}\) are assigned as task-specific parameters for the corresponding tasks. The task-specific parameters \(M^{(a)}\) and \(M^{(b)}\) are learned as scaling factors in changing the scales of shared parameters \(U\) and \(V\) according to each individual task, to better align with the singular vectors of the training data. The product is cumulative, associative and distributive so that the sequence of the tasks will not take effect on the final product.

The designed sharing mechanism can be extended naturally to multitask learning of more than \(2\) tasks by adding task-specific diagonal matrices, \(M^{(1)},\cdots,M^{(t)}\). For \(t\) numbers of tasks, the diagonal matrix \(M\) can be expanded and shared across all tasks as

\[M:=M^{(1)}\circ\cdots\circ M^{(t)} \tag{2}\]

where \(M\) is computed from the matrices \(M^{(1)},\cdots,M^{(t)}\) and \(\circ\) is the Hadamard product or standard matrix product, which are the same for the diagonal matrices.

The overparameterisation and the parameter sharing mechanisms are shown in Figure 1.

### Iterative Training Strategy

During the training phase of the overparameterised MTL model, instead of training the weight matrix \(W\), the factorised matrices \(U\), \(M^{(1)}\) to \(M^{(t)}\) and \(V\) are trained. The \(U\) and \(V\) matrices are initialised using the same initialisation method as the original parameter matrix, and \(M^{(1)}\) to \(M^{(t)}\) are initialised into identity matrices. The trained weight matrices are contracted back to \(W\) according to Eq. (1) for inference with fewer parameter counts and computational costs.

In order to train shared and task-specific parameters separately, we propose an iterative training strategy which consists of two training processes for each epoch of training, as shown in Algorithm 1

1. Choose a small subset of the training data (\(\sim 3\%\) in our experiments), for each task \(j\) from all \(t\) numbers of tasks to be learned together, only train its task-specific factors \(M^{(j)}\) by task loss \(L^{(j)}\), where \(1\leq j\leq t\).
2. All task-specific factors \(M^{(0)}\) to \(M^{(t)}\) are frozen. The other factors \(U\), \(V\) and the parameters of other unfactorised layers are trained by a multi-task learning loss \(L=\sum_{i\in t}\alpha_{i}L^{(i)}\), where \(\alpha\) represents the fixed or adaptive loss weights.

```
0: Initialised parameters: \(\theta\), loss weights \(\alpha\) while not convergence do  sample a mini-batch data for\(i\in\{1,\dots,t\}\)do  update \(M^{(i)}\) by \(L^{(i)}\) endfor  update \(\theta-\{M^{(1)},\dots,M^{(t)}\}\) by \(L=\sum_{i\in t}\alpha_{i}L^{(i)}\) endwhile
```

**Algorithm 1** The Iterative Training Process

During each training process, Frobenius decay is applied as a penalty on factorised matrices \(U\), \(V\) and \(M\) to regularise the models for better generalisation, implemented as in [11]. The Frobenius decay rate is set to \(1\)e-4.

## 3 Experiments

Experiments are conducted on public datasets NYUv2 [21] and COCO [15] including a series of ablation studies and comparisons. Note that as the proposed overparameterized model is contracted back to its original size for inference, the FLOPS, latency, and parameter size during inference following the proposed methods remain identical to the original model without overparameterization.

### Semantic Segmentation

To first demonstrate the effectiveness of matrix factorisation as an overparameterisation method, we compare the method with other state-of-the-art overparameterisation methods, including RepVGG [4] and ExpandNet [19], to perform a single Semantic Segmentation task. All experiments are conducted on NYUv2 dataset using SegNet model [1], a smaller SegNet model (tagged with -S in the result tables) with all filter sizes for all convolutional layers (except for the final layer) cut into half. The overparameterisation is applied to all convolutional layers except for the final layer. The implementation of SegNet and all experiments on NYUv2 dataset follow the same training hyperparameters and configurations as in [15]. The batch size is \(8\). The total training epoch is \(200\). A learning rate scheduler is used to reduce the learning rate to

Figure 1: Overparameterisation and Parameter Sharing Mechanisms
half at the \(150\)th epoch. Note that we use the multi-branch architecture including additional \(1\times 1\) convolutional and identity branches proposed in RepVGG to overparameterise the SegNet model, though it is not intended to be used as a drop-in replacement. Results in Table 1 show that the matrix factorisation method in this paper outperforms all other methods.

### Semantic Segmentation, Depth Estimation and Surface Normal Estimation

The proposed parameter-sharing mechanisms and iterative training strategy for MTL are experimented on NYUv2 dataset using SegNet, small SegNet and DeepLabv3 [3] with ResNet50 backbone (tagged with -ResNet). For some experiments (tagged with Pretrain), the ResNet50 backbone is pre-trained with ImageNet 1k dataset. All convolutional layers in the shared backbone are overparameterised for SegNet models. For DeepLabv3 model, all the first convolutional layers in ResNet bottleneck modules are overparameterised following the similar design in YOLOv7 [22]. The initialisation of the overparameterised parameters follows the spectral initialisation introduced in [13]. The performance is compared with multiple MTL methods including MTL baseline with only shared backbones and task-specific heads, cross-stitch [17] and MTAN [14]. The batch size for MTAN and the proposed method applied to MTAN is set to \(2\), and learning rate is \(5\)e\(-5\). Note that the task-specific diagonal matrices are trained without weight decay in the optimiser.

Ablation studies are also conducted to compare the proposed method (**Fac**) with a series of variants:

* The shared backbone is overparameterised according to Figure 1, but all factors are trained with the combined task losses.
* Split the channels factorised in factors \(U\) and \(V\) evenly across tasks, and stack the split matrix back to recover \(U\) and \(V\).
* Split the core diagonal matrix \(M\) evenly across tasks, and stack the split matrix back to recover \(M\).

Experiment results are shown in Table 3. The proposed method outperforms all baselines except for some tasks under the cross-stitch model, which has increased expressiveness by enlarging the model size (\(\sim 3\) times larger in inference model parameter size). The proposed method showcases comparable or superior performance in comparison to a backbone model pretrained with the ImageNet dataset, even when the backbone structure experiences changes due to overparameterisation, resulting in the loss of the pretrained weights

### Instance and Semantic Segmentation

We also evaluate the performance of our proposed method using Panoptic Feature Pyramid Network (PanopticFPN) [15] to perform instance segmentation and semantic segmentation tasks simultaneously. We follow the same experiment configurations described in [15] using COCO dataset. By overparameterise the PanopticFPN and following the iterative training strategy, our model also performs better than the Baseline model without pretraining on other datasets, as shown in Table 2.

## 4 Related Works

After investigating the advantages of optimisation in overparameterised linear and nonlinear networks in [1], several studies have proposed overparameterised deep learning models based on fully-connected or convolutional neural networks. However, none of these studies designed their overparameterisation under the context of MTL, accounting for the effective parameter sharing among multiple tasks. Among these studies, ExpandNet [1], DO-Conv [1] and Overcomplete Knowledge Distillation [13] designed different methods to overparameterise the kernels of convolutional layers, either over their channel axes [1] or spatial axes [1]. A more recent design named RepVGG [16] is a special case of overparameterisation, which instead of designing a drop-in replacement for convolutional layers of any neural architecture as previous studies, presented an effective convolutional neural network model with an overparameterised training network and a VGG-like inference network.

## 5 Conclusion

In this paper, we propose a parameter-sharing scheme and an iterative training for deep multitask learning that effectively share parameters using overparameterised models during training, while the model architecture stayed slim and compact during inference. Compared to the state-of-the-art, the scheme has demonstrated its potential in various datasets and model architecture. The design is particularly well-suited for embedded computer vision applications where there are tight budget constraints in both memory and computational resources.

## Acknowledgements

This study is supported under the RIE2020 Industry Alignment Fund - Industry Collaboration Projects (IAF-ICP) Funding Initiative, as well as cash and in-kind contribution from the industry partner(s).

\begin{table}
\begin{tabular}{l c c} \hline \hline \multirow{2}{*}{**Method**} & \multicolumn{2}{c}{**SemSeg**} \\ \cline{2-3}  & mIoU \(\uparrow\) & Pix Acc \(\uparrow\) \\ \hline SegNet Baseline & 17.85 & 55.36 \\
**Matrix Factorisation** & **19.74** & **56.42** \\ RepVGG & 19.26 & 56.30 \\ ExpandNet & 16.27 & 51.31 \\ SegNet-S Baseline & 17.22 & 54.29 \\
**SegNet-S Matrix Factorisation** & **19.58** & **56.28** \\ SegNet-S RepVGG & 17.25 & 54.26 \\ SegNet-S ExpandNet & 16.59 & 51.44 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Matrix Factorisation and SOTA Comparison on NYUv2

\begin{table}
\begin{tabular}{l c c c c c c c c} \hline \hline \multirow{2}{*}{**Method**} & \multicolumn{2}{c}{**SemSeg**} & \multicolumn{2}{c}{**Bias**} & \multicolumn{2}{c}{**Instance Seg**} & \multicolumn{2}{c}{**PanopticSeg**} \\ \cline{2-9}  & mIoU \(\uparrow\) & Pix Acc \(\uparrow\) & \(\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{ \mathbf{\mathbf{\mathbf{ \mathbf{ \mathbf{ }}}}}}}}}}}}}\) & \(\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{ \mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbfmathbfmathbfmathbfmathbfmathbfmathbfmathbfmathbf{ } }}}}}}}}}}}}}}\) & & & \\ \hline Baseline & 21.34 & 60.15 & 11.30 & 11.18 & 10.89 & 9.88 & 16.59 & 61.69 & 21.35 \\
**Fac** & **23.44** & **70.22** & **15.62** & **15.11** & **15.11** & **15.97** & **20.94** & **64.40** & **25.64** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Instance and Semantic Segmentation using PanopticFPN