PaperlayerView exact revision on arXiv

Source-linked AI summary

Revisiting Model Stitching to Compare Neural Representations

Yamini Bansal, Preetum Nakkiran, Boaz Barak

arXiv:2106.07682v1cs.LGstat.ML

TL;DR

The paper addresses limited understanding of how neural representations relate and whether existing similarity metrics capture their compatibility. It revisits model stitching, connecting frozen networks through a simple trainable layer, and finds compatibility across varied training conditions, plug-in gains from stronger representations, and stitching connectivity among typical SGD minima.

  • Problem

    Neural networks learn rich intermediate representations, but their properties and relationships—and which training factors shape them—remain poorly understood, while similarity metrics such as CKA provide limited operational interpretation.

  • Method

    The paper studies representations by stitching a candidate bottom representation to a network’s top layers through a simple, low-capacity trainable alignment layer.

  • Results

    Model stitching finds compatible representations across different initializations, data subsets, tasks, widths, and training durations, including supervised versus self-supervised models, while stronger representations can improve weaker models and typical SGD minima exhibit stitching connectivity.

  • Takeaways & Limitations

    Model stitching supplies quantitative, operational evidence about neural representations that can differ qualitatively from conclusions based on CKA and related similarity measures.

  • Takeaways & Limitations

    Stitching requires training an additional network and can be difficult across architectures because it requires a careful choice of stitching family.

Abstract

from arXiv · show

We revisit and extend model stitching (Lenc & Vedaldi 2015) as a methodology to study the internal representations of neural networks. Given two trained and frozen models $A$ and $B$, we consider a "stitched model'' formed by connecting the bottom-layers of $A$ to the top-layers of $B$, with a simple trainable layer between them. We argue that model stitching is a powerful and perhaps under-appreciated tool, which reveals aspects of representations that measures such as centered kernel alignment (CKA) cannot. Through extensive experiments, we use model stitching to obtain quantitative verifications for intuitive statements such as "good networks learn similar representations'', by demonstrating that good networks of the same architecture, but trained in very different ways (e.g.: supervised vs. self-supervised learning), can be stitched to each other without drop in performance. We also give evidence for the intuition that "more is better'' by showing that representations learnt with (1) more data, (2) bigger width, or (3) more training time can be "plugged in'' to weaker models to improve performance. Finally, our experiments reveal a new structural property of SGD which we call "stitching connectivity'', akin to mode-connectivity: typical minima reached by SGD can all be stitched to each other with minimal change in accuracy.

1 Introduction

The paper asks whether successful neural networks learn compatible internal representations and revisits model stitching as a quantitative way to test this question. Experiments examine compatibility across training conditions, representation quality, and SGD solutions, while contrasting stitching with similarity metrics such as CKA.

  • Research question: The paper frames a central uncertainty: successful models may learn either incompatible internals or roughly similar representations despite differing training conditions.These alternatives are presented as the “snowflakes” and “Anna Karenina” scenarios.
  • Experimental predictions: Networks trained with different initializations, datasets, tasks, widths, or training durations can be compared by testing whether their representations transfer through stitching.The experiments target both “all roads lead to Rome” and “more is better” predictions.
  • Method: Model stitching connects bottom layers from one network to top layers of another through a simple, low-capacity trainable alignment layer.The stitching layer is intended to align representations rather than improve the model independently.
  • Paper scope: The paper’s scope includes comparisons with prior representation measures, formal stitching-connectivity experiments, and empirical tests of cross-training compatibility and representation improvement.The introduction organizes these results across Sections 3–6.
  • Why stitching: The method provides operational, quantitative evidence about representation compatibility and can yield conclusions that differ qualitatively from CKA and other similarity measures.Stitching penalties inherit the meaningful units of the task loss, unlike abstract similarity scores.
  • Main findings: The experiments report stitching connectivity among typical SGD minima and improved performance when stronger representations trained with more data, larger width, or more epochs are plugged into weaker models.The paper presents these findings as quantitative evidence for practitioner intuitions about similar and improving representations.

2 Model Stitching

Model stitching evaluates whether a candidate representation can replace early layers of a network through a simple trainable alignment layer. Its low-capacity design distinguishes representation compatibility from gains produced by flexible learning or full fine-tuning.

  • Operational definition: Model stitching replaces the first ℓ layers of network A with representation r, using a simple stitching layer before A’s remaining layers.The stitching family can include linear 1 × 1 convolutions for convolutional networks.
  • Operational definition: The stitching objective minimizes the task loss after composing A’s top layers with a stitcher and r.The stitching penalty compares this optimized loss with A’s original loss.
  • Evaluation: In practice, the stitching layer is optimized on the training set, while the stitching penalty is estimated on the test set.For fully connected networks, linear stitching tests whether r can be transformed into A’s relevant representation subspace.
  • Design constraint: The stitching family should remain simple so it aligns representations rather than learning an arbitrary transformation between them.The authors use an untrained randomly initialized network as a check against stitching layers learning useful representations.
  • Interpretation: Model stitching occupies a middle ground between direct plugging-in and full fine-tuning, while enabling early-layer analysis with powerful nonlinear decoders unlike linear probes.Linear probe accuracy can make an early layer from a higher-quality network appear worse than a later layer from a weaker network.

3 Stitching vs. representational similarity

The paper contrasts model stitching with representation-similarity measures such as CKA and CCA-based methods. Stitching evaluates compatibility and task-relevant quality in operational terms, including asymmetric cases where one representation improves another model.

  • Similarity measures: Prior representation studies use CCA, SVCCA, PWCCA, and CKA to compare neural representations.CKA was proposed to address several issues associated with CCA.
  • Similarity measures: CKA computes a covariance-based similarity between two representation functions using the test distribution.Its linear form normalizes the squared Frobenius norm of cross-covariance by the covariance norms of both representations.
  • Comparison: Table 1 compares CKA closeness, indicated by CKA ≈1, with stitching closeness, indicated by error ≈0%.“Varies” denotes no consistent conclusion across architectures and tasks.
  • Comparison: Stitching can reveal that one representation is better than another, whereas CKA can only report that the two representations are far apart in such cases.The comparison reflects stitching’s asymmetric, task-based evaluation rather than a symmetric distance.
  • Task relevance: Similarity measures may be affected by useless or random coordinates, while stitching focuses on whether the representation supports the downstream task.Appending 1000 useless Gaussian coordinates to a representation can reduce CKA without establishing a corresponding task-performance change.
  • Interpretability: Stitching penalties use the loss’s meaningful units, giving a 3% CIFAR-10 accuracy penalty the operational meaning of at most a 3% accuracy decrease after replacement.This makes the result easier to interpret than a numerical difference between CKA values such as 0.9 and 0.8.
  • Invariance: The invariances of similarity measures are not always natural for representation quality, whereas the stitching family can be selected to encode desired invariances.The paper notes that CCA and CKA are invariant to orthogonal transformations, including transformations such as pixel permutations.

4 Stitching Connectivity

Stitching connectivity describes a low-penalty sequence of layerwise stitched models between independently trained networks. The paper reports this property for identically distributed networks and presents it as a structural property of typical SGD solutions.

  • Empirical property: Two independently initialized networks with the same architecture and data distribution can be stitched at all layers with close to 0 penalty.The authors call this empirical property stitching connectivity.
  • Construction: For networks A and B with L layers, S_i replaces A’s first i layers with B’s corresponding layers and optimizes the stitching layer.S_0 equals A and S_L equals B, producing a layerwise sequence between the endpoints.
  • Definition: A and B are stitching-connected when every intermediate stitched model has test loss comparable to A.Networks with very different internal representations would fail this criterion.
  • Claim: The paper’s main claim is that, for a fixed data distribution, almost all minima reached by SGD are stitching-connected to one another.This is stated as an informal conjecture for independently seeded, identically trained networks.
  • Relation to mode connectivity: Stitching connectivity complements mode connectivity by giving a discrete, layer-interpretable sequence rather than a parameter-space path.The sequence contains as many steps as layers and can connect networks trained on different tasks.
  • Interpretation: The property is especially relevant to overparameterized models, where exponentially many global training minima could otherwise implement classification differently.The authors report that SGD is empirically biased toward minima with essentially the same internal representations.

5 All Roads Lead to Rome

Model stitching tests whether networks trained with different methods or label quality learn compatible representations. Across supervised and self-supervised training, and even weak labels, early representations remain similar while later layers can differ.

  • Comparing self-supervised and supervised methods: Self-supervised and supervised ResNet-50 models trained on ImageNet are stitching connected at all layers despite substantially different training procedures.The compared methods generally achieve 75% ± 1 test accuracy, except SimCLR at 68.8%.
  • Comparing self-supervised and supervised methods: CKA ranges from 0.35 −0.9 across layers for SimCLR and supervised ResNet-18 models, while the maximum stitching penalty is 3%.
  • Changing the label distribution p(y|x): Networks trained with coarse labels or 10% and 50% label noise have first-half layers as good as a standard CIFAR-10 network when stitched.The exception is the network trained with 100% noisy labels.
  • Changing the label distribution p(y|x): Early layers learn general-purpose features such as curve detectors that can support a large variety of vision tasks.The passage presents this as consistent with prior work on vision.

6 More is Better

Model stitching supports the intuition that more data, training time, and width produce better representations that remain compatible with weaker models. The findings also suggest layer-specific sample complexity and a modular structure in which better layers can be plugged in.

  • Number of samples: 25K-sample representations can be plugged into a 10K-sample model and improve the stitched network relative to the 10K network.Models trained with 5K, 10K, and 25K CIFAR-10 samples were compared by stitching their lower parts to a 10K-sample model.
  • Number of samples: Some layers are more data hungry than others: performance is similar at the first few layers but degrades rapidly with fewer samples in mid-layers.
  • Number of samples: Training the first three and last three layers with 5K samples, freezing them, and training the remainder with all samples produces a network within 2% accuracy of the original.
  • Training time: Representations from 40, 80, and 160 training epochs remain stitching compatible with the 80-epoch model, while earlier layers converge faster.
  • Width: Higher-width models can be stitched to lower-width models to improve performance, but the reverse direction does not work similarly.The width multipliers tested were 0.25×, 1×, and 2×.
  • Conclusion: Together, the results suggest neural networks have a modularity in which better layers can be plugged in without retraining the entire network.

7 Conclusion and Future Work

Model stitching is presented as an interpretable tool for comparing neural representations and assembling models from pretrained components. Its main costs are additional training and difficulty stitching different architectures, while the experiments remain limited to vision.

  • Conclusion: Model stitching compares representations in interpretable units but requires training, making it more computationally expensive than measures such as CKA.
  • Conclusion: Stitching representations from different architectures can be tricky and requires careful selection of the stitching family.
  • Future work: Although the experiments stitch first layers, the method can also plug in intermediate layers or parts of layers to assemble models from pretrained components.
  • Future work: The reported results are all in the vision domain, motivating future comparisons of representations in natural language processing and adversarially trained networks.

A Experimental Details

The experiments use several CIFAR-10 architectures and pretrained models, with convolutional stitching layers or a ViT embedding transform. Training protocols and stitching optimization are specified across the experimental families.

  • ResNets: The main convolutional models are CIFAR-10-adapted ResNet-18, wider ResNet-w2x, narrower ResNet-0.5x, and deep ResNet-164.The first convolution uses 64, 128, or 32 filters for the base, wider, and narrower models respectively.
  • ResNets: Sample-count and epoch experiments train the base ResNet-18 with the specified numbers of samples and epochs.
  • Label distribution: Label-distribution experiments use binary Object-Animal labels and random labels assigned to 10%, 50%, or 100% of the training set.
  • Optimization: ResNets are trained for 64K SGD steps with 0.9 momentum, scheduled learning-rate drops, weight decay 0.0001, and standard crop and flip augmentation.
  • Vision Transformer: The ViT uses patch size 4, depth 12, 12 attention heads, and dimension 768, with two initializations trained on CIFAR-5m and stitched on CIFAR-10.
  • Myrtle CNN: Myrtle CNNs are five-layer Conv-BatchNorm-ReLU-MaxPool networks with a fully connected output layer, trained for 80K steps at learning rate 1.
  • Stitching layers: Convolutional stitching uses BatchNorm2D and convolutional components, is placed between ResNet blocks, while ViT stitching uses a 768 × 768 linear transform.
  • Optimization: All stitching layers are optimized with Adam, cosine learning-rate scheduling, and initial learning rate 0.001.

B.1 Ablations

The ablations test whether stitching behavior depends on kernel size and whether the stitching layer merely learns a useful transformation. Kernel size has little effect, while early layers can stitch successfully to a random network.

  • Kernel sizes from 1 to 9 have minimal impact on the stitched network's test error.Both models are ResNet-18 networks trained on CIFAR-10 with different random initializations.
  • Early layers have low stitching penalties when a randomly initialized network is connected to a trained top model.This indicates that the early layers behave similarly to a randomly initialized network in this test.
  • SimCLR and end-to-end supervised representations have identical stitching performance despite differing more under CKA.

B.2.3 More is better

The section compares CKA and stitching across representation changes associated with width, data, and training time. Stitching often finds compatibility where CKA reports substantial differences, while longer training can preserve similarity.

  • Width: CKA can be as low as 0.7 for early layers comparing ResNet-18-w0.25x with ResNet-18-w2x, while early-layer stitching penalties remain small.
  • Samples: CKA can reach 0 when comparing networks trained with 5K versus 25K samples, whereas stitching shows a negative penalty for the more-data layers.
  • Training time: CKA finds similar representations between models trained for 80 and 160 epochs.
  • Comparison with fine-tuning: Finetuning can overestimate representation similarity relative to model stitching when comparing trained and randomly initialized ResNet-164 models.The comparison freezes the first l layers of the untrained bottom network and trains reinitialized top layers.

B.4 Freeze training with fewer samples

The freeze-training experiment uses stitching penalties to identify layers with lower sample complexity, then freezes those layers while training the rest with the full dataset. The resulting network reaches good performance.

  • Certain layers have smaller sample complexity because they reach stitching-connected representations with fewer samples.
  • Freezing five selected layers and training the remaining layers with the full dataset produces test performance of up to ≈3%.The selected layers are {0, 1} in one direction and {8, 9, 11} in the opposite direction.
  • Predicting which layers will have small sample complexity is identified as a direction for future research.
Loading 2106.07682v1…