Model Stitching by Functional Latent Alignment

TL;DR

This paper introduces Functional Latent Alignment (FuLA), a new method for assessing neural network functional similarity through model stitching, demonstrating improved reliability over previous techniques across various training scenarios.

Contribution

The paper proposes FuLA, a novel optimality condition for model stitching that enhances the detection of functional similarity in neural networks, surpassing prior methods in reliability.

Findings

FuLA outperforms previous methods in adversarial, shortcut, and cross-layer stitching.

FuLA is less affected by training artifacts and better captures true functional similarity.

Experiments show FuLA identifies meaningful alignments missed by stitch-level matching.

Abstract

Evaluating functional similarity involves quantifying the degree to which independently trained neural networks learn functionally similar representations. Reliably inferring the functional similarity of these networks remains an open problem with far-reaching implications for AI. Model stitching has emerged as a promising paradigm, where an optimal affine transformation aligns two models to solve a task, with the stitched model serving as a proxy for functional similarity. In this work, we draw inspiration from the knowledge distillation literature and propose Functional Latent Alignment (FuLA) as a novel optimality condition for model stitching. We revisit previously explored functional similarity testbeds and introduce a new one, based on which FuLA emerges as an overall more reliable method of functional similarity. Specifically, our experiments in (a) adversarial training, (b) shortcut training and, (c) cross-layer stitching, reveal that FuLA is less prone to artifacts tied to training on task cues while achieving non-trivial alignments that are missed by stitch-level matching.