[Paper Review] The Stem Cell Hypothesis: Dilemma behind Multi-Task Learning with Transformer Encoders

Han He, Jinho D. Choi|arXiv (Cornell University)|Sep 14, 2021

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper proposes the Stem Cell Hypothesis, explaining that shared attention heads in Transformer encoders are inherently skilled at multiple NLP tasks but become impaired when jointly fine-tuned due to conflicting optimization signals. Through extensive pruning and parameter-free probing across five core NLP tasks (POS, NER, DEP, CON, SRL), the authors show that multi-task learning (MTL) underperforms single-task learning because these 'stem cell' heads lose task-specific expertise, leading to degraded performance across all tasks.

ABSTRACT

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.

Motivation & Objective

• To investigate why multi-task learning (MTL) with shared Transformer encoders underperforms single-task learning on diverse core NLP tasks.
• To identify whether attention heads are shared across tasks and how their fine-tuning dynamics affect performance.
• To explain the underlying cause of MTL failure in terms of attention head competition and interference.
• To validate the existence of 'stem cell' attention heads—naturally talented for multiple tasks but impaired under joint training.
• To develop and apply parameter-free probing methods to analyze attention head behavior across single-task and multi-task settings.

Proposed method

• Conducted multi-task learning (MTL) on five core NLP tasks: POS, NER, DEP, CON, and SRL using shared BERT-based encoders.
• Performed extensive structured pruning to identify essential attention heads for each task, revealing high overlap across tasks.
• Proposed the Stem Cell Hypothesis: certain attention heads are inherently multi-task capable but degrade under joint fine-tuning due to conflicting gradients.
• Designed novel parameter-free probes to assess attention head performance and transformation across tasks without additional parameters.
• Used dynamic and static pruning to measure head utilization and performance trade-offs in single-task and multi-task settings.
• Visualized attention head utilization across layers and tasks using RGB-encoded run aggregations to track head allocation patterns.

Experimental results

Research questions

• RQ1 Why does multi-task learning with shared Transformer encoders lead to lower performance than single-task learning on diverse NLP tasks?
• RQ2 To what extent are attention heads shared across different NLP tasks, and how does this sharing affect model performance?
• RQ3 Do certain attention heads possess inherent multi-task capability, and if so, why do they fail under joint training?
• RQ4 How do probing results and head utilization patterns differ between single-task and multi-task fine-tuning?
• RQ5 Can parameter-free probing methods effectively reveal the behavior and degradation of attention heads in MTL?

Key findings

• Multi-task learning on all five tasks (MTL-5) consistently underperforms single-task learning, with the worst performance observed on SRL (83.56% vs. 83.52% in single-task).
• All five tasks rely on nearly identical sets of attention heads, leading to competition and interference during joint training.
• Parameter-free probing shows that heads not fine-tuned for any task still achieve high performance on linguistic structures, confirming the existence of naturally talented 'stem cell' heads.
• After single-task learning, probing accuracy improves in line with task performance, indicating that stem cells evolve into task-specific experts.
• In multi-task learning, both probing and task performance drop significantly, supporting the hypothesis that conflicting signals degrade head expertise.
• Head utilization analysis reveals that MTL-DP models use a higher percentage of heads (e.g., 53.47% in BERT) than single-task models, indicating increased competition and inefficient allocation.