anthology-of-the-sota

LIT-333: Git Re-Basin: Merging Models modulo Permutation Symmetries

3 min read

Properties5
StatusActive
Version2
Tagsmodel-stability
Cited as “paper” byNOTE-116 — Git Re-Basin: Merging Models modulo Permutation Symmetries
Cited as “source” bySOTA-217 — Align the hidden-unit permutation before averaging weights from separately trained networks, THEORY-010 — Most of the loss barrier between two independently trained networks is permutation, not disagreement

Graph View

Ainsworth et al. (2022) — ARXIV-2209.04836

Key takeaways

Develops three algorithms for finding the permutation that aligns two independently trained networks before merging their weights (§3):

  1. Activation matching. Per layer, a linear assignment problem on the cross-product of the two models’ activations over the training data. It needs a pass over the data, though “a full run through the training dataset may be unnecessary”.
  2. Weight matching (Algorithm 1). Coordinate descent over all layers. Each step solves one layer’s permutation as a linear assignment problem, holding its neighbours fixed. It uses no data.
  3. Straight-through estimator (Algorithm 2). Learns the permutation by gradient descent on the loss at the midpoint, projecting with Algorithm 1 in the forward pass.

STE is best but “very computationally expensive”. Weight and activation matching “perform similarly, although weight matching is orders of magnitude faster and does not rely on the input data distribution” (Figure 2).

  • Empirical — Loss barriers after matching. MNIST MLP: zero barrier with all three methods, and after weight and STE matching the test-loss interpolation is convex (§5.1). CIFAR-10: barriers after weight matching fall with width for VGG-16 and ResNet20 and reach zero for wide ResNet20. The paper calls this the first zero-barrier LMC between independently trained ResNets (§5.3, Figure 4). ImageNet ResNet50 (1×): no zero barrier, but a 67% decrease relative to naive interpolation (§5.1). Holds when: sufficient width. “1×-sized models did not seem to exhibit linear mode connectivity” (§5.3). LMC also emerges only gradually during training, and is absent early (§5.2, Figure 3).
  • Empirical — Merged model quality. ResNet50 on ImageNet, merged by weight matching with BatchNorm statistics recomputed: 51.01% top-1. OT-Fusion gets 1.38%, “only marginally improving over naïve averaging” (A.7.1, Figure 8). Two ResNet20s trained on disjoint, biased halves of CIFAR-100 merge into a model with lower test loss than either input, but not higher accuracy. It falls short of an ensemble or of full-data training, and is better calibrated (§5.4, Figures 5, 10, 11). Holds when: BatchNorm statistics recomputed after interpolation (A.4). The paper lists failure modes in A.1.

What the evidence does not cover

  • Weight matching solves an NP-hard problem approximately (Lemma 1, the “sum of bilinear assignments problem”). Algorithm 1 terminates (Lemma 2) but is not guaranteed to find the best permutation. A failure to find LMC “cannot rule out the existence of a satisfactory permutation” (§5).
  • Thin models: zero barrier was not reached at 1× width on CIFAR-10, or on ImageNet. The paper leaves open whether its methods fail on thin models or some invariance other than permutation is at play (§5.3, §7).
  • Known failure modes (A.1): insufficient width, early training, VGGs on MNIST, MNIST MLPs at badly chosen learning rates, and ConvNeXt (few permutation symmetries). BatchNorm needs its statistics recomputed after merging. GroupNorm is not permutation-invariant and was not tested (A.4).
  • Multi-model merging (Algorithm 3, MergeMany, A.10) is demonstrated only on MNIST MLPs, up to 32 models.
  • Vision only; no transformer is merged.

Standing in the anthology

Read — the reading is NOTE-116. Arrived in the imported batch, which brought in the permutation-symmetry line, which is why two independently trained networks cannot simply be averaged.

Backlinks