How not to Stitch Representations to Measure Similarity: Task Loss Matching versus Direct Matching

TL;DR

This paper critically evaluates model stitching methods for measuring neural network similarity, revealing that task loss matching can be misleading and proposing direct matching as a more reliable alternative.

Contribution

The paper identifies limitations of task loss matching and demonstrates that direct matching provides a more accurate measure of representation similarity in neural networks.

Findings

Task loss matching can indicate high similarity between distant layers.

Task loss matching may suggest non-corresponding layers are more similar.

Direct matching avoids out-of-distribution issues and aligns better with structural similarity.

Abstract

Measuring the similarity of the internal representations of deep neural networks is an important and challenging problem. Model stitching has been proposed as a possible approach, where two half-networks are connected by mapping the output of the first half-network to the input of the second one. The representations are considered functionally similar if the resulting stitched network achieves good task-specific performance. The mapping is normally created by training an affine stitching layer on the task at hand while freezing the two half-networks, a method called task loss matching. Here, we argue that task loss matching may be very misleading as a similarity index. For example, it can indicate very high similarity between very distant layers, whose representations are known to have different functional properties. Moreover, it can indicate very distant layers to be more similar than architecturally corresponding layers. Even more surprisingly, when comparing layers within the same network, task loss matching often indicates that some layers are more similar to a layer than itself. We argue that the main reason behind these problems is that task loss matching tends to create out-of-distribution representations to improve task-specific performance. We demonstrate that direct matching (when the mapping minimizes the distance between the stitched representations) does not suffer from these problems. We compare task loss matching, direct matching, and well-known similarity indices such as CCA and CKA. We conclude that direct matching strikes a good balance between the structural and functional requirements for a good similarity index.