# Less is More – Towards parsimonious multi-task models using structured
sparsity

Richa Upadhyay1, Ronald Phlypo2, Rajkumar Saini1, Marcus Liwicki1
1Luleå University of Technology, Sweden, 2University Grenoble Alpes, France
<EMAIL_ADDRESS><EMAIL_ADDRESS>rajkumar.saini@ltu,se,
<EMAIL_ADDRESS>
###### Abstract

Model sparsification in deep learning promotes simpler, more interpretable
models with fewer parameters. This not only reduces the model’s memory
footprint and computational needs but also shortens inference time. This work
focuses on creating sparse models optimized for multiple tasks with fewer
parameters. These parsimonious models also possess the potential to match or
outperform dense models in terms of performance. In this work, we introduce
channel-wise $l_{1}/l_{2}$ group sparsity in the shared convolutional layers
parameters (or weights) of the multi-task learning model. This approach
facilitates the removal of extraneous groups i.e., channels (due to $l_{1}$
regularization) and also imposes a penalty on the weights, further enhancing
the learning efficiency for all tasks (due to $l_{2}$ regularization). We
analyzed the results of group sparsity in both single-task and multi-task
settings on two widely-used Multi-Task Learning (MTL) datasets: NYU-v2 and
CelebAMask-HQ. On both datasets, which consist of three different computer
vision tasks each, multi-task models with approximately 70% sparsity
outperform their dense equivalents. We also investigate how changing the
degree of sparsification influences the model’s performance, the overall
sparsity percentage, the patterns of sparsity, and the inference time.

##  1 Introduction

(a) Unstructured

(b) Structured

(c) Channel-wise

Figure 1: An illustration of the various forms of sparsity that may be
introduced to a parameter vector of a CNN. (blue denotes non-zero values).
Figure (a) depicts unstructured sparsity, (b) shows structured (block or
group) sparsity, and (c) represents channel-wise structured (group) sparsity.

The principle of parsimony states that a model with a lower number of
parameters is preferred over a more intricate model with a higher number of
parameters, given that both models fit the data equally well [1]. Model
sparsification, one of the methods to obtain parsimonious models and model
compression, holds significant importance in the fields of Machine Learning
(ML) and Deep Learning (DL), often being explored through feature or parameter
selection techniques. Although, in recent times, the concept of over-
parameterization [2] i.e., model having more parameters than necessary to fit
its training data and DL have become intertwined [3]. The pursuit of sparse
models does not necessarily conflict with the concept of over-
parameterization. In large-scale, complex scenarios, the benefit of sparse
models is manifold. They provide increased interpretability, reduce
overfitting, efficient computation, and aid in identifying the most
informative features, leading to an efficient learning process [4]. While
over-parameterization enables networks to approximate complex mappings and
simplify the loss function, making optimization easier, model sparsification
provides a different perspective that prioritizes efficiency and
interpretability, both of which can be critical in specific applications. The
contrast between the efficiency of sparse models and the comprehensive
functionality of over-parameterized models showcases the continuous evolution
of DL. Continuing along these lines, our work delves into the adoption of
structured sparsity to attain parsimonious models. We focus on identifying and
leveraging the most significant features or parameters of a model, thereby
balancing the efficiency of sparse models with the vast capabilities of deep
(over-parameterized) models.

Another technique for reducing model size is parameter sharing, typically
utilized in the context of MTL. Sparsity within an MTL framework, which seeks
to train multiple tasks concurrently [5], can yield significant advantages. It
is because, in MTL, where the model complexity increases as the number of
tasks increases, a simple (sparse) model can be more interpretable and
computationally efficient. Moreover, only some features might be relevant to
all the tasks and oversharing of knowledge (features) among the tasks may
cause negative information transfer [6]. As a result, sparse models can help
to choose significant features for specific tasks while also assisting in
learning shared representations across multiple tasks. Therefore, the
motivation for this study is to leverage the advantages of structured sparsity
to optimize MTL models.

This work, therefore, introduces structured (group) sparsity in MTL, i.e., it
learns sparse shared features among multiple tasks. There are two important
reasons for this integration. First is the organization of parameters in a CNN
is inherently grouped into layers, channels, or filters; these naturally
grouped parameters provide an opportunity to apply structured sparsity [7].
Another reason is CNNs, especially deeper ones, can develop redundant filters
that extract similar features [8, 9]. This is because the number of filters in
deep CNNs is usually thousands, and it is inevitable that there exist a lot of
similar filters that extract the same or similar features [8]. Structured
sparsity optimizes these networks by targeting and pruning entire redundant
filters or channels of the weight matrix rather than just individual weights.
Figure 1 illustrates the types of sparsity induced in a 4D parameter tensor
($\text{number of
filters}\times\text{channels}\times\text{height}\times\text{width}$). In this
work, we introduce channel-wise group sparsity in the shared network among all
the tasks, eventually eliminating a significant amount of shared parameters
(i.e., groups or channels zeroed out). As a result, decreasing the memory
footprint of the model will lead to less computational expense at the time of
inference. Our approach is based on the understanding that the dense
prediction tasks vary in the granularity and type of features they require.
Some tasks might necessitate more low-level features (e.g., edge or texture
details for semantic segmentation) while others could demand more high-level
representations (like object or scene understanding for depth estimation).
Channel-wise sparsity may allow for the retention of task-specific channels
that are most pertinent while pruning the less relevant ones. The main
contributions of this paper are as follows:

  1. 1.

Introducing structured (group) sparsity in an MTL framework, particularly
channel-wise $l_{1}/l_{2}$ penalty to the shared (CNN) layer parameters to
solve complex computer vision tasks. Along with reducing the number of shared
parameters, it aids in significantly improving task performance.

  2. 2.

Design of an experiment framework to analyze the effect of group sparsity in
single-task and multi-task settings.

  3. 3.

A deep analysis of the model’s performance across different levels of
sparsification, taking into account the performance of the tasks, group
sparsity, parameter reduction, and inference time.

  4. 4.

A comparative analysis of the structured ($l_{1}/l_{2}$) vs unstructured
($l_{1}$) sparsity.

This article is organized as follows: The research works that used structured
sparsity in the context of MTL are discussed in Section 2. Section 3
introduces multi-task learning, the concept of $l_{1}/l_{2}$ group sparsity,
and presents the approach proposed in this work. Section 4 gives a detailed
experimental set-up, while the results, followed by an extensive performance
analysis, are presented in Section 5. At last, Section 6 draws the conclusions
and future work.

##  2 Related work

Figure 2: This block diagram is a simple illustration of the proposed work.
Here, all the shared network layers aim for sparsity, as their channel-wise
parameters are grouped and subjected to the $l_{1}$/$l_{2}$ penalty. The
objective is to learn sparse multi-task shared representations that help to
enhance performance across all tasks. Simultaneously, the task-specific
networks focus on learning representations tailored for each individual task.
This multi-task network is trained using backpropagation on a composite loss
derived from the individual losses of each task.

Parameter-sharing in MTL enables models to exploit similarities across various
tasks, enhancing generalization and learning efficiency. There are two types
[5]: hard parameter-sharing and soft parameter-sharing. Hard parameter-sharing
is caused by the network design, which shares some parameters across all tasks
[10, 11, 12]. In contrast, soft parameter-sharing encourages models to have
comparable but distinct parameters for the shared layers by including a
regularization term in the loss function that penalizes the variations in the
parameters [13]. However, as tasks increase, the parameters and, consequently,
the computational cost grow proportionally. Notably, not all tasks require
complex networks and static architectures might lead to suboptimal results and
often overfitting. Numerous studies in the past have optimized multi-task
network structures for each task to overcome these issues. In this section, we
review recent studies emphasizing parameter efficiency in MTL. We also explore
the literature on structured sparsity both in general contexts and
specifically within MTL.

Parameter-efficient MTL:  MTL heavily uses Neural Architecture Search (NAS)
[14] to optimize network design based on task relatedness and complexity of
numerous tasks. Most pruning solutions start with a sizeable multi-task
network and then propose a layer-wise architecture search process. [15]
tailored residual networks to develop a task-specific ensemble of sub-networks
in different depths, which are learned based on task similarities. Other
approaches begin with a small network and incrementally expand the
architecture, like [16], which starts with a thin multi-layer network and
dynamically uses greedy algorithms to widen it during training. In their work,
[17] proposed utilizing a differentiable tree-structured topological space to
optimize parameters and branching distributions that enable automatic search
of a hard parameter-sharing multi-task network. A method for learning the
layer-sharing pattern over several tasks is proposed in [18], which involves
first learning a task-specific policy distribution and then randomly selecting
a select-or-skip policy decision from that distribution. Several other works,
such as [19, 20, 21], propose different approaches to learning task-specific
sparsity inducing ‘masks’ for the parameter vectors, which help to reduce the
number of trainable parameters in a multi-task setting.

Concept of parameter sparsity in DL: [7] is a very detailed review that
discusses the types of sparsity and their application in DL. Broadly, they
categorize sparsity into two types: model sparsity, which involves pruning
weights or neurons, and ephemeral sparsity, which is applied on a per-instance
basis and deactivates neurons or weights, such as with dropout or activation
functions. Numerous studies in the literature have adapted these
sparsification methods for various applications; we highlight a few of these,
particularly ones involving structured sparsity. The work presented in [22])
introduced block-sparse regularization in the form of $l_{1}/l_{q}$ norms [23]
to acquire a low-dimensional representation that can be shared across a group
of related tasks. They proved that penalizing the trace norm of the parameter
matrix to make it low rank makes the non-convex group lasso optimization
problem convex. A combination of block-sparse and element-wise sparse
regularization is proposed in [24] to improve the block-sparse models since
their performance depends on the extent to which features are shared across
tasks. Similarly, [25] introduced channel-wise, stripe-wise, and group-wise
filter pruning as a means of introducing sparsity in Deep Neural Network
(DNN). The issue of addressing the estimation of multiple linear regression
equations and variable selections within a multi-task framework is discussed
in [26]. Therefore, many studies utilize structured sparsity in the context of
MTL, focusing on variable selection, identifying task-specific layers, and
learning sparse binary masks for the parameters of each task.

Positioning our work: While most of the works in the field of MTL employ the
concept of sparsity for feature selection and usually use synthetic data for
simple regression tasks [22, 24, 26], and learn sparse (unstructured) masks
for individual tasks [19]; our work employs model sparsity and ephemeral
sparsity while training a multi-task model for various heterogeneous tasks
[27], such as image-level tasks (e.g., classification) and pixel-level tasks
(e.g., segmentation, depth estimation). [28] is very close to what we have
done; however, they propose a fusion regularization term in their work for
task grouping and demonstrate the performance on synthetic data. Our primary
contribution lies not in the introduction of sparsity itself; instead, it
centres on the investigation and applicability of the group sparsity in MTL
for dense prediction heterogeneous computer vision tasks. In the context of
this work, model sparsity is applied in the form of channel-wise
$l_{1}$/$l_{2}$ regularization as shown in Figure 2, which involves pruning of
shared model weights during training, it impacts not only the forward pass of
the multi-task model but also its inference in terms of performance,
complexity, and prediction time. The shared CNN (i.e., ResNet, discussed in
Section 4) has ReLU activation functions after every convolutional layer. This
naturally introduces ephemeral sparsity per data point during training, which
only impacts the training stage [7]. Consequently, in this study, we
concentrate solely on the effects of model sparsity. It is important to note
that our work centres on the sparsification of hard-shared parameters,
specifically the backbone parameters typically shared by all tasks. However,
this concept can also be adapted for the soft-parameter sharing setting.

##  3 Methodology

In this section, we begin by establishing the foundational concepts of our
work, delving into MTL and group sparsity-inducing penalties. Subsequently, we
introduce our proposed approach.

Multi-task learning (MTL) : MTL prioritizes the simultaneous training of
multiple tasks instead of isolated training. This approach leverages common
knowledge and representations across multiple tasks, often leading to enhanced
performance and generalization [29, 27]. Consider a simple multi-task
architecture that demonstrates hard parameter sharing as in Figure 2, that
shares a common backbone network between many tasks and each task having a
separate task-specific network. For $N$ non-identical but related tasks
sampled from a task distribution $\mathcal{T}=\\{T_{i}\\}_{i=1}^{N}$, let
$\theta_{b}$ be the shared parameters of the backbone network (shared layers),
while $\\{\theta_{i}\\}_{i=1}^{N}$ be the task-specific parameters for $N$
tasks, such that $\theta_{b}\cap\\{\theta_{i}\\}_{i=1}^{N}=\emptyset$. In MTL,
the objective is to minimize the combined loss $L_{comb}$ of all the tasks by
finding the optimal network parameters
$\theta=\theta_{b}\cup\\{\theta_{i}\\}_{i=1}^{N}$. It can be expressed as:

$\theta^{*}=\operatorname*{arg\,min}_{\theta}L_{comb}(\theta,D^{tr}),$ (1)

where $D^{tr}=\\{D^{tr}_{i}\\}_{i=1}^{N}$ represents the training dataset of
$N$ tasks. Consider a function $\mathcal{F}$ that illustrates how the losses
from all the tasks are combined for being back-propagated to the multi-task
network, i.e., $L_{comb}=\mathcal{F}(L_{1},L_{2},..,L_{N})$. Various
techniques for aggregating the losses across multiple tasks have been
discussed in [30]. This work employs the concept of uncertainty weighing [31]
to balance the multiple losses. It can be mathematically represented as:

$L_{comb}=\sum_{i=1}^{N}\left(\dfrac{1}{2\sigma_{i}^{2}}~{}L_{i}(\theta_{b},\theta_{i},D^{tr}_{i})+log(\sigma_{i})\right),$
(2)

where $\\{\sigma_{i}\\}_{i=1}^{N}$ are learnable parameters that are optimized
along with the model parameters $\theta$ to minimize the combined loss.

Structured sparsity inducing penalty:  The group lasso penalty [32], which
combines the $l_{1}$ and $l_{2}$ norms (hereafter referred to as $l_{1}/l_{2}$
in this manuscript), is a form of regularization that induces structured
sparsity [33]. Its objective is to generate solutions that eliminate entire
groups of variables, in contrast to the most frequently used $l_{1}$ norm,
which results in unstructured sparsity. The group $l_{1}/l_{2}$ penalty is
introduced as a regularization term along with the loss function; the
optimization objective is then expressed as:

$\min_{\theta}~{}~{}L(\theta,D^{tr})+R(\theta),~{}~{}\text{where}~{}~{}R(\theta)=\lambda\sum_{g=1}^{G}\sqrt{n_{g}}~{}||\theta_{g}||_{2}.$
(3)

$R(\theta)$ is the $l_{1}/l_{2}$ penalty term which is the $l_{1}$ norm
(promotes sparsity) of the $l_{2}$ norm (penalizes weights) over the non-
overlapping parameter groups $(g)$, and $G$ is the total number of groups.
Here $n_{g}$ is the group size, and $\lambda$ is the regularization parameter
or strength or degree of regularization. The $l_{2}$ norm is given by
$||\theta_{g}||_{2}=\sqrt{\sum_{j=1}^{n_{g}}\theta_{j}^{2}}$. Since the
penalty term is non-differentiable, in order to minimize the objective in Eq.
3, proximal gradient descent updates of the parameters are required, as
explained in [34]. The proximal updates are based on the gradients of the
differentiable part of the composite loss function, i.e.,

$\theta_{t+1}\leftarrow prox_{\alpha
R}(\theta_{t}-\alpha\nabla_{\theta}L(\theta_{t},D^{tr})),$ (4)

where $prox_{\alpha R}$ is the proximal operator, and $\alpha$ is the learning
rate. In spite of its non-differentiable nature, the proximal operator for the
$l_{1}/l_{2}$ penalty term can be efficiently computed in a closed form [35],
given by:

$prox_{\alpha
R}(\theta_{g})=\begin{cases}\left[1-\dfrac{\alpha\lambda\sqrt{n_{g}}}{||\theta_{g}||_{2}}\right]\theta_{g}~{}~{};~{}~{}||\theta_{g}||_{2}>\alpha\lambda\sqrt{n_{g}}\\\
~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}0~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{};~{}~{}~{}||\theta_{g}||_{2}\leq\alpha\lambda\sqrt{n_{g}}\end{cases}.$
(5)

Given that the parameters are partitioned into $G$ dis-joint groups, the
proximal operator likewise adheres to this decomposition and thus can be
implemented on each parameter group. According to Eq. 5, if the norm of the
parameter group is less than the strength of regularization (i.e., a function
of regularization parameter $\lambda$, learning rate $\alpha$, and the number
of elements in the group $n_{g}$), the entire group of parameters is set to
zero resulting in a sparse solution.

Proposed approach: In this work, we apply group sparsity to the shared
parameters $\theta_{b}$ of a multi-task network. Given that the foundation for
MTL and group sparsity has been laid out (discussed above), we can now
seamlessly delve into their integration. The parameter vector of a
convolutional layer is a 4-D vector of dimension $(N,C,H,W)$, where $N$ is the
number of filters, $C$ is the number of channels, $H$ and $W$ are the spatial
height and width. The channel-wise structure sparsity is introduced in each
convolutional layer of the backbone by considering each channel as a group, as
demonstrated in [36]. Therefore, the multi-task optimization objective can now
be written as (from Eq. 1 and 3):

$\displaystyle\min_{\theta}~{}~{}L_{comb}(\theta,D^{tr})+\lambda\sum_{g=1}^{G}\sqrt{n_{g}}~{}||\theta_{b_{g}}||_{2}$
(6)

where
$||\theta_{b_{g}}||_{2}=\sum_{l=1}^{L}\left[\sum_{n_{c}=1}^{C_{l}}||\theta_{b}^{l}(:,n_{c},:,:)||_{2}\right]$.
Here $L$ is the total number of convolutional layers in the backbone network,
and $C_{l}$ is the number of channels in the $l^{th}$ layer. In Eq. 6, the
group lasso term i.e.,
$\lambda\sum_{g=1}^{G}\sqrt{n_{g}}~{}||\theta_{b_{g}}||_{2}$ zeroes out an
entire channel group, i.e., $\theta_{b}^{l}(:,n_{c},:,:)$ if its norm is below
the degree of regularization as per Eq. 5, rendering the backbone parameters
sparse. Therefore, $l_{1}$ sparsity leads to structured feature selection,
while $l_{2}$ sparsity term promotes within-group regularization, ensuring
that all features in that group are either jointly important or jointly
unimportant. While the combined multi-task loss term i.e.,
$L_{comb}(\theta,D^{tr})$ fosters the joint learning of multiple tasks, such
that the tasks help each other to learn better. The reduction of trainable
parameters in a multi-task model through dynamic sparsification during its
training phase can result in benefits during inference, such as decreased
memory usage, computation requirements, and prediction time, as well as
enhanced performance.

##  4 Experimental setup

Datasets and Tasks: We evaluate the concept of structured sparsity on multi-
task learning on two publicly available datasets, i.e., the NYU-v2 dataset
[37] and the CelebAMask-HQ dataset [38] (referred to as celebA dataset in this
work). In the NYU-v2 dataset, which contains images of indoor scenes, three
dense pixel-level tasks are chosen: semantic segmentation (comprising 40
classes), depth estimation, and surface normal estimation. For the celebA
dataset, which is a large-scale face image dataset, two binary classification
tasks of male/female and smile/no smile identification are considered. A
pixel-level semantic segmentation task has also been selected, featuring three
classes: skin, hair, and background. The celebA dataset is usually used in MTL
to investigate task inter-dependencies.

Network architecture and hyperparameters: A three-channel RGB image of size
$256\times 256$ is fed to a dilated ResNet-50 [39] backbone network in batches
of 16(32) for NYU-v2(celebA) dataset. The output of the backbone are shared
representations which are then given to task-specific networks (Figure 2). A
Deeplab-V3 network [40] is employed as the task head for the dense prediction
tasks. While the task-specific network for classification tasks uses a minimal
network with a convolution layer and two linear fully connected layers. This
work employs cross-entropy loss for the semantic segmentation task, inverse
cosine similarity loss for surface normal estimation, means square loss for
depth estimation, and binary-cross entropy loss for classification. Adaptive
optimizer ‘Adam’ [41] is used for optimizing the non-differentiable objective
function (Eq. 3) by calculating adaptive proximal gradients (explained in
[34]). In order to ensure adequate training for all tasks, a learning rate
$\alpha$ of 0.0001 is used. The regularization parameter $\lambda$ is a
crucial hyper-parameter in this work, as it determines the degree of sparsity.
Therefore, we conducted an ablation study considering a range of $\lambda$
values, i.e., $[1\times 10^{-7},1\times 10^{-6},1\times 10^{-5},1\times
10^{-4},1\times 10^{-3},1\times 10^{-2}]$. The dataset was divided into non-
overlapping training, validation, and test set; for NYU-v2 dataset the same
data split as [18, 42] is followed, while for the celebA dataset, we randomly
split the data in 60-20-20%. A uniform test set and similar hyper-parameters
are used to compare performance across all experiments. The training of the
models is conducted on NVIDIA A100 Tensor Core GPUs, equipped with 40 GB of
onboard HBM2 VRAM. To assess the consistency of the model, all experiments
were conducted five times using distinct random seeds. The outcomes are
presented in the form of the mean and standard deviation. The source code for
reproducibility can be found at https://github.com/ricupa/Less-is-More-
Towards-parsimonious-multi-task-models-using-structured-sparsity.git.

##  5 Results and discussion

To evaluate the performance of the proposed approach, two types of experiments
were designed: single-task and multi-task experiments (mostly all possible
task combinations are considered). Table 1 and Table 2 display the performance
of the various single-task and multi-task experiments for the NYU-v2 and
celebA datasets, respectively. Table 1 displays the results for three values
of $\lambda$, i.e., $0$ (no sparsity), $1\times 10^{-6}$ & $1\times 10^{-5}$,
while Table 2 for $\lambda=0$ (no sparsity) and $1\times 10^{-5}$.

Table 1: Single task and multi-task (test set) performance for all the three
tasks on the NYU-v2 dataset. The % group sparsity represents the fraction of
groups that are eliminated after training. The values of %group sparsity and
parameter sparsity are the mean values over 5 trials. Here IoU is the
intersection over union, CS is the cosine similarity, and MAE is the mean
absolute error; the $\uparrow$ represent a higher value is better and
$\downarrow$ represent a lower value is preferable. The values denoted in
italics indicate the optimal performance achieved across all tasks in a
single-task configuration, incorporating group sparsity. On the contrary, the
values highlighted in bold represent the top two performances attained in a
multi-task setup, also utilizing group sparsity.
Experiments | Lambda | Tasks | In the backbone network
---|---|---|---
$\lambda$ | Segmentation | Surface Normal | Depth |  % group sparsity | % parameters
| IoU $\uparrow$ | CS $\uparrow$ | MAE $\downarrow$ | reduced to zero
Semantic | 0 | 0.3140 ± 0.0350 |  |  | 0 | 0
Segmentation | $1\times 10^{-6}$ | 0.3203 ± 0.0096 |  |  | 0 | 0
| $1\times 10^{-5}$ | 0.3338 ± 0.0068 |  |  | 35.77 | 43.05
Depth | 0 |  |  | 0.1645 ± 0.0016 | 0 | 0
estimation | $1\times 10^{-6}$ |  |  | 0.1572 ± 0.0017 | 21.95 | 31.11
| $1\times 10^{-5}$ |  |  | 0.1625 ± 0.0015 | 89.90 | 93.57
Surface normal | 0 |  | 0.7077 ± 0.0047 |  | 0 | 0
estimation | $1\times 10^{-6}$ |  | 0.7903 ± 0.0041 |  | 17.87 | 20.42
| $1\times 10^{-5}$ |  | 0.7699 ± 0.0058 |  | 55.62 | 64.39
segmentation + | 0 | 0.2233 ± 0.0094 |  | 0.1656 ± 0.0030 | 0 | 0
depth estimation | $1\times 10^{-6}$ | 0.3308 ± 0.0049 |  | 0.1379 ± 0.0034 | 0 | 0
| $1\times 10^{-5}$ | 0.3345 ± 0.0073 |  | 0.1389 ± 0.0027 | 66.06 | 73.58
segmentation + | 0 | 0.2276 ± 0.0064 | 0.7072 ± 0.0087 |  | 0 | 0
surface normal | $1\times 10^{-6}$ | 0.3353 ± 0.0082 | 0.7797 ± 0.0128 |  | 13.51 | 15.34
estimation | $1\times 10^{-5}$ | 0.3480 ± 0.0180 | 0.7833 ± 0.0058 |  | 64.05 | 72.97
surface normal | 0 |  | 0.6905 ± 0.0037 | 0.1719 ± 0.0038 | 0 | 0
estimation + | $1\times 10^{-6}$ |  | 0.7609 ± 0.0279 | 0.1299 ± 0.0030 | 60.62 | 69.83
depth estimation | $1\times 10^{-5}$ |  | 0.7967 ± 0.0042 | 0.1332 ± 0.0015 | 88.62 | 93.38
segmentation + | 0 | 0.2171 ± 0.0128 | 0.6948 ± 0.0080 | 0.1657 ± 0.0034 | 0 | 0
surface normal + | $1\times 10^{-6}$ | 0.3418 ± 0.0062 | 0.7814 ± 0.0108 | 0.1301 ± 0.0064 | 21.80 | 23.85
depth estimation | $1\times 10^{-5}$ | 0.3394 ± 0.0107 | 0.7857 ± 0.0076 | 0.1336 ± 0.0029 | 69.72 | 76.72

Figure 3: Comparison of the number of (non-zero) channels per convolution
layer before (blue) and after sparsity (orange). The names of dilated
ResNet-50 convolution layers are on the horizontal axis, and on the vertical
axis are the number of channels. For all these plots, the value of
$\lambda=1\times 10^{-5}$.

(a) performance vs $\lambda$

(b) % sparsity vs $\lambda$

(c) sparsity profile

(d) inference time vs $\lambda$

Figure 4: These figures demonstrate (a) the variations in the task’s test
performance, (b) group and parameter sparsity percentage, (c) sparsity profile
during training, and (d) mean inference time for the multi-task scenario
involving all three tasks of the NYU-v2 dataset, across a spectrum of
regularization parameter values $\lambda$. The term “% group sparsity" refers
to the proportion of eliminated groups, while “% parameter sparsity" indicates
the ratio of parameters assigned a value of zero. Table 2: Single task and
multi-task (test set) performance for all the three tasks on the celebA
dataset. The % group sparsity represents the fraction of groups that are
eliminated after training. Here, IoU represents the metric intersection over
union, and the $\uparrow$ represents that a higher value of the metric (IoU
and accuracy) is better. The % group sparsity and % parameter sparsity (i.e.,
reduce to zero) contain the mean values obtained from five trials. The values
highlighted in bold represent the top two performances attained in a multi-
task setup, incorporating group sparsity.

|  | Tasks | In the backbone network
---|---|---|---
|  |  | classification |  |
|  | Segmentation | Male/female | smile/no smile |  |
Experiments | Lambda | IoU $\uparrow$ | Accuracy $\uparrow$ | Accuracy $\uparrow$ |  % group sparsity |  % parameters reduce to zero
| 0 | 0.8885 ± 0.0048 |  |  | 0 | 0
segmentation | $1\times 10^{-5}$ | 0.9084 ± 0.0033 |  |  | 58.1 | 71.94
| 0 |  | 0.6296 ± 0.0019 |  | 0 | 0
male/female | $1\times 10^{-5}$ |  | 0.6284 ± 0.0029 |  | 87.32 | 92.31
| 0 |  |  | 0.5204 ± 0.0026 | 0 | 0
smile/no smile | $1\times 10^{-5}$ |  |  | 0.5238 ± 0.0004 | 100 | 100
| 0 | 0.8363 ± 0.0028 | 0.9103 ± 0.0217 |  | 0 | 0
seg + male | $1\times 10^{-5}$ | 0.9136± 0.0024 | 0.9726± 0.0035 |  | 67.66 | 77.77
| 0 |  | 0.6284 ± 0.0108 | 0.5242 ± 0.0006 | 0 | 0
male + smile | $1\times 10^{-5}$ |  | 0.7275 ± 0.0016 | 0.5308 ± 0.0103 | 63.18 | 78.1
seg +male + | 0 | 0.8123 ± 0.0024 | 0.8641 ± 0.0269 | 0.5238 ± 0.0012 | 0 | 0
smile | $1\times 10^{-5}$ | 0.9040± 0.0035 | 0.8514± 0.1578 | 0.5241± 0.0760 | 77.1 | 83.41

Group sparsity enhances the performance:  For both datasets, the outcomes of
the experiments show that even with approximately 70% sparsity, the multi-task
models perform better than the non-sparse models; this is also true in the
case of single-task settings. Upon assessing the task performance for the
NYU-v2 dataset, as presented in Table 1, it is apparent that the
implementation of group sparsity yields a notable improvement in the
performance of all tasks, especially in the multi-task setting. For the
semantic segmentation task, the best IoU score is achieved with a $\lambda$
value of $1\times 10^{-5}$ in both single task (0.3338 ± 0.0068) and the
multi-task settings (0.3480 ± 0.0180 in combination with surface normal
estimation). The lowest MAE (mean absolute error) for the depth estimation
task is achieved with a $\lambda$ value of $1\times 10^{-6}$ (0.1572 ± 0.0017)
for single-task learning. In the multi-task setting, when depth estimation is
combined with surface normal estimation, the performance improved
significantly, with the lowest MAE being 0.1299 ± 0.0030 at a lambda of
$1\times 10^{-6}$. For surface normal estimation, in a single-task framework,
the best result (CS) is for $\lambda=1\times 10^{-6}$ (0.7903 ± 0.0041).
However, the best performance throughout is when it is combined with depth
estimation, specifically at a $\lambda=1\times 10^{-5}$ (0.7967 ± 0.0042). It
is evident that combining tasks can boost performance, but the effect is task-
dependent. Segmentation performs well on its own or when supplemented with
surface normal estimation, while depth estimation gains significantly from the
surface normal estimation task (closely related tasks, with the latter often
being derived from the former). In fact, combining depth estimation and
surface normal estimation results in better performance and much higher group
sparsity than when paired with semantic segmentation. In the absence of
sparsity, i.e., $\lambda=0$, the individual task performance surpasses that of
any of the multi-task performances across all tasks. The reason behind this
can probably be negative information transfer between tasks due to oversharing
of information. The introduction of group sparsity in multi-task experiments
yields a significant improvement in the performances of all tasks. Thus, it
can be inferred that the implementation of group sparsity within the shared
layers regulates the dissemination of information across tasks. Specifically,
the utilization of $l_{2}$ regularization in the penalty term outlined in
Equation 3 facilitates soft parameter sharing, while the application of
$l_{1}$ regularization promotes sparse hard parameter sharing. Thus, the
outcomes of the proposed approach show the effectiveness of group sparsity in
MTL.

Figure 3 illustrates the groups of parameters before and after pruning for all
the convolution layers of the shared backbone network. Comparing single-task
plots to multi-task plots shows how MTL facilitates learning shared features,
and sparsity eliminates unimportant parameter groups. In multi-task model
plots, initial layers are predominantly active, sharing low-level features
with more parameters assigned to them than to later layers. Conversely, most
intermediate layers are sparse, with only a few deep layers actively sharing
high-level features. Such a sparsity pattern suggests that group lasso
preserves vital network structures, such as residual connections of the ResNet
backbone network.

Similar findings can also be asserted for the celebA dataset, as presented in
Table 2. The tasks of segmentation and male/female classification exhibit
significant enhancements with sparsity and in MTL setting. The accuracy of the
smile/no smile classification task remains consistent (under-performing) even
with the introduction of sparsity. As expected, combining image-level
male/female classification with pixel-level semantic segmentation tasks
significantly improves the performance of the classification task. When all
the tasks are trained together, even with 77% sparsity, the tasks perform
better or sometimes equivalent to their no sparse and single-task experiments.
These results prove the significance of both MTL and group sparsity.

% Sparsity-performance trade-off:  Figure 4(a) illustrates the average test
performance of all three tasks in a multi-task setting when subjected to
various regularization strengths (i.e., $\lambda$ on the horizontal axis).
While Figure 4(b) displays the amount of % group sparsity and % parameter
sparsity for different values of $\lambda$. It can be concluded from both
these figures that performance improvements are observed as sparsity increases
to a certain point, beyond which performance begins to decline. As the value
of lambda rises, sparsity also increases, leading to improved task performance
(i.e., IoU and CS $\uparrow$, while MAE $\downarrow$) up to $\lambda=1\times
10^{-5}$; beyond this point, any further increase in lambda, or consequently
in % sparsity, results in declining performance. It is noteworthy that
increasing $\lambda$ from $1\times 10^{-6}$ to $1\times 10^{-5}$ results in
substantial growth in sparsity, while the performance metrics remain
relatively stable i.e., no significant change. The model exhibits complete
sparsity (i.e., all parameters groups are equal to zero) for higher values of
$\lambda$ $(1\times 10^{-3}$ and $1\times 10^{-2}$), resulting in two
implications. Firstly, the absence of shared parameters contradicts the notion
of MTL. Secondly, the model is so simple (sparse) that it cannot acquire
adequate knowledge, leading to underfitting. The present ablation study
examined the impact of increasing regularization intensity on overall task
performance. Therefore, there is a trade-off between task performance and the
degree of sparsity.

Faster inference time:  Model sparsity results in a simple model with fewer
parameters for inference; therefore, an increase in sparsity decreases
inference time (specifically the runtime for backbone only), as shown in
Figure 4(d). In this work, the mean inference time represents the time taken
by the backbone or shared network to process a batch of data, determined over
five experimental trials using different batches of the same size to validate
the robustness of the results. For $\lambda=1\times 10^{-6}$, optimal
performance corresponds to approximately 21% group sparsity, reducing the
inference time by 0.32ms compared to the dense network. In contrast, at
$\lambda=1\times 10^{-5}$ with around 70% group sparsity, despite a slight dip
in performance, the network’s speed increases by 0.74ms. Therefore, the choice
of $\lambda$ modulates the balance between model performance and inference
speed.

Dynamic sparsity: In this work we apply dynamic sparsity that involves
continuously adjusting which weights are pruned during training, while static
sparsity is a fixed pruning of neural network weights after training. Figure
4(c) shows the sparsification profiles for different values of $\lambda$ while
training all three tasks together for the NYU-v2 dataset. Smaller $\lambda$
values initially yield a dense model that gradually becomes sparse, improving
performance. Conversely, larger $\lambda$ values lead to high sparsity from
the start of training, limiting performance. Starting with all parameters and
then progressively pruning them is akin to synaptic pruning in the human brain
[7]). Just as the network begins with a rich set of dense parameters and
sparsifies over time, the brain initially forms an excess of neural
connections. As development proceeds, less vital synapses are eliminated,
optimizing the brain for efficiency and environmental adaptation.

Structured vs Unstructured sparsity: We also compared the performance of
$l_{1}/l_{2}$ (structured) sparsity with $l_{1}$ (unstructured) sparsity, the
results are shown in Table 3. Although applying solely $l_{1}$ sparsity
achieves similar levels of parameter sparsity, the performance with
$l_{1}/l_{2}$ sparsity, as shown in Table 1, remains superior. Fine-tuning the
regularization parameter ($\lambda$) could potentially lead to comparable or
improved performance and even greater parameter sparsity. However,
unstructured sparsity often faces challenges in hardware efficiency since it
is a fine-grain approach that involves removing parameters randomly (without a
pattern) based on some criteria. Conversely, while structured sparsity aligns
better with hardware optimization, it typically faces limitations in achieving
high levels of sparsity without adversely impacting performance because it is
a coarse-grained approach for pruning structured blocks of parameters [7]. The
choice between structured and unstructured sparsity depends on the use case,
balancing between performance and hardware efficiency. Furthermore, as shown
in Table 3, unstructured sparsity can also lead to the zeroing out of some
channels, thus contributing to a notably low percentage of group sparsity.

Table 3: Performance of MTL with $l_{1}$ sparse regularization on the NYU-v2
Dataset. The regularization parameter ($\lambda$) is set at $1\times 10^{-3}$,
which yields a level of parameter sparsity comparable to that achieved with
$l_{1}/l_{2}$ sparsity.

MTL Experiments | Segmentation | Surface normal est. | Depth est. | % group | % parameter
---|---|---|---|---|---
| IoU $\uparrow$ | CS $\uparrow$ | MAE $\downarrow$ | sparsity | sparsity
Segmentation + Surface Normal est. | 0.2962 ± 0.0050 | 0.7421 ± 0.0048 | - | 0.59 | 74.31
Depth est. + Surface Normal est. | - | 0.7285 ± 0.0096 | 0.1545 ± 0.0011 | 5.46 | 77.78
Segmentation + Depth est. + | 0.2900 ± 0.0258 | 0.7389 ± 0.0187 | 0.1520 ± 0.0052 | 1.68 | 69.92
Surface Normal est.

###

In general, the outcomes presented in this section for both datasets
demonstrate the efficacy of group sparsity in MTL. The aforementioned approach
is effective as it leverages the principle of inductive bias, wherein tasks in
MTL share features and mutually reinforce learning. Additionally, the
incorporation of $l_{1}/l_{2}$ group sparsity facilitates the removal of
redundant parameters that do not contribute to any of the tasks while also
regularizing the parameters. This approach yields a parsimonious backbone
model with reduced parameters that accommodates all the tasks within a multi-
task framework.

##  6 Conclusion and future scope

We present an approach for developing parsimonious models by employing dynamic
group sparsity in a multi-task setting. Through extensive experiments, we
demonstrate that sparse multi-task models perform as well as or better than
their dense counterparts. So, sparsifying the model during training yields a
more general model that offers faster inference times. Our proposed method can
be integrated with any multi-task models that undergo gradient-based training.
Furthermore, it is adaptable to various tasks, be they classification,
regression, or otherwise. Therefore, we put forth a model-agnostic and task-
agnostic approach for developing simple and interpretable multi-task models.
In this study, the regularization factor ($\lambda$), which is a hyper-
parameter, determines the level of sparsity. Finding the ideal value of
$\lambda$ is a challenging task, though. Therefore, a promising future
research direction might involve developing an approach for learning the
optimal value of $\lambda$ while training, potentially leading to enhanced
performance and optimal sparsity.

## Acknowledgements

The authors are grateful to the National Supercomputer Centre at Linköping
University for their access to the Berzelius supercomputing system and
acknowledge the generous support of the Knut and Alice Wallenberg Foundation.

## References

  * Vandekerckhove et al. [2015] Joachim Vandekerckhove, Dora Matzke, and Eric-Jan Wagenmakers.  Model comparison and the principle of parsimony.  _Oxford Handbook of Computational and Mathematical Psychology_ , pages 300–317, 2015.
  * Allen-Zhu et al. [2019] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International conference on machine learning_ , pages 242–252. PMLR, 2019.
  * He and Tao [2020] Fengxiang He and Dacheng Tao.  Recent advances in deep learning theory.  _arXiv preprint arXiv:2012.10931_ , 2020.
  * Le Thi Phuong and Geskus [2018] Thao Le Thi Phuong and Ronald Geskus.  A comparison of model selection methods for prediction in the presence of multiply imputed data.  _Biometrical Journal_ , 61, 10 2018.  doi:10.1002/bimj.201700232.
  * Crawshaw [2020] Michael Crawshaw.  Multi-task learning with deep neural networks: A survey.  _arXiv preprint arXiv:2009.09796_ , 2020.
  * Wu et al. [2020] Sen Wu, Hongyang R. Zhang, and Christopher Ré.  Understanding and improving information transfer in multi-task learning.  In _International Conference on Learning Representations_ , 2020.
  * Hoefler et al. [2021] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste.  Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks.  _J. Mach. Learn. Res._ , 22(1), jan 2021.  ISSN 1532-4435.
  * Geng and Niu [2022] Lili Geng and Baoning Niu.  Pruning convolutional neural networks via filter similarity analysis.  _Machine Learning_ , 111(9):3161–3180, 2022.
  * Kahatapitiya and Rodrigo [2021] Kumara Kahatapitiya and Ranga Rodrigo.  Exploiting the redundancy in convolutional filters for parameter reduction.  In _proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 1410–1420, 2021.
  * Long et al. [2017] Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Philip S. Yu.  Learning multiple tasks with multilinear relationship networks.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , NIPS’17, page 1593–1602, Red Hook, NY, USA, 2017. Curran Associates Inc.  ISBN 9781510860964.
  * Lu et al. [2016] Yongxi Lu, Abhishek Kumar, Shuangfei Zhai, Yu Cheng, Tara Javidi, and Rogério Schmidt Feris.  Fully-adaptive feature sharing in multi-task networks with applications in person attribute classification.  _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 1131–1140, 2016.
  * Ruder et al. [2019] Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard.  Latent multi-task architecture learning.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 33(01):4822–4829, Jul. 2019.  doi:10.1609/aaai.v33i01.33014822.
  * Duong et al. [2015] Long Duong, Trevor Cohn, Steven Bird, and Paul Cook.  Low resource dependency parsing: Cross-lingual parameter sharing in a neural network parser.  In _Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 2: Short Papers)_ , pages 845–850, Beijing, China, July 2015. Association for Computational Linguistics.  doi:10.3115/v1/P15-2139.
  * Elsken et al. [2019] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter.  Neural architecture search: A survey.  _The Journal of Machine Learning Research_ , 20(1):1997–2017, 2019.
  * Wang et al. [2022] Tianxin Wang, Fuzhen Zhuang, Ying Sun, Xiangliang Zhang, Leyu Lin, Feng Xia, Lei He, and Qing He.  Adaptively sharing multi-levels of distributed representations in multi-task learning.  _Information Sciences_ , 591:226–234, 2022.  ISSN 0020-0255.  doi:https://doi.org/10.1016/j.ins.2022.01.035.
  * Lu et al. [2017] Yongxi Lu, Abhishek Kumar, Shuangfei Zhai, Yu Cheng, Tara Javidi, and Rogerio Feris.  Fully-adaptive feature sharing in multi-task networks with applications in person attribute classification.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 1131–1140, 2017.  doi:10.1109/CVPR.2017.126.
  * Guo et al. [2020] Pengsheng Guo, Chen-Yu Lee, and Daniel Ulbricht.  Learning to branch for multi-task learning.  In _Proceedings of the 37th International Conference on Machine Learning_ , ICML’20. JMLR.org, 2020.
  * Sun et al. [2020a] Ximeng Sun, Rameswar Panda, Rogerio Feris, and Kate Saenko.  Adashare: Learning what to share for efficient deep multi-task learning.  _Advances in Neural Information Processing Systems_ , 33:8728–8740, 2020a.
  * Sun et al. [2020b] Tianxiang Sun, Yunfan Shao, Xiaonan Li, Pengfei Liu, Hang Yan, Xipeng Qiu, and Xuanjing Huang.  Learning sparse sharing architectures for multiple tasks.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, pages 8936–8943, 2020b.
  * Bai et al. [2022] Ting Bai, Yudong Xiao, Bin Wu, Guojun Yang, Hongyong Yu, and Jian-Yun Nie.  A contrastive sharing model for multi-task recommendation.  In _Proceedings of the ACM Web Conference 2022_ , pages 3239–3247, 2022.
  * Ding et al. [2021] Ke Ding, Xin Dong, Yong He, Lei Cheng, Chilin Fu, Zhaoxin Huan, Hai Li, Tan Yan, Liang Zhang, Xiaolu Zhang, and Linjian Mo.  Mssm: A multiple-level sparse sharing model for efficient multi-task learning.  In _Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval_ , SIGIR ’21, page 2237–2241, New York, NY, USA, 2021. Association for Computing Machinery.  ISBN 9781450380379.  doi:10.1145/3404835.3463022.
  * Argyriou et al. [2006] Andreas Argyriou, Theodoros Evgeniou, and Massimiliano Pontil.  Multi-task feature learning.  In B. Schölkopf, J. Platt, and T. Hoffman, editors, _Advances in Neural Information Processing Systems_ , volume 19. MIT Press, 2006.
  * Zhang and Huang [2008] Cun-Hui Zhang and Jian Huang.  The sparsity and bias of the Lasso selection in high-dimensional linear regression.  _The Annals of Statistics_ , 36(4):1567 – 1594, 2008.  doi:10.1214/07-AOS520.
  * Jalali et al. [2010] Ali Jalali, Sujay Sanghavi, Chao Ruan, and Pradeep Ravikumar.  A dirty model for multi-task learning.  In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, _Advances in Neural Information Processing Systems_ , volume 23. Curran Associates, Inc., 2010.
  * Meng et al. [2020] Fanxu Meng, Hao Cheng, Ke Li, Huixiang Luo, Xiaowei Guo, Guangming Lu, and Xing Sun.  Pruning filter in filter.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 17629–17640. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ccb1d45fb76f7c5a0bf619f979c6cf36-Paper.pdf.
  * Lounici et al. [2009] Karim Lounici, Massimiliano Pontil, Alexandre B Tsybakov, and Sara Van De Geer.  Taking advantage of sparsity in multi-task learning.  _arXiv preprint arXiv:0903.1468_ , 2009.
  * Upadhyay et al. [2023a] Richa Upadhyay, Ronald Phlypo, Rajkumar Saini, and Marcus Liwicki.  Sharing to learn and learning to share – fitting together meta-learning, multi-task learning, and transfer learning: A meta review, 2023a.  arXiv:2111.12146.
  * Kshirsagar et al. [2017] Meghana Kshirsagar, Eunho Yang, and Aurélie C Lozano.  Learning task structure via sparsity grouped multitask learning.  _arXiv preprint arXiv:1705.04886_ , 2017.
  * Caruana [1997] Rich Caruana.  Multitask learning.  _Machine Learning_ , 28:41–75, 1997.
  * Vandenhende et al. [2021] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Luc Van Gool.  Multi-task learning for dense prediction tasks: A survey.  _IEEE transactions on pattern analysis and machine intelligence_ , 2021.
  * Kendall et al. [2018] Alex Kendall, Yarin Gal, and Roberto Cipolla.  Multi-task learning using uncertainty to weigh losses for scene geometry and semantics.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 7482–7491, 2018.
  * Yuan and Lin [2005] Ming Yuan and Yi Lin.  Model Selection and Estimation in Regression with Grouped Variables.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 68(1):49–67, 12 2005.  ISSN 1369-7412.  doi:10.1111/j.1467-9868.2005.00532.x.
  * Bach et al. [2012] Francis Bach, Rodolphe Jenatton, Julien Mairal, Guillaume Obozinski, et al.  Optimization with sparsity-inducing penalties.  _Foundations and Trends® in Machine Learning_ , 4(1):1–106, 2012.
  * Deleu and Bengio [2021] Tristan Deleu and Yoshua Bengio.  Structured sparsity inducing adaptive optimizers for deep learning.  _CoRR_ , abs/2102.03869, 2021.
  * Combettes and Wajs [2005] Patrick L. Combettes and Valérie R. Wajs.  Signal recovery by proximal forward-backward splitting.  _Multiscale Modeling & Simulation_, 4(4):1168–1200, 2005.  doi:10.1137/050626090.
  * Wen et al. [2016] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li.  Learning structured sparsity in deep neural networks.  In _Proceedings of the 30th International Conference on Neural Information Processing Systems_ , NIPS’16, page 2082–2090, Red Hook, NY, USA, 2016. Curran Associates Inc.  ISBN 9781510838819.
  * Nathan Silberman and Fergus [2012] Pushmeet Kohli Nathan Silberman, Derek Hoiem and Rob Fergus.  Indoor segmentation and support inference from rgbd images.  In _ECCV_ , 2012.
  * Lee et al. [2020] Cheng-Han Lee, Ziwei Liu, Lingyun Wu, and Ping Luo.  Maskgan: Towards diverse and interactive facial image manipulation.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2020.
  * Yu et al. [2017] Fisher Yu, Vladlen Koltun, and Thomas Funkhouser.  Dilated residual networks.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2017.
  * Chen et al. [2017] Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam.  Rethinking atrous convolution for semantic image segmentation, 2017.  arXiv:1706.05587.
  * Kingma and Ba [2017] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2017.  arXiv:1412.6980.
  * Upadhyay et al. [2023b] Richa Upadhyay, Prakash Chandra Chhipa, Ronald Phlypo, Rajkumar Saini, and Marcus Liwicki.  Multi-task meta learning: learn how to adapt to unseen tasks.  In _2023 International Joint Conference on Neural Networks (IJCNN)_ , pages 1–10, 2023b.  doi:10.1109/IJCNN54540.2023.10191400.

  *[MTL]: Multi-Task Learning
  *[CNN]: Convolutional Neural Network
  *[ML]: Machine Learning
  *[DL]: Deep Learning
  *[NAS]: Neural Architecture Search
  *[DNN]: Deep Neural Network


# Regularized Training of Nearest Neighbor Language Models

Jean-François Ton*
University of Oxford
<EMAIL_ADDRESS>
&Walter Talbott
Apple&Shuangfei Zhai
Apple&Josh Susskind
Apple

###### Abstract

Including memory banks in a natural language processing architecture increases
model capacity by equipping it with additional data at inference time. In this
paper, we build upon $k$NN-LM (Khandelwal et al., 2020), which uses a pre-
trained language model together with an exhaustive $k$NN search through the
training data (memory bank) to achieve state-of-the-art results. We
investigate whether we can improve the $k$NN-LM performance by instead
training a LM with the knowledge that we will be using a $k$NN post-hoc. We
achieved significant improvement using our method on language modeling tasks
on WIKI-2 and WIKI-103. The main phenomenon that we encounter is that adding a
simple L2 regularization on the activations (not weights) of the model, a
transformer, improves the post-hoc $k$NN classification performance. We
explore some possible reasons for this improvement. In particular, we find
that the added L2 regularization seems to improve the performance for high-
frequency words without deteriorating the performance for low frequency ones.
††footnotetext: ∗Work done as intern at Apple

##  1 Introduction

The problem of language modelling (LM) usually consists of two main
challenges. Firstly, mapping the context, i.e. the sentence prefixes, to a
vector representation, and secondly using this representation to predict the
subsequent word. In (Khandelwal et al., 2020), the authors claim that the
first problem is much easier to solve. Hence, given a pre-trained LM, they
post-hoc modify the representation using a k-nearest neighbor scheme ($k$NN)
and achieve significant improvements on challenging datasets, such as
WIKI-103.

Given that $k$NN improves the overall language modelling of a pre-trained
network, we examine training strategies that can make the underlying network’s
representations more amenable to the $k$NN step. Our results show improvements
over applying $k$NN to a generic LM network.

We first explore a simple learning scheme for the language model, where during
training we intentionally push representations that predict the same word
closer to together in the L2 sense, using a Momentum Contrastive (MOCO) He et
al. (2020) style implementation. We go on to note that this MOCO style
learning can be replaced by simply adding L2 regularization to the activation
of the layer used for $k$NN, eliminating implementation complexity. We then
present some initial experiments toward understanding why this L2
regularization brings improved performance.

##  2 Background

Our work builds upon $k$NN-LM (Khandelwal et al., 2020). In essence, $k$NN-LM
tackles the problem of how to improve a trained LM’s representations, and how
to adapt LMs to capture non-frequent sentences that are usually forgotten by
the model during training. $k$NN-LMs achieve significantly higher performance
through a simple interpolation between the original LM predictions and the
$k$NN predictions.

At inference time, given a new context sentence, $k$NN-LM works as follows:

  1. 1.

The context sentence $c_{i}$ is passed through the pre-trained network to
produce a representation $r^{context}_{i}\in\mathbf{R}^{d}$ as well as the
corresponding logits $y_{i}^{LM}$ to predict the next word.

  2. 2.

$r^{context}_{i}$ is used to find the $k$-nearest neighbors in the training
data. The logits $y^{kNN}$ are computed by a weighted average of the
neighbors’ labels, using the inverse exponential distance as the weight for
each neighbor.

  3. 3.

The logits are interpolated to give the final prediction:

$y_{final}=\lambda y^{kNN}+(1-\lambda)y^{LM},$

where $\lambda$ is the interpolation parameter that can be tuned on validation
data.

This simple post-hoc implementation allows (Khandelwal et al., 2020) to
improve upon the SOTA in LM by a significant margin. One thing to note about
$k$NN-LM is that they do not need to retrain the LM and hence the whole
algorithm can be run on CPU only. Furthermore, $k$NN-LM use FAISS, which is an
efficient library that allows them to quickly find $k$NNs.

One interesting detail to note in (Khandelwal et al., 2020) which was crucial
for this work, was that the authors tried both the inner product and the L2
for their distance metric in $k$NN. They concluded that L2 worked
significantly better. This observation implies the fact that the default
training recipe of LMs implicitly prefers one distance over another. We then
ask the question whether one could train in such a way that we could adapt the
model to the post-hoc $k$NN with a given distance metric. In the next section,
we describe how to adapt the training of the LM, with $k$NN and the
corresponding L2 distance metric in mind.

##  3 Proposed Method

In our initial attempt, we experimented with the idea of explicitly minimizing
the L2 distance between context vectors which predict the same target word.
This strategy directly mirrors the use of context vectors at the $k$NN step,
and we hope that training the representations in a way similar to testing will
further improve the effectiveness of $k$NN LM. However, a naive implementation
of it is infeasible. We then resorted to a MOCO (He et al., 2020) style
training scheme. Specifically, for each target word $w$, we construct a queue
$Q$ of fixed length $L$, which stores the recent $L$ context representations
for $w$. During training, we optimize a regularized objective as follows:

$\displaystyle\mathcal{L}_{new}=L_{CE}+\omega\sum_{j=1}^{\texttt{N}}\sum_{i=1}^{\texttt{L}}||sg(Q^{w_{j}}_{i})-r_{j}||^{2},$
(1)

where $N$ is the batch size, $r_{j}$ is the context representation of the
$j$th word, $Q^{w_{j}}$ is the queue corresponding to the $j$th target word
$w_{j}$; $\omega$ is the regularization parameter; $sg(\cdot)$ is the stop
gradient operator. Specifically, $Q$ is updated with a momentum target
encoding network which is initialized with the same parameters of the LM,
similar to MOCO (He et al., 2020).

Empirically, we found that Equation 1 provides a practical solution and yields
improved representations for the $k$nn LM, as shown in Fig 1. However, the use
of the queue and momentum target network still adds overhead to a large scale
model training. Hence we tried to decrease $Q$ and $L$, which interestingly
did not decrease the performance at all and therefore, to promote efficiency,
we tested an even simpler formulation, where we replace $Q$ with all zero
vectors. This eliminates the need of explicitly constructing and updating the
queue, while instead encourages the model to learning conservative
representations w.r.t. the L2 norms of its context representations. The
corresponding loss is as follows:

$\displaystyle\mathcal{L}_{new}=L_{CE}+\omega\sum_{j=1}^{\texttt{N}}||r_{j}||^{2}.$
(2)

To our surprise, Equation 2 yields similar performance to Equation 1 in
practice, while being much easier to implement and tune. This is a new finding
which we will try to explain in the ablation study. We thus use Equation 2 as
the default loss function in our experiments unless otherwise mentioned.

##  4 Experiments

We tested our method on the WIKI-2 and WIKI-103 datasets, which are widely
used benchmarks for language modelling. We are interested in demonstrating two
empirical results: improved performance using our approach over that of $k$NN-
LM, and exploring a possible mechanism for this improved performance.

Figure 1: Validation perplexity on WIKI-2 of the LM before (blue) and after
(orange) adding $k$NN. NOTE: weight=0 corresponds to the standard version that
does not include our added MOCO-style regularization term i.e. original $k$NN-
LM from Khandelwal et al. (2020)

|  $k$NN-LM $(\omega\\!=\\!0)$ | $\omega\\!=\\!0.1$ | $\omega\\!=\\!1.0$ | $\omega\\!=\\!10.0$
---|---|---|---|---
Train Ppl. LM | 19.99 | 20.05 | 20.11 | 21.37
Valid Ppl. LM | 75.96 | 75.68 | 76.37 | 81,29
Valid Ppl. $k$NN-LM | 74.11 | 73.13 | 70.63 | 80.52

Table 1: Experiments on WIKI-2 with corresponding validation perplexity using
L2 regularization.

###  4.1 Experimental setup

Dataset

WIKI-2 is a benchmark with 30k word vocabulary and consists of 2M tokens.
WIKI-103 is a benchmark with 250k word vocabulary and consisting of 103M
tokens.

Language Model Architecture For the language model architecture, we will be
using the exact setup as described in (Khandelwal et al., 2020). This setup
consists of the language model by (Baevski and Auli, 2018), which consists of
16 layers, each with 16 self-attention heads, 1024 dimensional hidden states,
and 4096 dimensional feedforward layers. Thus, following (Baevski and Auli,
2018), this LM has adaptive inputs and an adaptive softmax (Joulin et al.,
2017) with tied weights (Press and Wolf, 2016) for all our experiments. We
trained the each language model on a Tesla V100 with 400GB if RAM.

In addition, we follow the exact same training procedure as in Khandelwal et
al. (2020) and refer to their paper for further details on the training
parameters. The only difference in terms of implementation is the MOCO style
learner as well as the L2 regularization added to the final layer. Lastly, we
would like to note that while crossvalidating though the interpolation
parameter $\lambda$ we note that for all models, $\lambda=0.3$ works the best
which is in accordance to the finding in Khandelwal et al. (2020).

###  4.2 Experiments on WIKI-2

We first apply our proposed method on the standard WIKI-2 dataset, where we
run each configuration 5 times and plot the standard deviation, as seen in
Fig. 1. Note that $\omega=0$ in Fig. 1 corresponds to the standard $k$NN-LM
version, i.e. without the added term in the loss. Comparing Figure 1 and Table
1 to see that the MOCO and L2 approaches produces similar results. From these
results, we note the following phenomena:

  1. 1.

A clear "U"-shape demonstrating the added benefit of our loss term on the
validation perplexity of the LM for moderate values of $\omega$.

  2. 2.

Training performance does not decrease for moderate values of $\omega$,
showing that the extra term does not destroy training and generalization of
the standard LM.

  3. 3.

There is no difference in terms of validation perplexity between the standard
LM and our version before applying $k$NN, but there is a significant
difference after applying $k$NN. Our approach likely finds a different local
minimum for the language model that is better suited for $k$NN.

|  $k$NN-LM $(\omega\\!=\\!0)$ |  $k$NN-LM $(\omega\\!=\\!1)$ |  $k$NN-LM $(\omega\\!=\\!10)$
---|---|---|---
Train Ppl. LM | 11.31 | 11.24 | 11.07
Valid Ppl. LM | 18.00 | 17.95 | 17.71
Valid Ppl. $k$NN-LM | 16.09 | 15.89 | 17.46

Table 2: Experiments on WIKI-103

###  4.3 Experiments on WIKI-103

We illustrate our findings on the more challenging WIKI-103 dataset and
demonstrate that our L2 fix significantly improves the performance of the LM.
In the Table 2, we illustrate that when changing the regularization strength
we again see a significant gain in performance when adding our regularization
during training of the LM. Due to the computational costs when training these
models, we resort to the same hyperparameters as in the WIKI-2 dataset and
hence present fair comparisons of the different variants of the model.

Note that again, we see significant improvements in terms of validation
perplexity when using the $k$NN-LM scheme by simply adding a L2 regularization
when training the language model.

On another note, when taking a closer look at the validation perplexity before
applying $k$NN, we note that $\omega=10$ seems to have lowest validation
perplexity. This better generalization phenomenon is interesting and has
recently been noted in the machine vision community in the context of
investigating the regularization effects of batch normalization in
classification settings (Dauphin and Cubuk, 2021). This also relates to the
findings in (Merity et al., 2017), who used L2 regularization in LSTMs. In
this paper, we have found initial indications that the L2 regularization on
the activations might also be useful for Transformer models.

Finally, we believe that these two standard benchmark datasets in language
modelling are sufficient evidence to demonstrate the merit our of findings.

###  4.4 Further investigations into the representations and possible
explanations

We first examine the effect that target word frequency has on the loss. Figure
2 shows a histogram of word frequency, where the color represents the
respective losses each word incurred.

Figure 2: Frequency/Loss histograms. The x-axis denotes the frequency of the
word with high-frequency words to the left. The y-axis denotes the number of
words with $x$ frequency and the colors of each bar represents the loss
accumulated. (LEFT) Standard LM after $k$NN, (RIGHT) Our L2 regularized LM
after $k$NN.

Note that firstly, there is little difference in loss for the less frequent
words on the right side of the graph. If we shift our attention to the more
frequent words however, we see a different picture. Looking at our L2
regularized model, we note that for the most frequent words, our model seems
to incur lower loss (see the brighter colors bars at the peak of the
histograms) compared to the standard LM with $k$NN. This observation suggests
that the main differences in terms of representations come from the frequent
words rather than rare ones.

Secondly, knowing that the main differences are within the words that are most
frequent, we investigated these representations in more detail. In particular,
we analyzed the most frequent words and divided the data into "high scores"
meaning they contributed a lot to the loss (bad predictions) and "low scores"
meaning they are good predictions i.e. they contributed a little to the loss.

Figure 3: $x$-axis denotes the Loglikelihood under the Gaussian Mixture (GM).
$y$-axis denotes the normalized histogram. (TOP) Standard training of then LM
(BOTTOM) using a L2 regularization to train the LM.

We employed a simple mixture of Gaussians model (GMM) ($m=10$) and used the
log-likelihood as an indicator for how well the data are clustered. GMMs
allows us to put probability mass on each of the representations and given
that we are using a mixture of gaussians, we inherently captures clusters.
GMMs can be thought of as "putting multiple gaussian distributions on the
data". Intuitively, this means that if the likelihood of the GMM is high, the
representations can be easily captured using a mixture of gaussians, which is
indicative of being more clustered i.e. close to one of the gaussian mixture
means.

In Figure 3 we compare the distributions of the loglikelihoods for the
representations that have been trained using the standard LM and our modified
L2 regularization. In particular, for each representations we obtain the
corresponding likelihood from the GMM ($x$-axis on Figure 3). As mentioned
before, we split the words into "high scores" and "low scores" and plot their
histograms in blue and orange respectively. Fig. 3 demonstrates that when
using a L2 regularization (right), the representations that contribute less to
the loss (shown in orange) are better fitted by the GMM than when training the
model in the standard way. This can be seen by the increase in likelihood,
i.e. shift to the right on the $x$-axis. What is interesting is that the
difference between the likelihoods of the "high scores" and "low scores"
varies much more dramatically in the L2 regularized case. Hence, one of our
hypotheses is that $k$NN-LM improves the classification accuracy mostly for
the non-frequent words (Khandelwal et al., 2020), whereas our proposed method
with L2 regularization, in addition, also improves the classification accuracy
of the frequent words by clustering them closer together and hence allowing
$k$NN-LM to work better.

##  5 Conclusion

In conclusion, we propose a useful training mechanism that is inspired by the
fact that the post-hoc application of $k$NN seems to significantly improve the
performance of standard LMs. We have found that training a LM with L2
regularization at the final layer, i.e. layer which is used for the post-hoc
$k$NN search, improves validation performance. We have also found initial
indications that the L2 regularization mostly improves performance for the
most frequent, lower-loss words.

These findings motivate exploring L2 regularization in different Transformer
architectures and LM tasks. Furthermore, using training data at inference time
is an interesting direction for improving model performance, and we hope that
our approach encourages future work in this area.

## References

  * Baevski and Auli (2018) Alexei Baevski and Michael Auli. 2018.  Adaptive input representations for neural language modeling.  _arXiv preprint arXiv:1809.10853_.
  * Dauphin and Cubuk (2021) Yann N Dauphin and Ekin D Cubuk. 2021.  Deconstructing the regularization of batch-norm.  In _International Conference on Learning Representations (ICLR)_.
  * He et al. (2020) Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. 2020.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9729–9738.
  * Joulin et al. (2017) Armand Joulin, Moustapha Cissé, David Grangier, Hervé Jégou, et al. 2017\.  Efficient softmax approximation for gpus.  In _International Conference on Machine Learning_ , pages 1302–1310. PMLR.
  * Khandelwal et al. (2020) Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2020\.  Generalization through Memorization: Nearest Neighbor Language Models.  In _International Conference on Learning Representations (ICLR)_.
  * Merity et al. (2017) Stephen Merity, Nitish Shirish Keskar, and Richard Socher. 2017.  Regularizing and optimizing lstm language models.  _arXiv preprint arXiv:1708.02182_.
  * Press and Wolf (2016) Ofir Press and Lior Wolf. 2016.  Using the output embedding to improve language models.  _arXiv preprint arXiv:1608.05859_.