Technical Summary: Grounding Functional Similarity by Invariance-Aware Model Stitching

Problem Statement

Functional similarity evaluation aims to quantify the extent to which independently trained neural networks learn similar input–output relationships. A prominent method for this is model stitching, where the early layers (front model) of one network are connected to the later layers (end model) of another via a trainable alignment layer. If the resulting stitched composition maintains high performance, the models are deemed functionally similar.

However, recent studies (e.g., Smith et al., 2025) have identified a critical limitation in standard model stitching: it relies exclusively on forward compatibility. This means it only evaluates whether the representations of the front model can be aligned to solve the end model's task. Consequently, models that rely on fundamentally different information cues (e.g., color vs. shape, or robust vs. non-robust features) can still produce compatible representations, leading to "misleading" assessments of similarity. Furthermore, standard stitching can exploit previously unseen predictive cues (shortcuts) to achieve alignment, further obscuring true functional discrepancies.

The authors argue that standard stitching is inherently blind to the invariance properties of the models. Specifically, it fails to account for the backward direction: whether the end model processes the stitched representations in a manner consistent with how it processes its own native representations, particularly regarding shared invariances.

Methodology

1. Forward–Backward Compatibility Requirement

To address the limitations of forward-only evaluation, the authors introduce a forward–backward compatibility requirement. This framework posits that for two models to be functionally similar, they must satisfy two conditions:

• Forward Compatibility: The front model's representations can be aligned to the end model to solve the task.
• Backward Compatibility: For inputs that the front model represents identically (Indistinguishable Representations), the end model must process these representations similarly whether they originate from the front model or its own native path.

This requirement ensures that the stitched composition respects the invariance classes induced by the front model.

2. Invariance-Aware Model Stitching

The authors operationalize this requirement through Invariance-Aware Model Stitching. The core mechanism involves:

• Identically Represented Inputs (IRIs): Using a front model $f_{\le i}$, they generate a dataset $D_{IRIs}$ where inputs $x'$ are optimized to yield nearly identical representations to original inputs $x$ at the stitching layer ($||f_{\le i}(x') - f_{\le i}(x)|| \approx 0$).
• Training Protocol: The stitching layer is trained on $D_{IRIs}$ (using the front model's representations) but evaluated on the standard test set $D_{test}$.
• Logic: If the models are truly functionally similar, the stitching alignment learned on IRIs should generalize to the standard test set. If the alignment relies on cues specific to the IRIs generation or fails to respect the end model's invariances, performance will drop, revealing functional discrepancies.

3. Stitching Alignment Objectives

The paper evaluates several alignment objectives to establish the stitching transformation $T_\theta$:

• Hard Label Matching (HLM): Minimizes cross-entropy between the stitched output and ground truth labels.
• Soft Label Matching (SLM): Minimizes cross-entropy between the stitched output and the soft predictions of the standalone end model.
• Direct Matching (DM): Minimizes the $L_2$ distance between the front model's output and the end model's intermediate features at the stitch level.
• Functional Latent Alignment (FuLA): A novel objective proposed by the authors. It minimizes the distance between the stitched composition's internal features and the end model's features across multiple layers (from the stitch level up to the penultimate layer), effectively mimicking the end model's internal processing flow without relying on output labels.

Key Contributions

1. Identification of Failure Modes: The authors demonstrate that standard model stitching (even under SLM and DM, not just HLM) fails to distinguish between models relying on different information cues. They show that these methods can produce high similarity scores even when models utilize distinct visual cues (e.g., pixel shortcuts) or when one model is robust and the other is not.
2. Invariance-Aware Framework: They propose and validate invariance-aware model stitching as a more principled approach. By enforcing forward–backward compatibility, this method reveals functional discrepancies that standard stitching obscures.
3. Re-evaluation of Robustness: The paper revisits the functional similarity between robust (adversarially trained) and non-robust networks. While previous work suggested they are functionally similar, the authors show that under invariance-aware stitching, they exhibit significant functional incompatibilities, particularly regarding how they handle adversarial perturbations.
4. Analysis of Alignment Objectives: The study analyzes the interplay between different stitching objectives. It finds that while FuLA and SLM generally yield higher functional similarity, Direct Matching (DM) often reveals misalignment at the latent level. The authors highlight a trade-off: objectives that enforce stricter latent alignment (like DM) may be less prone to exploiting unseen cues but might yield lower overall similarity scores.

Experimental Results

The authors conducted experiments on image classification tasks using CIFAR-10, LS-ImageNet-10, and ImageNet-100 with architectures including ResNet-18, VGG-16, ViT-Tiny, and ViT-Base.

• Different Information Cues: When stitching models trained on different variants of CIFAR-10 (e.g., one trained with pixel shortcuts), standard stitching showed high similarity in deep layers. In contrast, invariance-aware stitching showed a sharp drop in similarity for the cross-bias configuration, correctly identifying the lack of functional similarity.
• Unseen Predictive Information: Standard stitching formulations were found to exploit "shortcuts" (predictive cues not present in the original training data) to align models. Invariance-aware stitching substantially mitigated this behavior, reducing "excessive agreement" to near zero when such cues were unavailable.
• Robust vs. Non-Robust Models: Standard stitching suggested robust and non-robust models were functionally similar. Invariance-aware stitching revealed that robust models share a more consistent representation structure with each other, while non-robust models develop idiosyncratic invariances that are not shared with robust models.
• Objective Comparison: FuLA and SLM generally achieved comparable functional similarity, while DM often resulted in lower similarity scores, suggesting that misalignment exists at the feature representation level even before the output layer.

Significance and Claims

The paper claims that invariance-aware model stitching provides a more grounded and principled metric for functional similarity. By moving beyond forward compatibility to include backward compatibility, the method prevents misleading conclusions where models appear similar simply because they can be aligned to solve a task, despite relying on different underlying mechanisms.

The authors emphasize that their approach reveals "functional discrepancies previously obscured," offering a finer-grained understanding of how neural networks operate internally. They conclude that evaluating similarity from multiple perspectives (forward and backward) is essential for accurately assessing the inner workings of deep neural networks, particularly in contexts involving robustness and diverse training objectives.

Limitations: The authors acknowledge the absence of an "oracle" notion of functional similarity, meaning definitive conclusions about which setting is the "most faithful" remain challenging. Additionally, the study is currently limited to image classification models.