Optimal Transport Depth Up-Scaling

Mingzi Cao; Xi Wang; Nikolaos Aletras

cs.CLAug 11, 2025

Optimal Transport Depth Up-Scaling

Authors: Mingzi Cao, Xi Wang, Nikolaos Aletras

Organizations: University of Sheffield

Abstract

Pre-training Large Language Models (LLMs) from scratch at larger scales yields remarkable performance but incurs substantially high training costs. Depth up-scaling provides an efficient alternative by inserting new layers into a pre-trained LLM, avoiding training from scratch. However, most existing methods copying or averaging base layers for new layer, which misalign functionally corresponding neurons, leading to neuron permutation mismatch that harms performance. To address this issue, we propose Optimal Transport Depth Up-Scaling (OT-DUS), which leverages Optimal Transport (OT) theory to align and fuse functionally corresponding neurons module by module in adjacent base layers for new layer construction. OT-DUS achieves better overall performance in both general and specialized domains than existing methods for continual pre-training and supervised fine-tuning across different model sizes and model families. Our analysis of insertion strategies further finds that inserting new layers at higher positions yields not only stronger performance but also improved training efficiency.

Figures & tables

Figure 1: Illustration of our proposed OT-DUS. OT-DUS uses optimal transport to initialize new layers, each derived from two adjacent base layers

f_{i}

and

f_{i+1}

. It first aligns each module

b

to previous module

b-1

f_{i}

(i.e. Alignment within Layer), then aligns it to

b

f_{i+1}

(i.e. Alignment across Layer).

Figure 2: State-of-the-art depth up-scaling methods and our proposed OT-DUS.

Figure 3: Neuron permutation mismatch caused by average and duplicate. Each column represents the order of neurons within layer. Multiple colours of each neuron represent multiple functions it contributes to.

	Normalization		Attention				MLP
Transport Matrix	Pre-Attn	Pre-MLP	Query	Key	Value	Output	Gate	Up	Down
$\mathbf{T}_{\text{in}}$	$\mathbf{I}$	$\frac{1}{2}(\mathbf{T}_{O}+\mathbf{I})$	$\mathbf{I}$	$\mathbf{I}$	$\mathbf{I}$	$\underline{\mathbf{I}}$	$\mathbf{T}_{O}$	$\mathbf{T}_{O}$	$\underline{\mathbf{I}}$
$\mathbf{T}_{\text{out}}$	$\mathbf{I}$	$\frac{1}{2}(\mathbf{T}_{O}+\mathbf{I})$	$\mathbf{T}_{Q}$	$\mathbf{T}_{K}$	$\mathbf{T}_{V}$	$\mathbf{T}_{O}$	$\mathbf{T}_{\text{Gate}}$	$\mathbf{T}_{\text{Up}}$	$\mathbf{T}_{\text{Down}}$

Table 1: Transport Matrix Flow. We manually set

\mathbf{T}_{\text{in}}

to each module for alignment within layer.

\mathbf{T}_{\text{out}}

is calculated through OT for alignment across layers (except Normalization component). Modules highlighted in underline deviate from Imfeld et al. (2024) .

		Perplexity $\downarrow$	Zero-shot Performance $\uparrow$
	Methods	Wiki-PPL	ARC	LogiQA	Wino	CSQA	BoolQ	PIQA	MMLU	Average
CPT	Base-8B	$8.35_{0.00}$	$79.91_{0.05}$	$26.68_{0.18}$	$72.19_{0.30}$	$65.19_{0.08}$	$81.85_{0.04}$	$79.04_{0.26}$	$58.86_{0.22}$	$66.25_{0.04}$
	SOLAR-11.5B	$9.90_{0.00}$	$80.02_{0.15}$	$26.88_{0.00}$	$71.80_{0.24}$	$57.36_{0.05}$	$80.76_{0.05}$	$78.51_{0.05}$	$54.32_{0.05}$	$64.24_{0.05}$
	LLaMA PRO-11.5B	$7.81_{0.00}$	$81.45_{0.15}$	$\mathbf{29.39}_{\mathbf{0.09}}$	$73.93_{0.24}$	$70.71_{0.21}$	$81.71_{0.06}$	$79.89_{0.08}$	$62.59_{0.03}$	$68.52_{0.04}$
	LESA-11.5B	$7.73_{0.00}$	$81.92_{0.15}$	$28.06_{0.09}$	$74.09_{0.12}$	$\mathbf{72.32}_{\mathbf{0.14}}$	$81.89_{0.04}$	$80.36_{0.06}$	$62.61_{0.03}$	$68.75_{0.03}$
	OT-DUS-11.5B (Ours)	$\mathbf{7.73}_{\mathbf{0.00}}$	$\mathbf{82.06}_{\mathbf{0.06}}$	$27.34_{0.00}$	$\mathbf{74.77}_{\mathbf{0.20}}$	$71.96_{0.33}$	$\mathbf{82.18}_{\mathbf{0.07}}$	$\mathbf{80.56}_{\mathbf{0.22}}$	$\mathbf{62.99}_{\mathbf{0.02}}$	$\mathbf{68.84}_{\mathbf{0.04}}$

Table 2: CPT on 1.5B tokens and SFT (after CPT) performance of 11.5B and 1.72B expanded models.

	Methods	MedMCQA	MedQA	Bio-MMLU	PubMedQA
Biomedical	Base	$48.92_{0.28}$	$53.26_{0.00}$	$68.35_{0.10}$	$76.33_{0.12}$
	SOLAR	$45.22_{0.13}$	$47.24_{0.05}$	$57.41_{0.14}$	$76.20_{0.00}$
	LLaMA PRO	$54.47_{0.22}$	$59.31_{0.14}$	$71.16_{0.27}$	$\mathbf{77.67}_{\mathbf{0.12}}$
	LESA	$56.73_{0.19}$	$60.07_{0.18}$	$\mathbf{71.40}_{\mathbf{0.36}}$	$76.00_{0.00}$
	OT-DUS	$\mathbf{56.79}_{\mathbf{0.16}}$	$\mathbf{60.36}_{\mathbf{0.09}}$	$71.30_{0.29}$	$76.33_{0.12}$
	Methods	Legal-MMLU	CaseHOLD		ECtHR

Table 3: Domain-specific CPT performance of 11.5B expanded models.

Interpolation Strategy	Average
OT-DUS-Btm	$68.39_{0.12}$
OT-DUS-Mid	$68.28_{0.01}$
OT-DUS-Top	$\mathbf{68.84}_{\mathbf{0.04}}$
OT-DUS-T&B	$68.44_{0.04}$

Table 4: Average performance of OT-DUS-11.5B using different interpolation strategies. Full results are presented in Table 11 (Appendix F ).

Methods	Trainable	Training Time
SOLAR-11.5B	11B	22:54:11 (+78.0%)
LLaMA PRO-11.5B	4B	14:58:34 (+16.4%)
LESA-11.5B	4B	12:54:07 (+0.3%)
OT-DUS-Btm-11.5B	4B	14:56:00 (+16.1%)
OT-DUS-Mid-11.5B	4B	13:53:14 (+7.9%)
OT-DUS-Top-11.5B	4B	12:52:04

Table 5: Training time for 11.5B models.

	Ablation Variants	Average
TMF	Attention Output $\mathbf{T}_{\text{in}}=\mathbf{T}_{Q}$ -1.72B	$51.96_{0.03}$
TMF	MLP Gate and Up $\mathbf{T}_{\text{in}}=\frac{1}{2}(\mathbf{T}_{O}+\mathbf{I})$ )-1.72B	$51.79_{0.02}$
OT	Average-1.72B	$51.59_{0.61}$
	Copy-1.72B	$50.13_{0.01}$
	Random-1.72B	$51.39_{0.13}$
Zero-Init	Average + Zero-Init-1.72B	$51.20_{0.61}$

Table 6: Average performance of ablation variants. Full Results are presented at Table 12 (Appendix G ).

Methods	Average
Base-8B	$73.59_{0.01}$
SOLAR-11.5B	$72.36_{0.00}$
LLaMA PRO-11.5B	$73.80_{0.04}$
LESA-11.5B	$73.76_{0.17}$
OT-DUS-11.5B (Ours)	$\mathbf{74.47}_{\mathbf{0.06}}$

Table 7: Average Performance of expanded models using Qwen3-8B as base model. Full results are presented in Table 13 (Appendix G ).

	Attention				MLP
Methods	Query	Key	Value	Output	Gate	Up	Down	Average
LLaMA PRO	12.23	9.99	10.29	16.64	14.17	14.20	16.63	13.45
Average	11.69	7.32	10.90	15.60	17.06	18.71	8.56	12.84
OT-DUS	13.25	10.32	11.57	16.64	16.78	18.36	16.64	14.79

Table 8: Layer-averaged Entropy of the transport matrix between base and new layers in the 11.5B expanded models. Normalization components are excluded as they do not involve cross-neuron mixing.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Number of benchmarks that achieve top performance during the training process of 11.5B expanded models. Sums may exceed 8 due to ties.

Model	Base	SOLAR	LLaMA PRO	LESA	OT-DUS
Llama-3.2-1B	11.57	16.64	11.57	871.50	11.57
Llama-3.1-8B	7.33	9.01	7.33	9.35	7.33
Mistral-24B	4.43 ∗	6.51 ∗	4.43	5.17 ∗	4.43
Qwen-2.5-32B	3.78 ∗	INF ∗	3.78	5.67 ∗	3.78
Llama-3-70B	1.98 ∗	4.21 ∗	1.98	2.62 ∗	1.98

Appendix

Table 9: PPL after 1.5x layer expansion initialization for different base models, along with PPL of base models. * denotes results from Yang et al. (2025b) , Results for LLaMA PRO and OT-DUS in bottom-half are obtained via reasonable extrapolation from top-half.

Expanded Model	Creating Time
LESA-1.72B	00:26:15
OT-DUS-1.72B	00:02:34
LESA-11.5B	04:52:13
OT-DUS-11.5B	00:37:16

Appendix

Table 10: Instantiation time for LESA and OT-DUS.

	Perplexity $\downarrow$	Zero-shot Performance $\uparrow$
Methods	Wiki-PPL	ARC	LogiQA	Wino	CSQA	BoolQ	PIQA	MMLU	Average
OT-DUS-Btm	$7.83_{0.00}$	$81.75_{0.10}$	$28.21_{0.09}$	$74.59_{0.41}$	$70.10_{0.14}$	$81.75_{0.02}$	$79.99_{0.13}$	$62.33_{0.08}$	$68.39_{0.12}$
OT-DUS-Mid	$\mathbf{7.70}_{\mathbf{0.00}}$	$\mathbf{82.07}_{\mathbf{0.00}}$	$27.65_{0.00}$	$74.32_{0.05}$	$70.25_{0.23}$	$81.00_{0.12}$	$80.20_{0.09}$	$62.53_{0.06}$	$68.28_{0.01}$
OT-DUS-Top	$7.73_{0.00}$	$82.06_{0.06}$	$27.34_{0.00}$	$\mathbf{74.77}_{\mathbf{0.20}}$	$\mathbf{71.96}_{\mathbf{0.33}}$	$\mathbf{82.18}_{\mathbf{0.07}}$	$\mathbf{80.56}_{\mathbf{0.22}}$	$\mathbf{62.99}_{\mathbf{0.02}}$	$\mathbf{68.84}_{\mathbf{0.04}}$
OT-DUS-T&B	$7.87_{0.00}$	$81.46_{0.10}$	$\mathbf{28.57}_{\mathbf{0.00}}$	$74.61_{0.18}$	$70.05_{0.05}$	$82.12_{0.02}$	$79.85_{0.03}$	$62.44_{0.03}$	$68.44_{0.04}$

Appendix

Table 11: Performance of 11.5B OT-DUS trained on 1.5B tokens using different interpolation strategies.

		Perplexity $\downarrow$	Zero-shot Performance $\uparrow$
	Methods	Wiki-PPL	ARC	LogiQA	Wino	CSQA	BoolQ	PIQA	MMLU	Average
TMF	Attention Output $\mathbf{T}_{\text{in}}=\mathbf{T}_{Q}$ -1.72B	$\mathbf{12.05}_{\mathbf{0.00}}$	$67.12_{0.02}$	$\mathbf{23.30}_{\mathbf{0.09}}$	$60.25_{0.09}$	$42.64_{0.05}$	$62.89_{0.04}$	$75.39_{0.03}$	$32.13_{0.05}$	$51.96_{0.03}$
TMF	MLP Gate and Up $\mathbf{T}_{\text{in}}=\frac{1}{2}(\mathbf{T}_{O}+\mathbf{I})$ )-1.72B	$12.09_{0.00}$	$66.81_{0.05}$	$22.79_{0.18}$	$\mathbf{60.46}_{\mathbf{0.14}}$	$42.59_{0.00}$	$62.22_{0.02}$	$\mathbf{75.48}_{\mathbf{0.03}}$	$32.20_{0.07}$	$51.79_{0.02}$
OT	Average-1.72B	$12.31_{0.27}$	$67.12_{0.53}$	$22.38_{0.23}$	$59.75_{0.50}$	$41.71_{2.16}$	$62.76_{0.22}$	$74.67_{0.14}$	$32.74_{1.75}$	$51.59_{0.61}$
	Copy-1.72B	$12.62_{0.00}$	$\mathbf{68.00}_{\mathbf{0.02}}$	$22.73_{0.00}$	$59.75_{0.14}$	$34.75_{0.10}$	$62.96_{0.02}$	$74.52_{0.16}$	$28.20_{0.09}$	$50.13_{0.01}$
	Random-1.72B	$16.40_{0.00}$	$66.79_{0.29}$	$22.27_{0.00}$	$59.85_{0.32}$	$40.95_{0.28}$	$62.69_{0.05}$	$74.99_{0.03}$	$32.19_{0.00}$	$51.39_{0.13}$

Appendix

Table 12: Ablation study on 1.72B and 11.5B expanded models.

	Perplexity $\downarrow$	Zero-shot Performance $\uparrow$
Methods	Wiki-PPL	ARC	LogiQA	Wino	CSQA	BoolQ	PIQA	MMLU	Average
Base-8B	$\mathbf{9.20}_{\mathbf{0.00}}$	$84.72_{0.15}$	$39.32_{0.00}$	$73.61_{0.32}$	$82.25_{0.19}$	$82.27_{0.02}$	$79.47_{0.13}$	$73.47_{0.00}$	$73.59_{0.01}$
SOLAR-11.5B	$10.47_{0.00}$	$84.25_{0.12}$	$39.32_{0.00}$	$72.11_{0.18}$	$77.01_{0.05}$	$84.56_{0.05}$	$78.44_{0.06}$	$70.84_{0.03}$	$72.36_{0.00}$
LLaMA PRO-11.5B	$9.24_{0.00}$	$\mathbf{85.49}_{\mathbf{0.05}}$	$37.94_{0.00}$	$\mathbf{74.11}_{\mathbf{0.00}}$	$81.71_{0.19}$	$84.25_{0.00}$	$79.49_{0.09}$	$73.60_{0.03}$	$73.80_{0.04}$
LESA-11.5B	$10.75_{0.00}$	$82.70_{0.07}$	$40.30_{0.09}$	$73.48_{1.09}$	$81.63_{0.05}$	$84.80_{0.05}$	$79.52_{0.03}$	$\mathbf{73.88}_{\mathbf{0.05}}$	$73.76_{0.17}$
OT-DUS-11.5B (Ours)	$10.58_{0.00}$	$83.38_{0.00}$	$\mathbf{41.83}_{\mathbf{0.09}}$	$73.48_{0.41}$	$\mathbf{84.08}_{\mathbf{0.09}}$	$\mathbf{85.15}_{\mathbf{0.02}}$	$\mathbf{79.54}_{\mathbf{0.19}}$	$73.84_{0.02}$	$\mathbf{74.47}_{\mathbf{0.06}}$

Appendix

Table 13: CPT Performance of Expanded models using Qwen3-8B as base model.

Language Model Scaling

Language model scaling studies how model size, training data, compute, and inference-time resources affect capability and efficiency.

See all 60 papers

Universal YOCO for Efficient Depth Scaling
Apr 1, 2026Yutao Sun, Li Dong, Tianzhu Ye +3Efficient Language Model Inference Inference-Time Scaling
Microsoft Research · Tsinghua University
Adaptive Depth Sparse Framework: Similarity-Driven Resource Allocation for Pre-Trained LLMs
Jul 23, 2026Yidu Wu, Xiang Wang, Kejie Zhao +3LLM Pruning LLM Inference Acceleration
Department of Electronic and Electrical Engineering, Southern University of Science and Technology, Shenzhen, China · ACS Lab, Huawei Technologies Co., Ltd., Shenzhen, China · Department of Computer Science and Engineering, Southern University of Science and Technology, Shenzhen, China
Test-Time Scaling Makes Overtraining Compute-Optimal
Apr 1, 2026Nicholas Roberts, Sungjun Cho, Zhiqi Gao +7Scaling Laws Language Model Scaling Laws
University of Wisconsin-Madison · Stanford University

Apr 1, 2026J/K move · Enter open · S save

Universal YOCO for Efficient Depth Scaling

Yutao Sun, Li Dong, Tianzhu Ye, Shaohan Huang+2

Microsoft Research · Tsinghua University

Efficient Language Model Inference Inference-Time Scaling Transformer

Open paper AskPDF arXiv

The rise of test-time scaling has remarkably boosted the reasoning and agentic proficiency of Large Language Models (LLMs). Yet, standard Transformers struggle to scale inference-time compute efficiently, as conventional looping strategies suffer from high computational overhead and a KV cache that inflates alongside model depth. We present Universal YOCO (YOCO-U), which combines the YOCO decoder-decoder architecture with recursive computation to achieve a synergistic effect greater than either alone. Built on the YOCO framework, YOCO-U implements a Universal Self-Decoder that performs multiple iterations via parameter sharing, while confining the iterative process to shallow, efficient-attention layers. This combination yields a favorable capability-efficiency tradeoff that neither YOCO nor recursion achieves independently. The YOCO architecture provides a constant global KV cache and linear pre-filling, while partial recursion enhances representational depth with limited overhead. Together, YOCO-U improves token utility and scaling behavior while maintaining efficient inference. Empirical results confirm that YOCO-U remains highly competitive in general and long-context benchmarks, demonstrating that the integration of efficient-attention architectures and recursive computation is a promising direction for scalable LLMs.

Abstract

Language Model Scaling

Explore similar work

Universal YOCO for Efficient Depth Scaling

Adaptive Depth Sparse Framework: Similarity-Driven Resource Allocation for Pre-Trained LLMs

Test-Time Scaling Makes Overtraining Compute-Optimal