CORTEX: Learning to Share and Specialize in Dense Language Models

Chuiyang Meng; Ming Tang; Vincent W. S. Wong

cs.AISep 28, 2026

CORTEX: Learning to Share and Specialize in Dense Language Models

Authors: Chuiyang Meng, Ming Tang, Vincent W. S. Wong

Organizations: Simon Fraser University · Southern University of Science and Technology · The University of British Columbia

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.

Figures & tables

Figure 1: Visualization of CORTEX modularization on the 160M backbone model trained on the synthetic benchmark. The heatmaps correspond to the shared and six knowledge domains, including (a) Shared, (b) COPY, (c) REVERSE, (d) ADD, (e) LOOKUP, (f) LOOKUP $\rightarrow$ ADD, and (g) REVERSE $\rightarrow$ LOOKUP, respectively. The columns denote Transformer blocks from 1 to 12, and the rows denote parameter groups within attention heads (Attn.) and the feed-forward network (FFN).

Figure 2: A schematic illustration of CORTEX, including (a) the relationship among knowledge domains, parameter groups, and specialized modules, and (b) the modularization process of CORTEX.

Backbone / Metric	Approach	Domain 1	Domain 2	Domain 3	Domain 4	Domain 5	Domain 6
		COPY	REVERSE	ADD	LOOKUP	LOOKUP $\rightarrow$ ADD	REVERSE $\rightarrow$ LOOKUP
160M EM (%)	Best Baseline	$98.67{\pm}0.21$ (MoM)	$97.84{\pm}0.27$ (MoM)	$96.21{\pm}0.33$ (MoM)	$96.48{\pm}0.32$ (UpIT)	$90.82{\pm}0.49$ (MoM)	$95.71{\pm}0.37$ (Self-MoE)
160M EM (%)	CORTEX	$\mathbf{99.42{\pm}0.18}$	$\mathbf{98.83{\pm}0.22}$	$\mathbf{97.62{\pm}0.29}$	$\mathbf{98.14{\pm}0.26}$	$\mathbf{94.76{\pm}0.41}$	$\mathbf{96.88{\pm}0.33}$
		General	Code	Math	Biomedical	Legal	Reasoning
Qwen3-8B NLL	Best Baseline	$1.6257{\pm}0.0108$ (UpIT)	$1.6578{\pm}0.0124$ (MoM)	$1.9058{\pm}0.0148$ (Self-MoE)	$1.7709{\pm}0.0134$ (UpIT)	$1.8171{\pm}0.0140$ (UpIT)	$\mathbf{1.7822{\pm}0.0128}$ (Self-MoE)
Qwen3-8B NLL	CORTEX	$\mathbf{1.6199{\pm}0.0097}$	$\mathbf{1.6096{\pm}0.0115}$	$\mathbf{1.8694{\pm}0.0137}$	$\mathbf{1.6778{\pm}0.0126}$	$\mathbf{1.7791{\pm}0.0131}$	$1.8258{\pm}0.0124$

Table 1: Comparison between CORTEX and the best baseline. We show exact match (EM, in %) for the synthetic benchmark and negative log likelihood (NLL) for the real-domain benchmarks.

Figure 3: Comparison of the PPL reduction over Dense on (a) the 160M model, (b) Qwen3-8B, and (c) Qwen3-32B. The PPL reduction of Dense is zero.

Backbone	Approach	General	Code	Math	Biomedical	Legal	Reasoning
Qwen3-8B Downstream	Best Baseline	$0.5785{\pm}0.0062$ (UpIT)	$\mathbf{0.3097{\pm}0.0082}$ (MoM)	$0.3414{\pm}0.0098$ (Self-MoE)	$0.7220{\pm}0.0059$ (UpIT)	$0.6219{\pm}0.0069$ (UpIT)	$\mathbf{0.4992{\pm}0.0091}$ (Self-MoE)
Qwen3-8B Downstream	CORTEX	$\mathbf{0.5900{\pm}0.0058}$	$0.3074{\pm}0.0085$	$\mathbf{0.3754{\pm}0.0102}$	$\mathbf{0.7295{\pm}0.0055}$	$\mathbf{0.6472{\pm}0.0064}$	$0.4929{\pm}0.0088$
Qwen3-32B Downstream	Best Baseline	$\mathbf{0.6732{\pm}0.0050}$ (UpIT)	$0.3930{\pm}0.0068$ (MoE)	$0.5322{\pm}0.0083$ (MoM)	$0.7950{\pm}0.0053$ (MoE)	$0.6939{\pm}0.0059$ (UpIT)	$0.6211{\pm}0.0076$ (Self-MoE)
Qwen3-32B Downstream	CORTEX	$0.6641{\pm}0.0047$	$\mathbf{0.4253{\pm}0.0071}$	$\mathbf{0.5471{\pm}0.0080}$	$\mathbf{0.8280{\pm}0.0049}$	$\mathbf{0.6974{\pm}0.0056}$	$\mathbf{0.6494{\pm}0.0072}$

Table 2: Comparison of downstream scores between CORTEX and the best baseline on each domain. We show the average

\pm

standard deviation over four seeds.

Variant	PPL	NLL	EM	$I^{mod}$	$\mathrm{SLS}$
Variant	PPL	NLL	EM	$I^{mod}$	COPY	REVERSE	ADD	LOOKUP	LOOKUP $\rightarrow$ ADD	REV $\rightarrow$ LOOKUP
CORTEX w/o Shared	$1.1302$	$0.1224$	$95.56$	$1.1182$	$0.1482$	$0.1667$	$0.1694$	$0.1318$	$0.1975$	$0.1558$
CORTEX	$\mathbf{1.0762}$	$\mathbf{0.0734}$	$\mathbf{97.61}$	$\mathbf{1.4560}$	$\mathbf{0.1976}$	$\mathbf{0.1742}$	$\mathbf{0.1755}$	$\mathbf{0.1769}$	$\mathbf{0.2054}$	$\mathbf{0.1726}$

Table 3: Effect of the shared module on the synthetic benchmark.

Figure 4: Selective lesion scores ( $\mathrm{SLS}$ ) of different modules on (a) 160M model, (b) Qwen3-8B, and (c) Qwen3-32B. Each entry measures the difference between the risk change on a domain and the average change on the remaining domains after removing a module.

Backbone	Method	NLL $\downarrow$	$I^{\mathrm{mod}}$	$\Delta^{\mathrm{les}}_{m,k}$
Backbone	Method	NLL $\downarrow$	$I^{\mathrm{mod}}$	Target	Other
	Random	0.4082	1.3860	0.0097	0.0086
160M	CORTEX	0.0734	1.4560	0.1901	0.0064
	Random	1.9179	1.0238	0.0374	0.0108
Qwen3-8B	CORTEX	1.7303	1.1534	0.1557	0.0106
	Random	1.6571	1.1996	0.0298	0.0097

Table 4: NLL,

I^{\mathrm{mod}}

, and lesion effects. Target and Other denote the average NLL changes after module removal on the paired domain and other domains, respectively. Results are averaged across seeds.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Meaning
$C$	Knowledge-domain label.
$d_{j}^{\mathrm{in}},d_{j}^{\mathrm{out}}$	Input and output dimensions of the $j$ -th trainable weight matrix.
$\mathcal{G}_{j},G_{j}$	Index set and the number of parameter groups in $\mathbf{W}_{j}$ .
$\mathbf{G}_{j,g,k}$	Domain-conditioned gradient of the $g$ -th parameter group in the $j$ -th trainable parameter matrix w.r.t. the $k$ -th knowledge domain.
$\mathcal{J},J$	Set and the number of trainable weight matrices in the backbone model.
$\mathcal{K},K$	Set and the number of knowledge domains.

Appendix

Table 6: List of key notations.

	$\displaystyle u^{t}_{j,g,0}=$	$\displaystyle\>\frac{2}{K(K-1)}\sum_{1\leq p<q\leq K}\frac{\langle\mathbf{G}^{t}_{j,g,p},\mathbf{G}^{t}_{j,g,q}\rangle_{F}}{\\|\mathbf{G}^{t}_{j,g,p}\\|_{F}\\|\mathbf{G}^{t}_{j,g,q}\\|_{F}},\quad j\in\mathcal{J},\ g\in\mathcal{G}_{j},$
	$\displaystyle u^{t}_{j,g,m}=$	$\displaystyle\>\frac{\\|\mathbf{G}^{t}_{j,g,\kappa(m)}\\|_{F}}{\sum_{k^{\prime}\in\mathcal{K}}\\|\mathbf{G}^{t}_{j,g,k^{\prime}}\\|_{F}}-\frac{1}{K-1}\sum_{q\in\mathcal{K}\setminus\{\kappa(m)\}}\frac{\langle\mathbf{G}^{t}_{j,g,\kappa(m)},\mathbf{G}^{t}_{j,g,q}\rangle_{F}}{\\|\mathbf{G}^{t}_{j,g,\kappa(m)}\\|_{F}\\|\mathbf{G}^{t}_{j,g,q}\\|_{F}},$
		$\displaystyle\quad j\in\mathcal{J},\ g\in\mathcal{G}_{j},m\in\mathcal{M}.$

Appendix

Algorithm 1 CORTEX

Hyperparameter	Synthetic 160M
Domains	6
Backbone	160M decoder-only
Optimizer	AdamW
Selected learning rate	$3{\times}10^{-4}$
Learning-rate candidates	$\{5{\times}10^{-4},4{\times}10^{-4},3{\times}10^{-4},2{\times}10^{-4},1{\times}10^{-4}\}$
Mini-batch size / samples per update	32 / 32

Appendix

Table 7: Training and modularization hyperparameters for the synthetic benchmark.

Hyperparameter	Qwen3-8B
Domains	6
Backbone	Qwen3-8B
Optimizer	Adafactor
Selected learning rate	$5{\times}10^{-5}$
Learning-rate candidates	$\{1{\times}10^{-4},9{\times}10^{-5},8{\times}10^{-5},7{\times}10^{-5},6{\times}10^{-5},$ $5{\times}10^{-5},4{\times}10^{-5},3{\times}10^{-5},2{\times}10^{-5},1{\times}10^{-5}\}$
Mini-batch size / samples per update	1 / 48

Appendix

Table 8: Training and modularization hyperparameters for Qwen3-8B.

Hyperparameter	Qwen3-32B
Domains	6
Backbone	Qwen3-32B
Optimizer	Adafactor
Selected learning rate	$2{\times}10^{-5}$
Learning-rate candidates	$\{1{\times}10^{-4},9{\times}10^{-5},8{\times}10^{-5},7{\times}10^{-5},6{\times}10^{-5},$ $5{\times}10^{-5},4{\times}10^{-5},3{\times}10^{-5},2{\times}10^{-5},1{\times}10^{-5}\}$
Mini-batch size / samples per update	1 / 48

Appendix

Table 9: Training and modularization hyperparameters for Qwen3-32B.

Baseline	COPY	REVERSE	ADD	LOOKUP	LOOKUP $\rightarrow$ ADD	REVERSE $\rightarrow$ LOOKUP
Dense	$96.82{\pm}0.31$	$91.36{\pm}0.44$	$88.27{\pm}0.52$	$92.74{\pm}0.39$	$84.91{\pm}0.61$	$86.38{\pm}0.57$
Random	$88.54{\pm}0.66$	$87.92{\pm}0.58$	$80.76{\pm}0.73$	$91.18{\pm}0.47$	$73.64{\pm}0.85$	$86.95{\pm}0.60$
MoE	$98.12{\pm}0.25$	$95.06{\pm}0.34$	$94.38{\pm}0.41$	$95.71{\pm}0.33$	$87.46{\pm}0.56$	$91.64{\pm}0.46$
MoM	$98.67{\pm}0.21$	$97.84{\pm}0.27$	$96.21{\pm}0.33$	$96.14{\pm}0.31$	$90.82{\pm}0.49$	$94.76{\pm}0.39$
UpIT	$96.21{\pm}0.36$	$93.71{\pm}0.42$	$92.04{\pm}0.48$	$96.48{\pm}0.32$	$86.39{\pm}0.59$	$90.43{\pm}0.50$
Self-MoE	$97.84{\pm}0.28$	$96.52{\pm}0.31$	$93.91{\pm}0.43$	$95.18{\pm}0.35$	$89.34{\pm}0.52$	$95.71{\pm}0.37$

Appendix

Table 10: Synthetic benchmark EM (%). We show the mean

\pm

standard deviation over 10 different seeds.

Baseline	General	Code	Math	Biomedical	Legal	Reasoning
Dense	$1.6766{\pm}0.0124$	$1.7949{\pm}0.0141$	$2.0047{\pm}0.0168$	$1.8220{\pm}0.0149$	$1.9425{\pm}0.0161$	$1.9683{\pm}0.0157$
Random	$1.7997{\pm}0.0175$	$1.7924{\pm}0.0158$	$2.0172{\pm}0.0194$	$1.9534{\pm}0.0182$	$2.0053{\pm}0.0197$	$1.9392{\pm}0.0171$
MoE	$1.6642{\pm}0.0119$	$1.7878{\pm}0.0136$	$1.9262{\pm}0.0155$	$1.7947{\pm}0.0142$	$1.9054{\pm}0.0153$	$1.9133{\pm}0.0149$
MoM	$1.6441{\pm}0.0112$	$1.6578{\pm}0.0124$	$1.9290{\pm}0.0151$	$1.7995{\pm}0.0138$	$1.8714{\pm}0.0147$	$1.8385{\pm}0.0136$
UpIT	$1.6257{\pm}0.0108$	$1.7658{\pm}0.0131$	$2.0039{\pm}0.0162$	$1.7709{\pm}0.0134$	$1.8171{\pm}0.0140$	$1.8288{\pm}0.0132$
Self-MoE	$1.6684{\pm}0.0121$	$1.7373{\pm}0.0129$	$1.9058{\pm}0.0148$	$1.8208{\pm}0.0145$	$1.8583{\pm}0.0146$	$\mathbf{1.7822{\pm}0.0128}$

Appendix

Table 11: Real-domain benchmark NLL on Qwen3-8B. We show the mean

\pm

standard deviation over 4 different seeds.

Baseline	General	Code	Math	Biomedical	Legal	Reasoning
Dense	$1.4202{\pm}0.0108$	$1.5311{\pm}0.0122$	$1.7400{\pm}0.0145$	$1.5648{\pm}0.0129$	$1.6712{\pm}0.0136$	$1.7221{\pm}0.0140$
Random	$1.5298{\pm}0.0149$	$1.5193{\pm}0.0135$	$1.7889{\pm}0.0161$	$1.6876{\pm}0.0152$	$1.7272{\pm}0.0157$	$1.6897{\pm}0.0148$
MoE	$1.3884{\pm}0.0099$	$1.4383{\pm}0.0107$	$1.6240{\pm}0.0128$	$1.4878{\pm}0.0114$	$1.6364{\pm}0.0130$	$1.6111{\pm}0.0125$
MoM	$1.3918{\pm}0.0101$	$1.4967{\pm}0.0119$	$1.5815{\pm}0.0122$	$1.5384{\pm}0.0120$	$1.5942{\pm}0.0127$	$1.6129{\pm}0.0124$
UpIT	$\mathbf{1.3578{\pm}0.0093}$	$1.4733{\pm}0.0112$	$1.6766{\pm}0.0134$	$1.5100{\pm}0.0117$	$1.5514{\pm}0.0119$	$1.5945{\pm}0.0121$
Self-MoE	$1.3871{\pm}0.0098$	$1.4395{\pm}0.0108$	$1.5845{\pm}0.0120$	$1.5180{\pm}0.0118$	$1.6011{\pm}0.0126$	$1.5597{\pm}0.0116$

Appendix

Table 12: Real-domain benchmark NLL on Qwen3-32B. We show the mean

\pm

standard deviation over 4 different seeds.

Backbone	General	Code	Math	Biomedical	Legal	Reasoning
Qwen3-8B	$+0.0115{\pm}0.0046$	$-0.0023{\pm}0.0065$	$+0.0340{\pm}0.0075$	$+0.0075{\pm}0.0044$	$+0.0253{\pm}0.0051$	$-0.0063{\pm}0.0069$
Qwen3-32B	$-0.0091{\pm}0.0038$	$+0.0323{\pm}0.0053$	$+0.0149{\pm}0.0061$	$+0.0330{\pm}0.0042$	$+0.0035{\pm}0.0046$	$+0.0283{\pm}0.0058$

Appendix

Table 13: Paired differences in downstream scores (CORTEX minus the best baseline in Table 2 ). We show the average

\pm

standard deviation of the differences over four paired training seeds.

Figure 5: PPL reduction over the Dense baseline, calculated from the average NLL across knowledge domains.

Backbone	Original	Matched shuffle	Difference
160M	$0.1837{\pm}0.0120$	$0.0548{\pm}0.0100$	$+0.1289{\pm}0.0115$
Qwen3-8B	$0.1452{\pm}0.0150$	$0.0427{\pm}0.0120$	$+0.1026{\pm}0.0140$

Appendix

Table 14: Average paired-domain SLS for original and matched-shuffle assignments. Difference denotes original minus matched-shuffle SLS. Values are averages

\pm

standard deviations across training seeds.

Figure 6: Domain-normalized gradient share ( $\mathrm{GS}_{m,k}/\mathrm{GS}_{k}^{\mathrm{domain}}$ ) on (a) 160M model, (b) Qwen3-8B, and (c) Qwen3-32B. The columns denote knowledge domains and the rows denote modules. Each column sums to one.

Figure 7: Visualization of modularization of CORTEX on the Qwen3-8B backbone model trained on the real-domain data mixture. The seven heatmaps correspond to the shared knowledge and six knowledge domains, including (a) Shared, (b) General, (c) Code, (d) Math, (e) Biomedical, (f) Legal, and (g) Reasoning, respectively. The columns denote Transformer blocks and the rows denote parameter groups within attention heads (Attn.) and the feed-forward network (FFN).

Figure 8: Visualization of modularization of CORTEX on the Qwen3-32B backbone model trained on the real-domain data mixture. The seven heatmaps correspond to the shared knowledge and six knowledge domains, including (a) Shared, (b) General, (c) Code, (d) Math, (e) Biomedical, (f) Legal, and (g) Reasoning, respectively. The columns denote Transformer blocks and the rows denote parameter groups within attention heads (Attn.) and the feed-forward network (FFN).

Figure 9: Comparison of training step time on the 160M, Qwen3-8B, and Qwen3-32B backbones.

Figure 10: Comparison of peak memory per GPU on the 160M, Qwen3-8B, and Qwen3-32B backbones.

Explore similar work

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts
Apr 20, 2026Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia +3Mixture-Of-Experts Post-Training
University of Washington · Allen Institute for AI · University of California, Berkeley
Decodable But Not Detachable: Training Data Granularity Determines Parametric Modularity in Large Language Models
Aug 10, 2026Marcus Armstrong, Navid Ayoobi, Arjun MukherjeeNeurons Selectivity
Department of Computer Science University of Houston Houston, TX 77204
EMO: Pretraining Mixture of Experts for Emergent Modularity
May 7, 2026Ryan Wang, Akshita Bhagia, Sewon MinMixture-Of-Experts Experts
University of California, Berkeley · Allen Institute for AI

Apr 20, 2026, cs.LGJ/K move · Enter open · S save

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts

Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia, Matei Zaharia+2

University of Washington · Allen Institute for AI · University of California, Berkeley

Mixture-Of-Experts Post-Training Retraining Modular Framework Merge Transport

Open paper AskPDF arXiv

Extending a fully post-trained language model with new domain capabilities is fundamentally limited by monolithic training paradigms: retraining from scratch is expensive and scales poorly, while continued training often degrades existing capabilities. We present BAR (Branch-Adapt-Route), which trains independent domain experts, each through its own mid-training, supervised finetuning, and reinforcement learning pipeline, and composes them via a Mixture-of-Experts architecture with lightweight router training. Unlike retraining approaches that mix all domains and require full reprocessing for any update (with cost scaling quadratically), BAR enables updating individual experts independently with linear cost scaling and no degradation to existing domains. At the 7B scale, with experts for math, code, tool use, and safety, BAR achieves an overall score of 49.1 (averaged across 7 evaluation categories), matching or exceeding re-training baselines (47.8 without mid-training, 50.5 with). We further show that modular training provides a structural advantage: by isolating each domain, it avoids the catastrophic forgetting that occurs when late-stage RL degrades capabilities from earlier training stages, while significantly reducing the cost and complexity of updating or adding a domain. Together, these results suggest that decoupled, expert-based training is a scalable alternative to monolithic retraining for extending language models.

Abstract

Explore similar work

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts

Decodable But Not Detachable: Training Data Granularity Determines Parametric Modularity in Large Language Models

EMO: Pretraining Mixture of Experts for Emergent Modularity