Source-linked AI summary

The Stem Cell Hypothesis: Dilemma behind Multi-Task Learning with Transformer Encoders

Han He, Jinho D. Choi

arXiv:2109.06939v1cs.CLcs.AI

TL;DR

Whether transformer-encoder MTL works as well across distinct core NLP tasks as it does across closely related tasks remains an open question, despite its potential computational efficiency. The paper evaluates five tasks, analyzes shared attention-head use, and finds that MTL-5 underperforms single-task models while probes indicate head expertise can be disrupted by joint training.

  • Problem

    MTL with transformer encoders has been studied less for core NLP tasks with distinct decoders than for closely related tasks, despite its potential to reduce computation by sharing an encoder.

    Source: §1 · p. 1
  • Method

    The authors analyze attention-head use across five core NLP tasks with pruning and parameter-free probes of individual heads during single-task and joint training.

    Source: §1 · p. 1
  • Results

    Across five core tasks, MTL-5 underperforms single-task models; probes show shared task-relevant heads and often lower probing and task performance under MTL.

    Source: §1 · p. 1
  • Takeaways & Limitations

    The proposed Stem Cell Hypothesis characterizes naturally talented attention heads as able to develop into task experts but not all-rounders.

    Source: §6 · p. 9

    The SRL analysis uses only predicate head words, and the authors note that not all words in an argument span are necessarily important to its predicate.

    Source: §5.1 · p. 6

Abstract

from arXiv · show

Multi-task learning with transformer encoders (MTL) has emerged as a powerful technique to improve performance on closely-related tasks for both accuracy and efficiency while a question still remains whether or not it would perform as well on tasks that are distinct in nature. We first present MTL results on five NLP tasks, POS, NER, DEP, CON, and SRL, and depict its deficiency over single-task learning. We then conduct an extensive pruning analysis to show that a certain set of attention heads get claimed by most tasks during MTL, who interfere with one another to fine-tune those heads for their own objectives. Based on this finding, we propose the Stem Cell Hypothesis to reveal the existence of attention heads naturally talented for many tasks that cannot be jointly trained to create adequate embeddings for all of those tasks. Finally, we design novel parameter-free probes to justify our hypothesis and demonstrate how attention heads are transformed across the five tasks during MTL through label analysis.

1 Introduction

Although sharing a transformer encoder can reduce computation, MTL on core NLP tasks with distinct decoders is less studied than MTL on closely related tasks. Across five tasks, joint training underperforms single-task models; the paper investigates shared-head interference and uses parameter-free probes to examine how heads change during training.

1 Introduction

  • The study addresses a gap in evidence on MTL across core NLP tasks with distinct decoders, beyond prior work on closely related tasks.Sharing one transformer encoder is also presented as a way to reduce computation by running it once to generate embeddings for all tasks.Source: §1 · p. 1
  • MTL-5 is outperformed by single-task models across all five core tasks, and performs worst among the tested task combinations.The tasks are POS, NER, DEP, CON, and SRL.Source: §1 · p. 1
  • All five tasks rely on almost the same attention heads, which compete during MTL and blur features extracted for individual tasks.The authors propose the Stem Cell Hypothesis: these naturally talented heads can be fine-tuned for distinct tasks, but not jointly for all of them.Source: §1 · p. 1
  • Parameter-free probes show that MTL often lowers both probing and task performance, consistent with heads losing expertise under potentially conflicting teaching signals.The probes track how individual attention heads are updated in individual versus joint training.Source: §1 · p. 1
  • The authors describe their work as a comprehensive head-level analysis of core-task MTL, with parameter-free probes and supporting experiments on RoBERTa, ELECTRA, and DeBERTa.Most analysis is based on BERT.Source: §1 · p. 2

2 Related Work

Prior work studies MTL for related tasks and analyzes linguistic information in transformer representations, including with parameter-free probes. The impact of fine-tuning on transformer encoders remains unclear, and this paper focuses its probes on explaining MTL’s impact.

2 Related Work

  • MTL research has studied pre-trained transformers in NLP, often designing architectures for message passing; this work instead uses a conventional architecture with batch sampling and loss balancing.The cited MTL studies include Wang et al. (2018), Clark et al. (2019b), Liu et al. (2019b), Kondratyuk and Straka (2019), Chen et al. (2020), and Lin et al. (2020).Source: §2 · p. 2
  • Prior BERT studies find linguistic structures in its representations using supervised and parameter-free probes, providing context for this paper’s probing methods.The passage discusses part-of-speech, syntactic chunks, semantic roles, and dependency structure.Source: §2 · p. 2
  • The impact of fine-tuning on transformer encoders remains unclear, while this paper’s probing methods focus on the impact of MTL.Zhao and Bethard (2020) compare probing performance before and after fine-tuning; this paper’s approach aligns with those unsupervised probes.Source: §2 · p. 2

3 Multi-Task Learning

The study shares a transformer encoder across five core NLP tasks while retaining task-specific decoders, and balances task sampling and losses during training. MTL-5 underperforms single-task models on every task, and none of the tested task pairs yields a win-win result.

3 Multi-Task Learning

  • The joint model shares one encoder across tasks, with a distinct decoder for each task making predictions from the same hidden states.The stated goal is to outperform single-task counterparts while being faster and more memory efficient.Source: §3 · p. 2

3.1 Shared Encoder

  • BERT is the shared encoder in the main experiments, with each token represented by the average of the last-layer hidden states of its subtokens.Random subtokens are replaced with [MASK] during training for word dropout.Source: §3.1 · p. 2

3.2 Task-Specific Decoders

  • The five tasks are POS, NER, DEP, CON, and SRL, each using an adopted state-of-the-art decoder except POS.The decoders are simplified to build an efficient model.Source: §3.2 · p. 2

3.3 Data and Loss Balancing

  • Task batches are shuffled and sampled in proportion to dataset size^0.8, while running-average loss normalization balances the tasks.The running average uses each task’s five most recent losses.Source: §3.3 · p. 3

3.4 MTL Experiments

  • The reported evaluation averages three runs and uses task-specific metrics on OntoNotes 5.Metrics are POS accuracy, NER span-level labeled F1, DEP labeled attachment score, CON constituent-level labeled F1, and SRL micro-averaged F1.Source: §3.4 · p. 3
  • MTL-5 is outperformed by its single-task counterparts on all five tasks, and none of the ten pairwise combinations yields a win-win result.NER generally improves with MTL while the other tasks degrade; SRL benefits from CON, but CON does not benefit from SRL.Source: §3.4 · p. 3, Table 1 · p. 3

4 Pruning Analysis

The pruning analysis identifies task-essential attention heads and finds substantial overlap in the heads used across tasks and between single-task and multi-task models. This overlap motivates the Stem Cell Hypothesis that some commonly used heads cannot be jointly trained to create adequate embeddings for dissimilar tasks.

4.1 Pruning based on L0-Regularization

  • L0-regularization encourages BERT to use fewer attention heads by learning a binary gate for each head and optimizing it with task-specific or balanced multi-task loss.Because binary gates are nondifferentiable, the method relaxes them with a continuous distribution and uses the Hard Concrete Distribution.Source: §4.1 · p. 4

4.2 Pruning Strategies

  • Dynamic pruning outperforms static pruning on accuracy across all five tasks and prunes more heads except for CON.Static pruning fine-tunes the encoder before freezing the decoder for pruning; dynamic pruning fine-tunes and prunes simultaneously while the decoder adapts to the pruned encoder.Source: §4.3 · p. 4, §4.2 · p. 4

4.3 Pruning Experiments

  • STL-DP performs as well as or slightly better than unpruned STL for POS, NER, and SRL while using approximately 50% fewer heads.The passage attributes the performance to L0-regularization.Source: §4.3 · p. 4
  • SRL requires more heads than DEP and CON, which require more than POS; NER requires more than CON, attributed to its need to capture more scattered world knowledge.The first ordering aligns with the paper's semantic-over-syntactic-over-lexical relation difficulty intuition.Source: §4.3 · p. 5
  • MTL-DP prunes 50% of the heads and has slightly higher accuracy than MTL-5, except for NER.The comparison suggests that about half of the heads are underutilized during MTL training.Source: §4.3 · p. 5

4.4 Pruning Visualization

  • Head utilization is similar across runs, especially for DEP, CON, and SRL, while POS utilization appears more random because high performance can use a small set of heads.The authors describe consistency across runs as a prerequisite for subsequent analyses.Source: §4.4 · p. 5
  • Across tasks, the pruned models mostly use heads from layers 5 to 8, which the authors associate with rich surface and semantic features.This contrasts with Jawahar et al. (2019) and Tenney et al. (2019), and aligns with Liu et al. (2019a) on POS and chunking performance using middle-layer heads.Source: §4.4 · pp. 5–6
  • STL-DP and MTL-DP show almost identical head-utilization patterns, motivating the hypothesis that commonly used heads cannot be jointly trained to create adequate embeddings for dissimilar tasks.The hypothesis proposes that individual heads have limited representation capacity and are suited only to specific tasks.Source: §4.4 · p. 6

5 Probing Analysis

The section uses parameter-free probes to examine individual attention heads before and after fine-tuning. The probes find naturally capable heads, sparse MTL gains across labels, and task-specific changes in head performance.

5.1 Probing Methods

  • Probes evaluate individual heads using attention links for DEP/SRL and attended-value representations to label POS/NER/CON spans.The attended-value representation is formed as H = AV, where attention weights A are applied to values V.Source: §5.1 · pp. 6–7
  • For CON, the probe uses only height-3 constituents; SRL analysis uses predicate-head words, limiting coverage of verb-particle constructions and argument-span words.Higher CON levels did not correlate well with model performance as spans became longer and noisier.Source: §5.1 · pp. 6–7

5.2 Probing Experiments

  • Across pretrained BERT, STL, pairwise MTL, and MTL-5, experiments select each label’s top-probing head from 144 heads and average results across three seeds.The 144 heads comprise 12 heads in each of 12 layers.Source: §5.2 · p. 7
  • MTL improves probing performance on only a few labels, while MTL-5 provides no benefit for most labels.Most heads perform better after fine-tuning, but joint training does not consistently improve them across tasks.Source: §5.2 · p. 7

5.2.1 Pluripotent Stem Cells

  • The 203 probing tasks exceed BERT-base’s 144 heads, and the best ROOT head reaches 96.25% unsupervised accuracy, with ROOT identified mostly in layers 2–3.The results suggest that some heads play multiple roles and that syntax can be encoded in early transformer layers.Source: §5.2.1 · pp. 7–8, Figure 2 · p. 7
  • SRL recognizes numbered arguments in layers 5–7 and five modifier roles above 80% accuracy in layers 8–10; no role peaks before layer 5.The authors infer that semantic roles take more effort to learn than syntactic dependencies.Source: §5.2.1 · p. 8, Figure 3 · p. 8

5.2.2 Stem Cells Specialization

  • POS MTL degrades most labels except WP$, which improves with NER, DEP, and SRL; pretrained BERT exceeds 90% on six NER types, and POS/DEP improve more entity types than other tasks.The POS and NER comparisons are shown in Figures 4 and 5.Source: §5.2.2 · p. 8, Figure 4 · p. 8, Figure 5 · p. 8
  • POS improves the most constituent types but causes the largest drop for NP; NER, DEP, and SRL show smaller, task-specific constituency gains.The paper suggests NP’s lexical complexity may contribute to its decline; without pruning, stem-cell knowledge can transfer to non-stem-cell heads that specialize.Source: §5.2.2 · p. 9, Figure 6 · p. 9

6 Conclusion

The conclusion presents the Stem Cell Hypothesis as an account of attention-head interference and reports its validation through parameter-free probes.

6 Conclusion

  • The hypothesis is that naturally talented heads can become task experts but cannot create adequate shared embeddings for dissimilar tasks; parameter-free probes reveal interfered representations.The study focuses on five core tasks and identifies broader task and language coverage as future work.Source: §6 · p. 9

A.1 Corpus Statistics

The appendix identifies the datasets whose statistics are reported in Tables 4 and 5.

A.1 Corpus Statistics

  • Table 4 reports dataset statistics for POS, DEP, CON, and SRL; Table 5 reports NER statistics.The tables label training, development, and evaluation splits as TRN, DEV, and TST.Source: §A.1 · p. 12, Table 4 · p. 12, Table 5 · p. 12

A.2 Hyper-Parameter Configuration

The appendix points to Table 6 for the models’ hyper-parameter settings.

A.2 Hyper-Parameter Configuration

A.3.1 Dependency Parsing

Probing comparisons show relation-specific changes and generally lower MTL probing scores; across three additional transformer encoders, MTL-5 underperforms single-task models while pruning retains comparable performance.

A.3.1 Dependency Parsing

  • MTL improves only selected dependency relations, including dobj and expl, while other tasks show different relation-level gains.NER favors noun modifiers; CON modifiers and complements; SRL subjects and clausal relations.Source: Figure 7 · p. 12, §A.3.1 · p. 12

A.3.2 Semantic Role Labeling

  • Compared with STL, POS and CON mainly improve ARG0-3 and selected modifiers, while NER and DEP improve referent arguments and modifiers.The paper suggests POS and CON reveal syntactic functions that coordinate attention to semantic roles; it attributes NER and DEP improvements possibly to their biaffine decoders.Source: §A.3.2 · p. 13

A.3.3 Detailed Probing Results

  • MTL probing scores are generally lower than STL scores, while pre-trained BERT probing scores are usually much lower than STL.The probing tables report mean results over three runs, with STL as the baseline.Source: §A.3.3 · p. 13, Table 7 · p. 14, Table 8 · p. 15

A.4 Results for Other Transformers

  • Across all transformer encoders, MTL-5 is outperformed by its single-task counterparts on every task, as shown in Table 12, raising a dilemma for transformer-based multi-task learning.Table 12 reports single-task, five-task MTL-5, and pairwise-task results.Source: §A.4 · p. 19, Table 12 · p. 20
  • Dynamic pruning removes roughly 50% of heads while retaining comparable performance across the additional transformer encoders.The base encoders had 12 layers and 144 heads; BERT hyperparameters were reused without model-specific tuning.Source: §A.4 · p. 19, Table 13 · p. 20
  • MTL-DP reuses a very similar set of heads to STL-DP across the three additional encoders.Head-utilization visualizations are provided for RoBERTa, ELECTRA, and DeBERTa.Source: §A.4 · p. 19, Figure 9 · p. 21, Figure 11 · p. 22
  • DEP probing is already very high on the additional encoders before dependency-treebank fine-tuning; stem-cell heads specialize in STL but lose specialities in MTL.The paper presents this pattern as supporting the Stem Cell Hypothesis.Source: §A.4 · p. 19, Figure 10 · p. 21, Figure 12 · p. 22
Loading 2109.06939v1…