Cortex learns to share and specialize modules inside language models
CORTEX: Learning to Share and Specialize in Dense Language Models
Artificial Intelligence
Summary
Large language models cover many topics but need to balance shared knowledge with specific skills for each area. The authors propose Cortex, a new approach that splits the model’s parameters into groups and learns which groups should focus on which topics by analyzing training updates from different domains. This method helps the model form clear modules specialized for different knowledge areas, improving performance on certain synthetic tests while staying strong on real-world data. Cortex learns these modules internally during training instead of requiring extra components or post-training analysis.
What this means in practice
- •For language model engineers: Train language models that automatically form domain-specialized modules to improve accuracy and interpretability.
- •For ai system integrators: Use modular language models with learned specialization to better tailor AI outputs to specific knowledge areas or synthetic domains.
Authors
Chuiyang Meng, Ming Tang, Vincent W. S. Wong
Abstract
Large language models are trained on heterogeneous data mixtures, where different knowledge domains require both shared knowledge and specialization. Existing modular approaches typically impose explicit components or discover modules through interpretability analysis after training. In this work, we propose CORTEX, a learning dynamics-inspired framework that learns internal modularization within dense language models. CORTEX partitions trainable matrices into parameter groups and learns module assignments from domain-conditioned gradient and cross-domain gradient similarity. We introduce the selective lesion score and module-domain mutual information to characterize the target-domain lesion effects and alignment, and analyze how module assignment affects the trade-off between assignment bias and update magnitude. Experiments with 160M, Qwen3-8B, and Qwen3-32B backbone models show that CORTEX achieves the highest synthetic-domain exact match and largest average perplexity reduction, while remaining competitive on real-domain evaluations and forming identifiable modules.