Multi-Task Structural Learning using Local Task Similarity induced Neuron Creation and Removal

We consider the problem of structural learning where the MTN architecture and its parameters are learned simultaneously. Given the set of $\mathcal{T}$ tasks that each has its own single network with $L$ layers, our algorithm results in a single MTN capable of inferring all the $\mathcal{T}$ tasks accurately without the need for retraining. First, we establish the terminologies that are used in the rest of the paper. A node is a layer that connects one branch to another branch (or to a node), and a branch is a sequence of layers that follow a node. Initially, the first layer of each single task network is the task node, while the rest of the task network excluding the task head is called the task branch. Similarly, a group of tasks will have a group node and a group branch. A task node is of particular significance to our algorithm, as tasks can only be fused at the task node. Also, only task nodes that are connected to the same group branch or group node can be fused. At the start of the training, all task nodes are connected to the input image and can be fused. $T$ and $G$ denote a task and a group, respectively. Additionally, $F$ and $\mathcal{F}$ denote the output features of the task node and the group node, respectively. Figure 1 provides the overall schematic of our approach, where the leftmost column illustrates the initial state of our setup using the terminologies defined so far. In the following sections, we discuss the different components of our approach.

As is evident from Figure 1, we begin training from single-task networks. Since the encoder of each task is initialized with ImageNet weights, there exists a correspondence between task nodes initially. During training, task nodes would learn concepts that minimize particular task loss independent of other tasks. This independence likely breaks any correspondence between parameters mapped one-to-one between any two task nodes. This behavior comes from the permutation invariance of neural networks that leads to no guarantee on the order in which concepts are learned (Wang et al., 2020a). Therefore, MTSL aligns the concepts learned by the task nodes and locally increases their similarity using Centered Kernel Alignment (CKA) (Kornblith et al., 2019).

CKA is used to measure similarity between two feature representations and has been shown to provide meaningful similarity scores. During training, we introduce a CKA-based regularization term between task nodes branching from the same group node/branch (or the input). This regularization term, as shown in the alignment part of Figure 1, is included between all pairs of task node features (indicated by bi-directional arrows) and enforces the task representations to align by serving as an alignment constraint. We use the unbiased CKA estimator (Nguyen et al., 2021) to facilitate reliable estimates of CKA with small batch sizes used during training.

The overall loss that is used for training in the task learning phase of our algorithm is provided in Equation 1, where the first term ($\mathcal{L}_{MTL}$) represents the multi-task loss which is a weighted sum of all individual task losses. The second term represents the CKA regularization term ($\mathcal{L}_{CKA}$) that is included with a balancing factor $\lambda$ and a negative sign to maximize alignment between tasks.

The overall loss $\mathcal{L}$ used during the task learning phase (discussed in Section 3.2) leads tasks to learn similar features while also minimizing the concerned task loss. After the task learning phase, MTSL begins the structural learning phase to first leverage neuron creation. In the brain, local neuronal activity can affect the structure of the neural circuitry (Luhmann et al., 2016) and play a role in learning experiences (Kudithipudi et al., 2022). Taking cues from these notions of locality, we use CKA to gauge the similarity between task node features that represent the local activity of task neurons. These local task similarities are used to induce the creation of group nodes.

First, CKA between all pairs of task node features is calculated after which all possible groups of task nodes are listed. From these groups, a set of groups that maximize the total similarity is chosen and the groups that satisfy a minimum similarity of $\gamma$ induce the creation of a group neuron. For instance, in Figure 1, we see that the groups picked are $[T_{1},T_{2}]$ and $T_{3}$ assuming that the total number of tasks $\mathcal{T}$ is three. More details regarding the grouping algorithm are provided in the Appendix.

After the task nodes have been grouped based on local task similarity, for each group, a group node is created. The learned knowledge in the task neurons is used to initialize the created group node using a two-step process. In the first step, the weights of the group node are obtained by averaging the parameters of the concerned task nodes. This averaging is justified by the alignment constraint used in the task learning phase that ensures that the corresponding parameters learn similar concepts. Figure 1 depicts the averaging initialization using a plus symbol. In the second step, MTSL distills the information learned by multiple task nodes into the group node using an attention-based feature amalgamation method (Ye et al., 2019) referred to as ATT. Figure 1 depicts this amalgamation process by using arrows from task node features $F_{1}$ and $F_{2}$ to group node feature $\mathcal{F}$.

The knowledge amalgamation objective $\mathcal{L}_{KA}$ is provided in Equation 2 assuming that there are $N$ tasks grouped together. $ATT^{net}$ denotes the attention network consisting of two linear layers with an intermediate ReLU activation and a final sigmoid activation. $ATT_{i}^{net}(\mathcal{F})$ provides a 1$\times$$C$ dimensional attention vector which acts as a weight for the different channels of $\mathcal{F}$ allowing the selective distillation of task features into the group feature.

Starting from a dense set of neurons as in the initial state of MTSL provides the opportunity to leverage a rich information flow originating from diverse task information. Using neuron removal, MTSL moves towards a sparser architecture by removing task nodes that learn similar representations. These locally similar task nodes become redundant once they transfer their knowledge to the group node. The task branch is then disconnected from these redundant task nodes and connected to the group node. As defined in Section 3.1, the neurons in the task branch that now connect to the group node become task nodes. These changes are evident in the depicted next state in Figure 1.

Algorithm 1 presents the different phases involved in the MTSL algorithm, namely the task learning phase, the structural learning phase, and a fine-tuning phase. The task learning phase and the structural learning phase occur alternatively for $n$ number of times, followed by a final fine-tuning phase. In the task learning phase, the entire network is trained to minimize multi-task loss and maximize similarity among task nodes as described in Section 3.2. The structural learning phase involves neuron creation and neuron removal as discussed in Section 3.3 and in Section 3.4, respectively. $E_{t}$ determines the number of epochs for which each subsequent task learning phase is executed. Similarly, $E_{s}$ determines the epochs for ATT-based knowledge transfer. Considering a total training budget of $E$ epochs, the task learning phase is executed up to $E-f$ epochs where $f$ is the minimum epochs allocated for the fine-tuning phase during which the task nodes are no longer forced to align.

We compare MTSL with existing multi-task architectures and methods to learn multi-task architecture branching. The implementations of Cross-stitch (Misra et al., 2016) and MTAN (Liu et al., 2019) have been taken from the repository of (Vandenhende et al., 2021) while the implementation of Learning-to-Branch (LTB; (Guo et al., 2020)) has been taken from LibMTL (Lin & Zhang, 2022). We update LTB with the resource loss proposed by Branched Multi-Task Architecture Search BMTAS (Bruggemann et al., 2020) from their open source repository to obtain the implementation of BMTAS. LTB-R and BMTAS-R refer to the results obtained by retraining the converged architecture from scratch. For all methods we use the same hyperparameter settings and use ResNet18 backbone with DeepLab head.

In Table 1, we observe that Cross-stitch obtains the best improvement (highlighted in bold) over Single Task Networks (STNs). Cross-stitch retains all separate task networks and learns task-specific adapter parameters. These adapter parameters enable the selective transfer of information from other tasks. Additionally, since there is no explicit parameter sharing, task interference can be mitigated, leading to increased performance. However, Cross-stitch loses the inference time and memory advantage of using shared parameters. When comparing MTSL with methods that learn branching structures (last group), MTSL outperforms both LTB and BMTAS by a considerable margin in both datasets. For instance, in $\Delta_{MTL}^{\mathcal{S}\mathcal{D}}$, MTSL outperforms LTB by 2.72% and 4.72% and BMTAS by 3.61% and 3.94% in Cityscapes and NYUv2 datasets. Also, with the aid of CKA regularization, MTSL leads to a more inference efficient architecture (visualized in Section 4.4) as evidenced by GMac. LTB-R, BMTAS-R and MTSL-R are versions of LTB, BMTAS and MTSL, respectively, where the final converged architecture is reinitialized and retrained. With additional training, LTB-R and BMTAS-R can obtain a marginal gain over MTSL. However, they deviate from structural learning, as the original trained weights are no longer relevant and are discarded. MTSL-R is more efficient in terms of inference and shows a performance comparable to that of LTB-R and BMTAS-R. Note that MTSL is designed for structural learning and shows clear performance and inference efficiency improvements over LTB and BMTAS.

The methods discussed so far only consider learning branching structure in the encoder. For the upcoming experiments, we extend MTSL to UniNet (Gurulingan et al., 2021), an encoder-decoder architecture with a total of ten locations, six in the encoder and four in the decoder, where task nodes can participate in neuron creation and removal. The encoder is based on ResNet18, and the decoder has ResNet blocks.

MTSL algorithm begins from STNs and could potentially end in a network with all encoder and decoder layers shared between all tasks (referred to as One-Net). These two possible networks are considered as the baselines to evaluate our approach. Table 2 shows the generalization performance of the baselines and MTSL. MTSL achieves better generalization than One-Net in most cases, while being close to One-Net in inference efficiency, as seen in the GMac and parameter count columns. In NYUv2, MTSL even rivals the performance of STN in both $\mathcal{S}$ and $\mathcal{D}$ tasks. We hypothesize that the improved generalization can be attributed to brain-inspired aspects of the MTSL algorithm, such as local task similarity and the change of dense to sparse architecture.

The training time of MTSL is higher than that of One-Net because of the training epochs required for knowledge amalgamation. However, this only adds 34 additional epochs in training where only the network parts up until the task nodes is involved in the computation. For a fair comparison, we train One-Net for 34 more epochs and get the One-Net-L baseline. We note that MTSL also improves generalization over One-Net-L. Overall, we see that with only a fractional increase in training costs, MTSL algorithm provides a learned network with better generalization.

MTSL provides an improved generalization, likely due to its motivation originating from abstract notions of the brain. Given that the brain is effective in discerning noise from semantics, a natural question to ask is whether MTSL can lead to improvements in robustness to natural corruptions. Table 3 shows the robustness of the baselines and MTSL to various natural corruptions (Hendrycks & Dietterich, 2019) categorized into four types, namely noise, blur, weather, and digital. Under each corruption category, the average across five different severity levels is taken. We see that the MTSL network exhibits better robustness compared to One-Net in most cases, especially in the weather and digital category of the Cityscapes dataset. These results further demonstrate that MTSL presents a compelling case for the utility of drawing inspiration from the brain.

The assumptions about the data used for training play a pivotal role in determining the learned representations. In the brain, exposure to the nature of experiences determines the way neural circuitry develops (Kudithipudi et al., 2022). Thus, the resultant architecture obtained with MTSL on two datasets would likely differ. Figure 2 visualizes the One-Net architecture and the learned architectures in Cityscapes and NYUv2. Evidently, after the first branch from layer six, the task groups and the branching structure emerging on the two datasets are different. Semantic segmentation ($\mathcal{S}$) and edge detection ($\mathcal{E}$) grouped together in both datasets follow intuition, as edge detection requires predictions of semantic edges. However, the emergence of certain groups of tasks such as depth and autoencoder in NYUv2 is counterintuitive. In addition to task relationships, these results show that local task similarity and inherent biases in the dataset can also impact the optimal architecture. In Section E, we analyze the sensitivity of the converged architecture to other factors.

In Table 1, we provide results for state-of-the-art LTB and BMTAS methods. Here, we visualize the converged architectures of LTB and BMTAS along with those of MTSL and provide inferences. Figure 3 illustrates the converged architectures on the Cityscapes dataset (and Figure 4 in Appendix on NYUv2 dataset). We see that LTB and BMTAS generally converged to a larger architecture compared to MTSL. MTSL is able to learn a more efficient architecture, likely due to the use of the CKA regularization term, which forces tasks to learn similar representations, leading to more sharing between tasks.

To determine the effectiveness of the different components involved in MTSL, we perform systematic evaluations and provide the results in Table 4. When alignment is not used, the task representations diverge largely around layers 3 and 4. As a result, the task nodes no longer become sufficiently similar to be fused. Therefore, the resultant architectures in the first three rows perform similarly and have inference requirements similar to STN. On the other hand, in the last three rows, with the help of alignment, the resultant architectures undergo neuron creation and removal in more layers and approach close to the inference efficiency of One-Net. Of these, the last row constitutes all the components used in MTSL and results in the best performance. Note that although we use CKA for alignment and knowledge amalgamation, various other methods can also be used.