Technical Summary: Partial Fusion of Neural Networks

Problem Statement

Neural network ensembles typically outperform individual models by combining their outputs, offering increased robustness and prediction performance. However, ensembles incur significant computational costs, requiring multiple models to be stored in memory and increasing inference time. Conversely, weight aggregation methods (such as weight averaging) produce single, cost-efficient models but often suffer from reduced accuracy, particularly when the networks being merged have not been aligned or were trained on different data distributions. Existing weight aggregation techniques often rely on strict permutation invariance (matching layers of equal size) or require complex alignment procedures that may not generalize to heterogeneous architectures or uneven layer sizes.

The core challenge addressed is finding an efficient tradeoff between the high performance of ensembles and the low computational cost of weight aggregation. Current methods tend to be binary: either fully aggregate weights (potentially losing performance) or maintain a full ensemble (incurring high costs). There is a need for a flexible mechanism to interpolate between these extremes.

Methodology

The paper introduces Partial Fusion, a method that interpolates between weight aggregation and ensembles by merging only the most similar neurons across networks while keeping the remaining neurons isolated. This approach allows for a flexible, layer-wise tradeoff between model size and performance.

1. Partial Fusion via Partial Optimal Transport

The primary method proposed extends existing weight aggregation techniques based on neuron-level similarity.

• Concept: Instead of aligning all neurons between two networks (as in standard Optimal Transport (OT) Fusion), partial fusion identifies and matches only the most similar neurons. The unmatched (most dissimilar) neurons are left isolated, preserving their original weights.
• Mechanism: The authors utilize Partial Optimal Transport (Partial OT). Given two networks $A$ and $B$, a coupling $\tilde{\pi}$ is computed that matches only a fraction $(1-\alpha)$ of the mass between the neurons of the two networks, where $\alpha \in [0, 1]$ is a parameter controlling the degree of fusion.
  • $\alpha = 0$: Corresponds to full weight aggregation (all neurons matched).
  • $\alpha = 1$: Corresponds to an ensemble (no neurons matched).
• Implementation: The method encodes layers as discrete probability distributions based on neuron features (e.g., activation vectors or weight matrices). It solves a partial OT problem to find the optimal matching of the most similar neurons. The resulting fused network $C$ contains three components per layer: isolated neurons from $A$, fused neurons (averaged weights), and isolated neurons from $B$.
• Global Optimization: To improve alignment, the method incorporates a fixed-point iteration approach (inspired by Ainsworth et al., 2023) that optimizes permutations/matchings globally across layers rather than layer-by-layer, accounting for dependencies between layers.

2. Generalized Pruning

The paper frames partial fusion from a complementary perspective: Generalized Pruning of ensemble models.

• Concept: Instead of starting with weight aggregation and relaxing it, this view starts with an ensemble (a large concatenated model) and applies a pruning operation. Unlike standard pruning which only deletes neurons, generalized pruning allows neurons to be linearly combined.
• Formalism: This is formalized as transforming the weights of a large ensemble model $E$ into a smaller network $S$ via linear transformations (stochastic kernels).
• Clustering Approach: A specific realization of generalized pruning is proposed using clustering. Neurons from the ensemble are clustered, and neurons within a cluster are linearly combined. This offers more degrees of freedom than partial fusion, as it does not restrict combinations to pairs from different networks.
• Inductive Bias: The authors note that partial fusion imposes a specific inductive bias: it only linearly combines neurons if they originate from different networks and limits combinations to pairs. This bias is hypothesized to be beneficial for Convolutional Neural Networks (CNNs), whereas the more flexible clustering approach may be better suited for Multi-Layer Perceptrons (MLPs).

Key Contributions

1. Partial Fusion Framework: The introduction of a method to interpolate between weight aggregation and ensembles by merging only the most similar neurons, providing a tunable tradeoff between computational cost and performance.
2. Generalized Pruning Perspective: The theoretical unification of partial fusion and model compression under the concept of generalized pruning, where neurons can be linearly combined rather than just deleted. This extends the applicability of the method to single-model compression.
3. Concrete Algorithms:
   • Partial OT Fusion: An extension of OT-based model fusion using partial optimal transport to handle unequal matching and layer sizes.
   • Clustering-based Pruning: A method for generalized pruning using hierarchical clustering to linearly combine neurons, offering flexibility in model compression.
4. Empirical Validation: Experiments on MLPs and CNNs (VGG11, ResNet18) across various datasets (MNIST, CIFAR10) demonstrating that partial fusion achieves performance closer to ensembles with significantly fewer parameters.

Results

The experiments demonstrate the efficacy of the proposed methods in several regimes:

• Split Dataset Aggregation (Heterogeneous Data): When merging models trained on different data splits (e.g., one specialized, one general), Partial OT Fusion significantly outperforms standard OT fusion. Increasing the number of isolated neurons (increasing $\alpha$) leads to monotonic performance improvements, bridging the gap between the aggregated model and the ensemble. For instance, a 20-40% increase in neuron count yielded clear accuracy gains over standard fusion.
• Fine-Tuning: Partially fused models, when fine-tuned with a small subset of data (1% for MNIST, 5% for CIFAR10), consistently surpass both individual models and the full ensemble in accuracy.
• Same Data Aggregation: When merging models trained on the same data with different seeds, full weight aggregation often performs poorly for CNNs (like VGG11) but better for ResNets due to residual connections. Partial fusion provides a robust interpolation, recovering performance lost in full fusion.
• Single Model Compression (Generalized Pruning):
  • For MLPs, generalized pruning via clustering significantly outperforms both unstructured pruning and pruning with post-processing (OT Fusion). The authors attribute this to clustering's ability to balance the errors of "deletion" (pruning) and "blurring" (linear combination).
  • For CNNs (VGG11), Partial OT Fusion often outperforms clustering, suggesting the inductive bias of matching specific pairs of neurons is more effective for convolutional layers.
• Parameter Efficiency: The methods allow for model sizes between that of a single network and a full ensemble. For example, a partially fused model with $\alpha=0.5$ uses approximately 1.45$\times$ the parameters of a single network (compared to 2$\times$ for an ensemble) while achieving accuracy much closer to the ensemble.

Significance and Claims

The paper claims that the primary insight is the significance of differences in neuron-level similarities. Neurons vary in their uniqueness; some are highly similar across networks, while others play unique roles. The proposed methods exploit these differences:

• Partial Fusion leverages the fact that keeping the most dissimilar neurons isolated increases the average similarity of the remaining aggregated neurons, leading to more effective weight averaging.
• Generalized Pruning exploits the tradeoff between deleting neurons (losing information) and linearly combining them (blurring information).

The authors position their work as a step toward a structured study of merging operations in the weight space of neural networks. They emphasize that the methods provide a flexible, efficient tradeoff between model size and performance, often yielding disproportionate performance gains for small increases in model size. The paper remains modest regarding future implications, noting that scaling to larger models and developing principled criteria for layer-wise decisions (e.g., which layers to fuse vs. isolate) are open directions for future work. The code is made available to facilitate further research.