Title: Dissecting Attention and MLP Roles: A Study of Domain Specialization in Large Language Models
Primary area: interpretability and explainable AI
Keywords: Mechanistic Interpretability, Large Language Models, Domain control, layer analysis, activation patching

Abstract: Large language models (LLMs) perform well across diverse domains such as programming, medicine, and law, yet it remains unclear how domain information is represented and distributed within their internal mechanisms. A key open question is the $\textit{division of labor}$ between the Transformer's core components: self-attention and MLP layers. We address this question through a mechanistic study that dissects their roles by integrating three complementary analyses: representation separability via probes, parameter change under adaptation, and causal effects from activation swaps. We propose a clear division of labor: attention layers route domain identity, while MLP layers implement domain-specific computation. Causal interventions strongly support our claims. For instance, swapping attention activations at specific mid-depth layers (e.g., for Python $\leftrightarrow$ C++) reliably shifts the next-token distribution, whereas layers with low domain separability have a negligible effect. In contrast, while finetuning, MLP layers exhibit relatively larger weight changes, consistent with domain-specific knowledge being stored there. This pattern holds consistently across four models and six domains. In a supplementary experiment, we demonstrate that selecting a few components highlighted by our study can accelerate domain adaptation, indicating the potential for more focused fine-tuning.

TL;DR: In LLMs, attention layers act as domain "routers" while MLP layers store and perform the domain-specific "computation" in high-level domain abstraction