Transformers as Multi-task Learners: Decoupling Features in Hidden Markov Models

2025/06/02 by Yifan Hao, Hao, Yifan, Chenlu Ye +5 · 1 voice
Computer Science · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Natural Language Processing Techniques #Speech Recognition and Synthesis #cs.AI #cs.LG

paper · pdf · doi:10.48550/arxiv.2506.01919

openalex publication_date 2025/06/02 · arxiv published 2025/06/02 · arxiv updated 2025/06/02 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Transformer based models have shown remarkable capabilities in sequence learning across a wide range of tasks, often performing well on specific task by leveraging input-output examples. Despite their empirical success, a comprehensive theoretical understanding of this phenomenon remains limited. In this work, we investigate the layerwise behavior of Transformers to uncover the mechanisms underlying their multi-task generalization ability. Taking explorations on a typical sequence model, i.e, Hidden Markov Models, which are fundamental to many language tasks, we observe that: first, lower layers of Transformers focus on extracting feature representations, primarily influenced by neighboring tokens; second, on the upper layers, features become decoupled, exhibiting a high degree of time disentanglement. Building on these empirical insights, we provide theoretical analysis for the expressiveness power of Transformers. Our explicit constructions align closely with empirical observations, providing theoretical support for the Transformer's effectiveness and efficiency on sequence learning across diverse tasks.

Citations

• Transformers and Their Roles as Time Series Foundation Models
• The Llama 3 Herd of Models
• Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process
• How Transformers Learn Causal Structure with Gradient Descent
• How do Transformers perform In-Context Autoregressive Learning?
• Provably Efficient High-Dimensional Bandit Learning with Batched Feedbacks
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations
• Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining
• One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
• Provably Efficient Representation Learning with Tractable Planning in Low-Rank POMDP
• Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection
• A Latent Space Theory for Emergent Abilities in Large Language Models
• GPT-4 Technical Report
• LLaMA: Open and Efficient Foundation Language Models
• Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers
• Transformers learn in-context by gradient descent
• What learning algorithm is in-context learning? Investigations with linear models
• GEC: A Unified Framework for Interactive Decision Making in MDP, POMDP, and Beyond
• Transformers Learn Shortcuts to Automata
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
• PAC Reinforcement Learning for Predictive State Representations
• Provably Efficient Reinforcement Learning in Partially Observable Dynamical Systems
• Emergent Abilities of Large Language Models
• Embed to Control Partially Observed Systems: Representation Learning with Provable Sample Efficiency
• When Is Partially Observable Reinforcement Learning Not Scary?
• Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?
• Low-Rank Constraints for Fast Inference in Structured Models
• An Explanation of In-context Learning as Implicit Bayesian Inference
• Representation Learning for Online and Offline RL in Low-rank MDPs
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs
• Language Models are Few-Shot Learners
• Learning higher-order sequential structure with cloned HMMs
• Reduced-Rank Hidden Markov Models
• A tutorial on hidden Markov models and selected applications in speech recognition
• Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning

Discussions

• Transformers as Multi-Task Learners: Decoupling Features in Hidden Markov Models [hn, 2 points, 1 comments]

Related

• Systematic Generalization and Emergent Structures in Transformers Trained on Structured Tasks
• Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformer
• Transformers on Markov Data: Constant Depth Suffices
• Can Transformers Learn Sequential Function Classes In Context?
• On Limitation of Transformer for Learning HMMs