---
title: MSE-Optimal Layer Fusion
url: https://www.emergentmind.com/topics/mse-optimal-layer-fusion
type: topic
---

# MSE-Optimal Layer Fusion

MSE-optimal layer fusion refers to the family of algorithms, structural results, and analytic formulations that produce a single neural network layer (or set of layers) to minimize the mean squared error (MSE) between the network’s output—at a specific intermediate, embedding, or final layer—and the reference outputs of several networks, modules, or signal sources. Recent research has clarified both the structural optimality and the computational tractability of such fusion procedures for a variety of operations, architectures, and learning contexts including multi-task model merging, model compression, weight-space initialization, and multi-modal integration.

## 1. Formal Definitions and Mathematical Framework

The general problem of MSE-optimal layer fusion can be formulated as follows: Given a collection of $K$ networks or models $\{W_k\}_{k=1}^K$—each possibly trained on a separate (task, data, or modality)—and a prescribed layerwise decomposition, seek a new set of weights $\widehat{W}^{(\ell)}$ for each layer $\ell$ such that the aggregate mean-squared error of the transformed activations (or outputs) is minimized across models and input data distributions.

**Canonical objective:** For input $x$ and some feature function $f_k^{\ell}(W; x)$ (the representation at layer $\ell$ within expert $k$), define:

$$
\Delta f_k^{\ell} = f_k^{\ell}(W_{fused}; X_k^{\ell}) - f_k^{\ell}(W_{k}; X_k^{\ell})
$$

where $X_k^{\ell}$ is a matrix of calibration/example inputs at layer $\ell$. The total drift to be minimized is

$$
J^{\ell}(T^{\ell}) = \sum_{k=1}^K \left\| \Delta f_k^{\ell} \right\|_F^2 = \sum_{k=1}^K \|f_k^{\ell}(W_{fused}; X_k^{\ell}) - f_k^{\ell}(W_{k}; X_k^{\ell})\|_F^2
$$

This yields a convex, quadratic problem per layer, typically with closed-form solutions for standard layer types, as demonstrated in [2505.23859], [2001.10509], and [2508.16059].

## 2. MSE-Optimal Fusion in Model Merging and Knowledge Integration

**Layer-wise Optimal Task Vector Merging (LOT Merging):** For model merging, especially in multi-task consolidation, the LOT Merging strategy proceeds by minimizing feature drift at each layer:

- For **linear (matrix-multiply) layers:** The function is $f_k^{\ell}(W) = X_k^{\ell} W^{\ell}$, and the optimal update $T^{\ell *}$ to the task vector is

  $$
  T^{\ell *} = \Big(\sum_k (X_k^{\ell})^\top X_k^{\ell} \Big)^{\dagger} \sum_k (X_k^{\ell})^\top X_k^{\ell} T_k^{\ell}
  $$

  where $T_k^{\ell}$ are the task vectors from fine-tuned experts, and $\dagger$ denotes the Moore-Penrose pseudoinverse.

- For **elementwise scale (Hadamard) layers** (e.g., LayerNorm $\gamma$):

  $$
  T^{\ell *}[d] = \frac{ \sum_{k, x} x[d]^2 T_k^{\ell}[d] }{ \sum_{k, x} x[d]^2 }
  $$

- For **bias (shift) layers:**

  $$
  T^{\ell *} = \frac{1}{K} \sum_{k=1}^K T_k^{\ell}
  $$

This procedure, combining layerwise analytics and cheap matrix operations, enables efficient consolidation of fine-tuned models without retraining. Its formal convexity and solution structure are detailed in [2505.23859].

### Connection to Theoretical Optimality

LOT Merging is MSE-optimal under an additive-drift model, which assumes that layerwise mean squared drift dominates multi-task loss transfer and that the loss surface locally preserves Lipschitz continuity in the feature drift. The per-layer minimization guarantees the unique minimizer for the sum-squared-drift objective [2505.23859].

## 3. Fusion via Optimal Transport and Neuron/Filter Alignment

An alternative and influential approach involves neuron- or channel-alignment via optimal transport (OT) before parameter averaging [1910.05653]. This method is crucial in arbitrary-width networks and networks with permutation-symmetric layers.

- **Problem formulation:** At layer $\ell$, for each model $k$ with weight $W^{(k)}$, find soft assignment matrices $T^{(k)}$ (relaxing permutation matrices to the set of doubly stochastic matrices via entropic regularization). Then compute the MSE-optimal fused weights as

  $$
  \widehat{W}^{(\ell)} = \frac{1}{K} \sum_{k=1}^K W^{(k)} T^{(k)}
  $$

- **Sinkhorn Algorithm:** Soft alignments are found via the Sinkhorn-Knopp iterative algorithm over a cost matrix based on parameter distance or activation dissimilarity. Once the assignments are computed, the aligned weights are averaged.

This layerwise OT-based fusion is MSE-optimal among all matched/soft-aligned reparameterizations [1910.05653].

## 4. Fusion in Multi-Modal, Multi-Embedding, and Multi-Layer Settings

MSE-optimal fusion extends beyond classical parameter merging to scenarios involving modality integration or structural compression.

**Multi-layer Steerable Embedding Fusion (MSEF):** In time-series forecasting with LLMs, MSEF injects layer-specific “steering vectors” and time-series embeddings at every Transformer block, constructing the per-layer fused input by concatenation:

$$
F^{(\ell)} = [ s^{(\ell)} ; e_{TS}^{(\ell)} ; h^{(\ell-1)} ]
$$

Only these layerwise steering vectors and the terminal output layer are trained, with all backbone weights frozen. The objective is MSE on predicted series, and repeated multi-layer fusion is shown to dramatically improve preservation of time-series information throughout the model [2508.16059].

## 5. MSE-Optimal Layer Fusion for Network Compression and Initialization

The layer-fusion framework also appears in network shrinking, where two adjacent layers are replaced by a single layer mapping, minimizing MSE over a finite dataset [2001.10509]:

- **Dense $\to$ Dense:** Fuse two FC layers $H_1(x)$ then $W_2$, $b_2$ into a single FC layer, with

  $$
  W^* = W_2 C_{a_1,x} C_x^{-1}
  $$
  $$
  b^* = W_2 \mu_1 + b_2 - W^* \mu_0
  $$

  where $C_{a_1,x}$ is the covariance between post- and pre-activation features, and $\mu_0, \mu_1$ are empirical means.

- **Convolutional variants:** Analogous closed forms exist for conv$\to$conv or conv$\to$dense, using cross-covariance contractions and channel-wise Toeplitz matrices.

When used as an initialization (e.g., FuseInit), this approach yields shallower networks that attain or surpass the performance of their deeper teachers after modest fine-tuning [2001.10509].

## 6. Empirical Performance and Computational Considerations

Across vision (ViT-B/32, ViT-L/14), vision-language (BLIP), and time-series tasks, MSE-optimal layer fusion substantially improves performance versus previous non-training or naive averaging baselines:

- **LOT Merging:** On ViT-B/32, average accuracy $82.7\%$ versus $78.3\%$ for the best baseline; on ViT-L/14, $90.5\%$ versus $89.6\%$ [2505.23859].
- **OT Fusion:** On CIFAR10 with VGG11, OT fusion achieves $86.0\%$ accuracy one-shot, compared to $17.0\%$ for vanilla averaging and $91.34\%$ for ensembles [1910.05653].
- **MSEF:** Mean MSE reduction of $31.8\%$ versus the best-external TSF model on long-term forecasting [2508.16059].

These methods are also computationally efficient—e.g., for LOT merging, merging eight ViT-B/32 task vectors takes 44 seconds on a single RTX 3090 [2505.23859].

## 7. Limitations and Extensions

While the theoretical guarantees of MSE-optimal layer fusion are rigorous under specific data/model assumptions (convexity, channel independence, data representativity), practical limitations include:

- **Representational mismatch:** Fusion is optimal for the empirical, calibration-set MSE, not necessarily for downstream task loss.
- **Covariance estimation:** Some methods require well-conditioned covariance matrices, enforceable via regularization.
- **Architectural restrictions:** Extension to highly non-linear layers or non-standard modules may be non-trivial.
- **Dependence on calibration/sample diversity:** The quality of fusion is sensitive to how representative the calibration set or support batch is for the data domain.

Nonetheless, these layerwise, MSE-optimal procedures constitute a principled, scalable, and theoretically sound toolkit for multi-model, multi-task, and multi-modal integration, with extensive demonstration on state-of-the-art deep architectures [2505.23859], [1910.05653], [2508.16059], [2001.10509].

Source: https://www.emergentmind.com/topics/mse-optimal-layer-fusion