OpT-DeUS: Optimal Transport Depth Up-Scaling

• The paper introduces OT-based insertion to construct new Transformer layers by aligning neurons between adjacent blocks, overcoming permutation mismatch.
• It employs blockwise OT alignment via Sinkhorn-Knopp and Transport Matrix Flow to fuse weights and preserve function while freezing original layers.
• Empirical results show that OpT-DeUS outperforms naive methods with improved zero-shot and fine-tuning scores, alongside faster model creation.

Optimal Transport Depth Up-Scaling (OpT-DeUS) is a progressive depth expansion method for Transformers and LLMs in which new layers are created by aligning adjacent pretrained Transformer blocks with optimal transport before fusing them. The method addresses the setting in which a pretrained model with $n$ Transformer layers is expanded to $n+k$ layers, after which only the newly inserted layers are trained during the expansion stage. Its central premise is that conventional copying or coordinate-wise averaging neglects neuron permutation mismatch between adjacent layers, whereas OT-based alignment can construct inserted layers from matched rather than merely index-aligned neurons (Cao et al., 11 Aug 2025).

1. Problem setting and conceptual basis

The immediate target of OpT-DeUS is depth up-scaling for pretrained Transformer stacks rather than geometric depth estimation. The paper studies the regime in which a base model $\mathcal M$ with layers $f_1,\dots,f_n$ is expanded into a deeper model $\mathcal M'$ with $m=n+k$ layers by inserting new blocks between existing adjacent layers. The inserted layer $f'_i$ is built from the neighboring base layers $f_i$ and $f_{i+1}$, and the expansion is followed by continual pre-training or supervised fine-tuning. In the progressive variant emphasized by the paper, only the inserted layers are trainable during the depth up-scaling stage, while the original pretrained layers remain frozen (Cao et al., 11 Aug 2025).

The motivating failure mode is neuron permutation mismatch. The paper states that Transformer blocks can implement similar computations using different internal neuron orderings, so two neighboring layers can be similar up to permutation. Under that view, naive copying and naive averaging are limited because they combine weights coordinate-wise without first aligning functionally corresponding neurons. The initialization family is presented explicitly as

$$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$

At a formal level, OpT-DeUS instantiates the standard discrete OT template on a neuron set rather than on pixels or points. In the broader OT literature, a transport matrix $n+k$0 minimizes a linear cost under marginal constraints on discrete supports (Solomon, 2018). OpT-DeUS adopts that structure block-wise: neurons or channels in one layer serve as source support points, neurons or channels in the next layer serve as target support points, and the resulting transport plan is used to align one block to the other before averaging (Cao et al., 11 Aug 2025).

2. OT-based alignment and inserted-layer construction

For each corresponding block pair $n+k$1 and $n+k$2, the method first defines two discrete measures with uniform marginals, denoted $n+k$3 and $n+k$4. The paper states that neurons are treated equally, so the marginals are initialized uniformly. A support function $n+k$5 represents a neuron by its weights, using the weight-based representation from Singh and Jaggi. The block-level cost matrix is then

$n+k$6

where $n+k$7 and $n+k$8 are the $n+k$9-th and $\mathcal M$0-th neurons of the two block matrices (Cao et al., 11 Aug 2025).

The OT problem solved by OpT-DeUS is written as

$\mathcal M$1

subject to

$\mathcal M$2

The resulting transport matrix is used in two distinct alignment steps. First, a within-layer alignment applies a transport inherited from preceding blocks: $\mathcal M$3 Second, an across-layer alignment applies the OT solution itself: $\mathcal M$4 After alignment, the inserted block is initialized by simple averaging: $\mathcal M$5 This averaging is therefore not raw interpolation but interpolation after neuron matching (Cao et al., 11 Aug 2025).

The paper states that the OT problem is solved with the Sinkhorn-Knopp algorithm and that the regularization parameter is set to $\mathcal M$6, following prior work. It does not print the entropically regularized objective or the Sinkhorn iteration equations, but that solver choice is consistent with the established role of Sinkhorn as a standard large-scale regularized OT workhorse in the numerical OT literature (Cao et al., 11 Aug 2025, Tupitsa et al., 2022).

A final function-preserving step zero-initializes two difficult blocks: $\mathcal M$7 These are the attention output projection $\mathcal M$8 and the MLP down projection $\mathcal M$9. The stated reason is that their inputs are influenced by multiple blocks and residual paths, making clean alignment harder; zero-initialization avoids misalignment while preserving function (Cao et al., 11 Aug 2025).

3. Transport Matrix Flow and blockwise propagation of alignment

A distinctive engineering component of OpT-DeUS is Transport Matrix Flow (TMF), which propagates alignment information through the internal structure of a Transformer block. The need for TMF arises because once one block is aligned or permuted, the basis expected by later blocks in the same layer changes. The paper therefore distinguishes two transport matrices: $f_1,\dots,f_n$0, which aligns a block to the previous block within the same layer, and $f_1,\dots,f_n$1, which aligns a block across adjacent layers (Cao et al., 11 Aug 2025).

The rules reported for TMF are structurally specific. At the layer entrance of the inserted block, $f_1,\dots,f_n$2. For the pre-attention RMSNorm block, $f_1,\dots,f_n$3 and this identity alignment propagates to the query, key, and value projections. For the post-attention or pre-MLP RMSNorm block, $f_1,\dots,f_n$4 is set by averaging the $f_1,\dots,f_n$5 matrices from both residual paths. For the attention output projection $f_1,\dots,f_n$6, the paper sets $f_1,\dots,f_n$7 for computational simplicity and subsequently zero-initializes the inserted $f_1,\dots,f_n$8. For the MLP gate projection $f_1,\dots,f_n$9 and up projection $\mathcal M'$0, $\mathcal M'$1 is set to $\mathcal M'$2 from the attention module. For the MLP down projection $\mathcal M'$3, $\mathcal M'$4, followed by zero-initialization of $\mathcal M'$5 (Cao et al., 11 Aug 2025).

This blockwise treatment is one of the main ways OpT-DeUS differs from a monolithic whole-layer fusion rule. The method operates on corresponding Attention and MLP sub-blocks rather than on a single concatenated parameter tensor. The paper explicitly discusses the query, key, value, attention output, MLP gate, MLP up, and MLP down projections, together with the normalization blocks. It does not describe a separate head-level OT matching procedure, and it does not describe embedding alignment or fusion (Cao et al., 11 Aug 2025).

A plausible implication is that TMF treats layer insertion not merely as interpolation between two parameter vectors but as constrained model fusion under local basis changes induced by sequential composition and residual connections. That interpretation is consistent with the paper’s explicit connection to permutation symmetry and OT-based model fusion, even though the paper keeps the implementation at the block level rather than formalizing a global layerwise transport system (Cao et al., 11 Aug 2025).

4. Expansion regime, insertion positions, and implementation setting

OpT-DeUS inserts new layers in the top half of the base model by default. For a base stack of $\mathcal M'$6 layers, Algorithm 1 loops over

$\mathcal M'$7

The paper also evaluates four insertion strategies: $\mathcal M'$8 corresponding respectively to insertion in the bottom half, middle portion, top half, and top-and-bottom quarters. The paper later argues experimentally that top insertion is preferable both for performance and for training efficiency because it yields shorter back-propagation time when only inserted layers are trainable (Cao et al., 11 Aug 2025).

The evaluated base models are Llama-3.1-8B with $\mathcal M'$9 layers and Llama-3.2-1B with $m=n+k$0 layers. The expanded models are $m=n+k$1B with $m=n+k$2 layers and $m=n+k$3B with $m=n+k$4 layers, corresponding to adding $m=n+k$5 and $m=n+k$6 layers respectively, or $m=n+k$7 more layers. Continual pre-training uses $m=n+k$8B tokens sampled from the CC-MAIN-2024-51 subset of FineWeb-Edu. Supervised fine-tuning uses Alpaca GPT4, and the paper notes that unlike progressive CPT, the SFT stage updates the whole model following LESA (Cao et al., 11 Aug 2025).

The reported optimization settings are batch size $m=n+k$9 and sequence length $f'_i$0. The CPT maximum learning rate is $f'_i$1 for $f'_i$2B expanded models and $f'_i$3 for $f'_i$4B expanded models. The SFT maximum learning rate is $f'_i$5 for $f'_i$6B and $f'_i$7 for $f'_i$8B. Hardware is also stated explicitly: $f'_i$9B expanded models are trained on $f_i$0 NVIDIA GH200 $f_i$1 GPUs, $f_i$2B expanded models on $f_i$3 NVIDIA A100 $f_i$4, and model creation on an AMD EPYC 7413 CPU with $f_i$5 NVIDIA A100 $f_i$6 (Cao et al., 11 Aug 2025).

The method applies OT once at model creation time. After the inserted layers are constructed, training proceeds in the progressive expansion regime or in the downstream SFT regime. The paper does not provide asymptotic time or memory complexity formulas for OT alignment; empirical creation-time measurements are used instead (Cao et al., 11 Aug 2025).

5. Empirical behavior, initialization stability, and efficiency

The central empirical claim is that OT-based alignment improves over both existing depth up-scaling baselines and plain adjacent-layer averaging. On continual pre-training for the $f_i$7B model, the average zero-shot score is $f_i$8 for OpT-DeUS, compared with $f_i$9 for LESA, $f_{i+1}$0 for LLaMA PRO, $f_{i+1}$1 for Avg-DeUS, $f_{i+1}$2 for SOLAR, and $f_{i+1}$3 for the Base-8B model. On supervised fine-tuning for the same scale, the average is $f_{i+1}$4 for OpT-DeUS, compared with $f_{i+1}$5 for LESA and $f_{i+1}$6 for Avg-DeUS. At the $f_{i+1}$7B scale, OpT-DeUS reaches $f_{i+1}$8 on CPT and $f_{i+1}$9 on SFT, compared with $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$0 and $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$1 for LESA and $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$2 and $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$3 for Avg-DeUS. The cleanest OT-specific ablation is the comparison against Avg-DeUS: the gain is $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$4 average points at $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$5B and $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$6 average points at $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$7B, which the paper presents as direct evidence that neuron alignment improves over unaligned interpolation (Cao et al., 11 Aug 2025).

The position study reinforces the top-half insertion rule. On $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$8B CPT, the reported average scores are $$f'_i \leftarrow \begin{cases} f_i, & \text{if copying} \ \mathrm{Avg}(f_i,f_{i+1}), & \text{if averaging} \ \mathrm{NN}(f_i,f_{i+1}), & \text{if predicting} \ \mathrm{OpT\text{-}DeUS}(f_i,f_{i+1}), & \text{if using OT}. \end{cases}$$9 for OpT-DeUS-Btm, $n+k$00 for OpT-DeUS-Mid, $n+k$01 for OpT-DeUS-Top, and $n+k$02 for OpT-DeUS-T{data}B. The training-time comparison shows the same ordering in efficiency: OpT-DeUS-Btm takes $n+k$03, OpT-DeUS-Mid $n+k$04, OpT-DeUS-Top $n+k$05, and OpT-DeUS-T{data}B $n+k$06. The paper attributes the efficiency advantage of higher insertion to shorter back-propagation time (Cao et al., 11 Aug 2025).

Creation time is another practical result. For the $n+k$07B expanded model, LESA requires $n+k$08 for creation, whereas OpT-DeUS requires $n+k$09. For the $n+k$10B expanded model, LESA requires $n+k$11 and OpT-DeUS $n+k$12. When creation and training are combined for the $n+k$13B model, LESA totals $n+k$14, whereas OpT-DeUS totals $n+k$15. The paper therefore presents OpT-DeUS as substantially faster to create than LESA while remaining slightly stronger on downstream averages (Cao et al., 11 Aug 2025).

The initialization-only perplexity results are used to argue strong function preservation. Immediately after $n+k$16 layer expansion and before training, OpT-DeUS matches the base or LLaMA PRO perplexity on all reported models: for Llama-3.2-1B, Base $n+k$17, LLaMA PRO $n+k$18, OpT-DeUS $n+k$19; for Llama-3.1-8B, Base $n+k$20, LLaMA PRO $n+k$21, OpT-DeUS $n+k$22; for Mistral-24B, Base $n+k$23, LLaMA PRO $n+k$24, OpT-DeUS $n+k$25; for Qwen-2.5-32B, Base $n+k$26, LLaMA PRO $n+k$27, OpT-DeUS $n+k$28; and for Llama-3-70B, Base $n+k$29, LLaMA PRO $n+k$30, OpT-DeUS $n+k$31. By contrast, LESA reaches $n+k$32 on Llama-3.2-1B in this initialization-only test. The paper interprets that instability as evidence that smaller models may provide too little layer data for LESA’s auxiliary predictor, while OpT-DeUS does not show that pathology (Cao et al., 11 Aug 2025).

6. Relation to the broader OT literature and terminological distinctions

Within OT methodology, OpT-DeUS is a discrete, blockwise, regularized matching procedure rather than a new transport solver. Its transport matrix is a standard discrete coupling under marginal constraints, and its use of Sinkhorn-Knopp places it within the now-standard entropic OT toolkit developed for large-scale discrete problems (Solomon, 2018, Tupitsa et al., 2022). In that sense, its novelty lies in the modeling target—neuron alignment between adjacent Transformer layers—rather than in a new OT objective class (Cao et al., 11 Aug 2025).

The title can be misunderstood because “depth” appears in several unrelated OT-adjacent literatures. In "Provable optimal transport with transformers: The essence of depth and prompt engineering" (Daneshmand, 2024), depth means the number of Transformer layers used to simulate more steps of adaptive gradient descent on the dual of entropically regularized Wasserstein-$n+k$33. That paper shows a different relation between OT and Transformer depth: deeper transformers solve a fixed OT problem more accurately because each layer corresponds to one optimization step. OpT-DeUS instead uses OT to create additional layers in a pretrained model (Cao et al., 11 Aug 2025, Daneshmand, 2024).

A second possible confusion comes from geometric depth estimation. OT has also been used for uncertainty-guided local depth-distribution supervision in sparse-view 3D Gaussian Splatting, where patch-wise entropic OT complements pointwise depth losses under uncertainty (Sun et al., 2024). That setting concerns geometric depth priors and local depth distributions, not Transformer depth expansion. OpT-DeUS is therefore not a depth super-resolution or monocular-depth method. It is a model growth method for LLMs that borrows OT as a neuron-matching mechanism (Cao et al., 11 Aug 2025, Sun et al., 2024).

This suggests that OpT-DeUS belongs most naturally to the intersection of three themes: permutation symmetry in neural networks, OT-based model fusion, and progressive layer expansion. The paper explicitly situates itself against copying, averaging, predictor-based insertion, and function-preserving zero-initialization methods, while borrowing the discrete alignment machinery of OT to resolve mismatch between adjacent learned bases (Cao et al., 11 Aug 2025).

7. Scope, assumptions, and limitations

The paper states or implies several restrictions on the scope of OpT-DeUS. First, it assumes adjacent layers are functionally similar enough that interpolation or fusion is meaningful. Second, it relies on block-wise weight similarity as the OT cost through the Euclidean weight-space metric

$n+k$34

and does not use an activation-based cost. Third, it simplifies difficult submodules by setting $n+k$35 for $n+k$36 and $n+k$37, then zero-initializing those inserted blocks. Fourth, it does not describe a separate attention-head-level alignment, embedding fusion, or a full asymptotic analysis of OT creation cost (Cao et al., 11 Aug 2025).

Several specific technical details are absent from the paper. It does not print the explicit entropically regularized OT objective in formula form. It does not provide the Sinkhorn iteration equations. It does not give asymptotic time or memory complexity formulas for the OT alignment. It does not provide a fully explicit formula for how the two residual-path transport matrices are averaged at post-attention RMSNorm. These omissions do not prevent understanding the method at the architectural level, but they constrain exact reconstruction of the transport subroutine from the paper alone (Cao et al., 11 Aug 2025).

A broader methodological limitation is that the empirical evidence establishes gains over Avg-DeUS, LESA, LLaMA PRO, and SOLAR in the reported settings, but the margin over LESA is modest at the larger scale. The paper nevertheless treats the Avg-DeUS comparison as the decisive ablation because it isolates the effect of OT alignment itself. This suggests that the method’s core claim is not that interpolation is ineffective, but that unaligned interpolation leaves measurable performance on the table when neuron permutation mismatch is present (Cao et al., 11 Aug 2025).

Overall, OpT-DeUS is best understood as a progressive Transformer expansion method that recasts inserted-layer construction as a blockwise OT alignment-and-fusion problem. Its contribution is to make depth up-scaling in LLMs sensitive to permutation symmetry: adjacent layers are not merely copied or averaged, but first matched by a transport plan, propagated through the block structure by TMF, and then fused in a way that preserves initialization stability while improving downstream continual pre-training and fine-tuning results (Cao et al., 11 Aug 2025).