PaperlayerView exact revision on arXiv

Source-linked AI summary

REPAIR: REnormalizing Permuted Activations for Interpolation Repair

Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, Behnam Neyshabur

arXiv:2211.08403v3cs.LGcs.AIcs.CVstat.ML

TL;DR

The paper examines why neuron alignment has not been sufficient to establish low-barrier interpolation between independently trained networks, despite the conjecture that permutation invariance may help. It proposes REPAIR, which rescales interpolated networks to mitigate variance collapse, and reports barrier reductions across architectures and datasets. These results support low-barrier permuted interpolation in additional settings, while the practical correction moves beyond strictly permuted linear interpolation.

  • Problem

    The paper investigates why neuron alignment alone has not established low-barrier interpolation, in the context of the conjecture that accounting for permutation invariance can yield such paths between SGD solutions.

    Source: §1 · pp. 2–3

  • Method

    REPAIR mitigates variance collapse by rescaling hidden units in interpolated networks so their statistics match those of the original networks.

    Source: §1 · p. 3

  • Results

    REPAIR reduces barriers across a wide variety of architectures and datasets, including reducing the ResNet18-on-CIFAR-10 barrier from 16% to 1.5%.

    Source: §1 · p. 3, §6 · p. 12

  • Takeaways & Limitations

    REPAIRed interpolations between networks trained on disjoint, biased dataset splits can outperform either endpoint network.

    Source: §6 · p. 12

    Low-barrier permuted interpolation with REPAIR requires rescaling preactivations, so it is not strictly permuted linear interpolation.

    Source: §6 · p. 12

Abstract

from arXiv · show

In this paper we look into the conjecture of Entezari et al. (2021) which states that if the permutation invariance of neural networks is taken into account, then there is likely no loss barrier to the linear interpolation between SGD solutions. First, we observe that neuron alignment methods alone are insufficient to establish low-barrier linear connectivity between SGD solutions due to a phenomenon we call variance collapse: interpolated deep networks suffer a collapse in the variance of their activations, causing poor performance. Next, we propose REPAIR (REnormalizing Permuted Activations for Interpolation Repair) which mitigates variance collapse by rescaling the preactivations of such interpolated networks. We explore the interaction between our method and the choice of normalization layer, network width, and depth, and demonstrate that using REPAIR on top of neuron alignment methods leads to 60%-100% relative barrier reduction across a wide variety of architecture families and tasks. In particular, we report a 74% barrier reduction for ResNet50 on ImageNet and 90% barrier reduction for ResNet18 on CIFAR10.

1 Introduction

The paper asks why aligning neurons has not been enough to make independently trained networks linearly interpolable, and identifies activation-variance collapse as a key obstacle. It proposes REPAIR to correct that collapse and reduce interpolation barriers.

1 Introduction

  • REPAIR reduces interpolation barriers across varied architectures, datasets, normalization layers, widths, and depths, including 74% for aligned ImageNet ResNet50s (76% to 20%).The figure compares interpolated networks before and after REPAIR; larger markers indicate wider networks.Source: Figure 1 · p. 2, §1 · p. 3
  • The work investigates Entezari et al. (2021)'s conjecture that accounting for permutation invariance can yield low-barrier linear paths between SGD solutions.The paper notes that reducing interpolation obstacles may matter for applications including model ensembling and model merging.Source: §1 · p. 2
  • Neuron alignment alone has not established low-barrier interpolation for standard deep convolutional networks.Prior work focused on improving alignment algorithms, including optimal-transport alignment and LayerNorm-based approaches.Source: §1 · p. 2
  • Aligned interpolated networks suffer variance collapse: hidden-unit activation variance becomes much smaller than in the original networks.The paper identifies this as a source of poor performance in interpolated networks.Source: §1 · p. 3
  • REPAIR mitigates variance collapse by rescaling hidden units so their statistics match those of the original networks.REPAIR stands for REnormalizing Permuted Activations for Interpolation Repair.Source: §1 · p. 3

2 Preliminaries

The preliminaries define linear interpolation and its loss barrier, explain why neuron permutations matter, and introduce alignment methods used to pair hidden units.

2.1 Linear interpolation of neural networks

  • The loss barrier is the maximum increase in loss along the linear path relative to the convex combination of endpoint losses.The barrier between independent, unmodified SGD solutions is typically large.Source: §2.1 · p. 3

2.2 Permutation invariance

  • For typical architectures, permuting neurons within a layer can preserve the represented function, so corresponding units in two networks may nevertheless be misaligned.Entezari et al. (2021) conjecture that accounting for this symmetry can yield small interpolation barriers for sufficiently wide networks.Source: §2.2 · p. 3
  • Figure 2 depicts activation variance declining with depth in averaged networks, with REPAIR intended to restore internal statistics toward endpoint levels.“Before REPAIR” denotes interpolation after aligning endpoint hidden units but before applying the correction.Source: Figure 2 · p. 4

2.3 Neuron alignment algorithms

  • Alignment methods pair hidden units using activation correlations, assignment optimization, Hessian approximations, or soft alignment; this paper uses the method introduced by Li et al. (2015).Residual networks require consistent permutations across layers feeding a shared residual stream.Source: §2.3 · p. 4

3 Variance collapse in interpolated networks

Aligned networks can still have high interpolation barriers, especially as depth increases; the paper investigates whether declining activation variance explains this pattern. It finds that variance collapse worsens through deep networks and gives a statistical account of the reduction.

3 Variance collapse in interpolated networks

  • Aligned-network barriers rise sharply with depth, and even strong optimal-transport matching is insufficient for low-barrier connectivity in standard ResNets.The paper investigates whether the depth-related increase observed by Entezari et al. (2021) reflects the same phenomenon behind high barriers in standard ResNets.Source: §3 · p. 5

3.1 Identifying the problem: Variance Collapse

  • In deep MLPs, interpolated hidden-unit activation variance progressively decays until later layers are nearly constant; final-layer variance also falls by more than 10× in VGG and ResNet50 experiments.The VGG and ResNet50 experiments use CIFAR-10 and ImageNet, respectively.Source: §3.1 · p. 5

3.2 Why does this phenomenon occur?

  • For a first-layer unit, interpolation combines endpoint preactivations linearly, and typical aligned-unit correlation of ≈0.4 yields midpoint variance of 0.7 times endpoint variance.The paper reports this as a 30% reduction relative to the endpoint networks.Source: §3.2 · pp. 5–6
  • The authors expect variance loss to compound with depth, consistent with experiments, but state that their statistical analysis cannot be rigorously extended to deeper layers.The observed worsening spans MLP, VGG, and ResNet50 networks.Source: §3.2 · p. 6

4 REPAIR

REPAIR addresses variance collapse in aligned weight interpolations by restoring selected units’ statistics. This section develops approximate and forward-pass exact variants and demonstrates results across MLPs, ResNets, normalization choices, and scales.

4 REPAIR

  • REPAIR applies affine rescaling and shifting to selected interpolated channels to restore their statistics toward endpoint values and counter variance collapse.For ResNets, selected outputs include convolutional preactivations and residual blocks.Source: §4 · p. 6

4.1 Closed-form approximate variant

  • The approximate variant uses endpoint statistics to compute rescaling coefficients, exact for first-layer units but approximate for deeper layers.On aligned MNIST MLPs, it enables over 90% accuracy through 27 hidden layers, versus a six-layer limit without correction.Source: §4.1 · p. 7

4.2 Forward-pass exact variant

  • At ResNet20 midpoints, REPAIR raises BatchNorm and LayerNorm accuracy to 92.0%, from 34.7% and 88.4%, respectively.LayerNorm networks uniquely reach relatively low barriers before REPAIR; the figure also compares FixUp and widths from 1 to 16.Source: Figure 4 · p. 8
  • The forward-pass exact variant uses added BatchNorm layers to match selected channels’ mean and standard deviation to endpoint-interpolated statistics.About 5,000 training examples estimate running statistics; the added layers are then set to evaluation mode.Source: §4.2 · p. 8
  • 56.5% ImageNet midpoint accuracy after full REPAIR (from less than 1% uncorrected); CIFAR-10 ResNet18 reaches 94.1% (from 66.0%).Figure 3 also reports the CIFAR-10 barrier reduced to 1.5%.Source: Figure 3 · p. 7, §4.2 · p. 8, Figure 5 · p. 9

5 Experimental results

Across the tested architectures and datasets, REPAIR reduces interpolation barriers, including on ImageNet and in a split-data merge. Its effects vary with normalization, width, and depth.

5.1 Normalization layer

  • At 16× width, REPAIR reduces BatchNorm- and LayerNorm-ResNet20 barriers below 0.05; at 8×, midpoint accuracy rises from 90.3% to 93.2% for LayerNorm and 32.3% to 94.4% for BatchNorm.LayerNorm-ResNet20 endpoints achieve 92.4% CIFAR-10 accuracy versus 93.6% for BatchNorm.Source: §5.1 · pp. 9–10
  • LayerNorm's test-time normalization may prevent variance reductions from compounding; the authors say this explains prior connectivity differences, at the cost of lower performance on most datasets.They contrast Singh and Jaggi (2020), who found a barrier for standard ResNets, with Ainsworth et al. (2022), who used LayerNorm-based ResNets.Source: §5.1 · p. 10

5.2 ImageNet

  • On ImageNet, REPAIR raises midpoint accuracy from below 1% to 41.1% (ResNet18), 56.5% (ResNet50), and 64.2% (double-width ResNet50).Even after correction, the ResNet50 test-error barrier is 20.1%; the double-width model's is 12.9%.Source: §5.2 · p. 10
  • Adding residual-block outputs to ResNet50 REPAIR raises midpoint accuracy from 53.2% to 56.5% and reduces the barrier from 23.4% to 20.1%.Applying REPAIR only to convolutional outputs is mathematically equivalent to resetting BatchNorm statistics for architectures with BatchNorm after each convolution.Source: §5.2 · p. 10

5.3 Network width and depth

  • Across four datasets, REPAIR reduces barriers for MLPs and ResNet20s, with barriers decreasing as width grows; VGG barriers rise with depth without REPAIR but remain small and roughly flat with it.For MNIST, even narrower networks reach nearly zero barrier.Source: §5.3 · pp. 10–11, Figure 6 · p. 11

5.4 Split data training

  • At α = 0.3, REPAIR reaches 1.30 CIFAR-100 test loss versus Ainsworth et al.'s 1.73; authors attribute the improvement partly to REPAIR and stronger standard ResNets.The REPAIRed interpolations also outperform both endpoints on the combined test set.Source: Figure 7 · p. 11, §5.4 · pp. 11–12
  • The split-data experiment trains CIFAR-100 models on disjoint, class-biased halves; REPAIRed interpolations outperform either endpoint on the combined test set.Each endpoint is more accurate on a different half of the classes.Source: §5.4 · p. 11, Figure 7 · p. 11

6 Discussion and future work

The discussion presents REPAIR as a correction for variance collapse that improves interpolated-network performance across varied settings. The authors qualify its support for permutation-based connectivity: practical low barriers require rescaling in some settings.

6 Discussion and future work

  • REPAIR mitigates variance collapse and improves interpolated-network performance across architectures and datasets, including ResNet18 on CIFAR-10 and ResNet50 on ImageNet.The reported barriers fall from 16% to 1.5% for CIFAR-10 ResNet18, while ImageNet ResNet50 midpoint accuracy rises from below 1% to 56.5%.Source: §6 · p. 12
  • These results support Entezari et al. (2021)'s conjecture in further scenarios, but REPAIR's preactivation rescaling means the corrected paths are not strictly permuted linear interpolations.The authors say such correction appears necessary at practical widths without LayerNorm.Source: §6 · p. 12
  • REPAIRed interpolations between networks trained on disjoint, biased dataset splits can outperform either endpoint.The authors suggest that further applications may include weight-space ensembling, checkpoint averaging, and robust finetuning.Source: §6 · p. 12

A Further discussion of related work

The related-work discussion describes alignment approaches based on neuron similarity, functional difference, and optimal transport, while noting that prior methods did not consistently yield low barriers.

A Further discussion of related work

B.1 Training hyperparameters

The training appendix summarizes shared optimizer settings and notes which hyperparameters remain fixed across experimental variations.

B.1 Training hyperparameters

  • All networks are trained using SGD with momentum 0.9; weight decay and learning rates vary by task.Task-specific values are listed in the hyperparameter table.Source: §B.1 · p. 16
  • MLP hyperparameters are held constant across widths, depths, and datasets.Source: §B.1 · p. 16
  • ResNet20 hyperparameters remain constant across width, dataset, and normalization choices, except that MNIST and SVHN training omit cutout and horizontal flip.Source: §B.1 · p. 16

B.2 PyTorch wrapper module pseudocode

The PyTorch wrappers collect activation statistics from endpoint networks and use them to set interpolated-network normalization statistics. The procedure applies to selected modules in the interpolated network.

B.2 PyTorch wrapper module pseudocode

  • TrackLayer wrappers collect endpoint-module output means and standard deviations after about 5,000 training examples pass through the tracked networks.The wrappers update running statistics while discarding the BatchNorm output.Source: §B.2 · p. 17
  • ResetLayer wraps each selected interpolated module with BatchNorm, setting its bias to the goal mean and its weight to the goal standard deviation.Its forward pass applies BatchNorm to the wrapped layer's output.Source: §B.2 · p. 17
  • For each corresponding endpoint and interpolated module triple, the ResetLayer statistics are set to the interpolation of the two TrackLayer statistics.The paper notes that approximately 5,000 training examples suffice to retrieve statistics from tracked endpoint modules.Source: §B.2 · p. 17

C Additional plots

The additional plots examine REPAIR on split CIFAR-100 data and without neuron alignment, and provide loss-based versions of earlier analyses. They also visualize how activation-correlation alignment permutes first-layer ResNet50 filters.

C Additional plots

  • The first-layer ResNet50 visualization shows activation-correlation permutation yields some nearly identical filter pairs, while other filters lack close matches.The figure contrasts model A's filters, permuted model B filters, and model B's original filter order.Source: Figure 9 · p. 19
  • Figures 10–12 cover normalization-layer effects, interpolated ResNet barriers, and network width/depth; test-loss barriers appear in Figure 10’s right panel, Figure 11, and Figure 12, while Figure 10’s left uses accuracy.Figure 11 reports interpolated ResNets on ImageNet and CIFAR-10.Source: Figure 10 · p. 19, Figure 11 · p. 19, Figure 12 · p. 20
  • REPAIRed interpolation yields checkpoints beating either endpoint on full CIFAR-100 test loss and accuracy, but corrected endpoints outperform its best gain.The experiment trains two networks on disjoint, biased subsets and compares interpolations with logit mixing and distribution-shift-corrected endpoints.Source: Figure 13 · p. 20
  • Without alignment, REPAIR helps near endpoints, while alignment is essential near midpoint; both are needed for a high-performing midpoint.The figure covers ResNet18, VGG11, and ResNet50, applying REPAIR to convolutional outputs and, for ResNet50, residual-block outputs.Source: Figure 14 · p. 21
Loading 2211.08403v3…