Model Stitching by Functional Latent Alignment

2025-05-26

AI Summary

This paper addresses the challenge of reliably quantifying functional similarity between independently trained neural networks. Existing methods suffer from susceptibility to task-irrelevant cues and difficulty in identifying nontrivial alignments. To overcome these limitations, we propose Functional Latent-space Alignment (FuLA), a novel model stitching optimization criterion inspired by knowledge distillation. FuLA achieves cross-model functional alignment via learnable affine transformations in the latent space, effectively suppressing biases induced by adversarial training, shortcut learning, and cross-layer stitching. Evaluated on multiple functional similarity benchmarks, FuLA demonstrates superior robustness and alignment quality compared to conventional approaches—particularly under substantial training disparities—and reveals previously overlooked deep functional equivalences.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Evaluating functional similarity in neural networks
• Improving model stitching for reliable alignment
• Addressing artifacts in training for better representations

Innovation

• Functional Latent Alignment for model stitching
• Optimal affine transformation aligns models
• FuLA improves reliability in similarity tests