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

2025-06-02

AI Summary

This work investigates the hierarchical mechanisms and multitask generalization origins of Transformers in Hidden Markov Model (HMM)-based sequence tasks. To address this, we combine theoretical analysis—proving expressive capacity for HMM modeling—with empirical diagnostics, including layer-wise activation visualization and attention pattern analysis. Our study is the first to systematically characterize functional differentiation across Transformer layers: lower layers specialize in local token feature extraction, while upper layers perform feature disentanglement and temporal decoupling. Building on this insight, we propose a novel HMM-inspired feature disentanglement paradigm, establishing a theoretical foundation for interpretable Transformer modeling. Both theoretical constructions and experimental results demonstrate strong consistency, confirming that this hierarchical decomposition underpins efficient cross-task sequence modeling. Our findings significantly advance the understanding of the fundamental principles governing Transformer-based sequential learning.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Interpretability, Analysis, and Evaluation of NLP Models
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

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.

Problem

• Understanding Transformers' multi-task generalization mechanisms
• Analyzing layerwise feature decoupling in Hidden Markov Models
• Providing theoretical support for Transformers' sequence learning efficiency

Innovation

• Transformers decouple features in upper layers
• Lower layers extract neighboring token features
• Theoretical analysis supports Transformer effectiveness