StitchNet: Composing Neural Networks from Pre-Trained Fragments

TL;DR

StitchNet introduces a method to assemble neural networks from pre-trained fragments, enabling high-performance models with reduced training resources and personalized, on-the-fly creation.

Contribution

This work presents a novel paradigm for neural network construction by stitching pre-trained fragments guided by CKA, reducing training costs and enabling personalized models.

Findings

Achieves comparable accuracy to traditional training methods

Reduces compute and data requirements significantly

Enables on-the-fly personalized model creation

Abstract

We propose StitchNet, a novel neural network creation paradigm that stitches together fragments (one or more consecutive network layers) from multiple pre-trained neural networks. StitchNet allows the creation of high-performing neural networks without the large compute and data requirements needed under traditional model creation processes via backpropagation training. We leverage Centered Kernel Alignment (CKA) as a compatibility measure to efficiently guide the selection of these fragments in composing a network for a given task tailored to specific accuracy needs and computing resource constraints. We then show that these fragments can be stitched together to create neural networks with accuracy comparable to that of traditionally trained networks at a fraction of computing resource and data requirements. Finally, we explore a novel on-the-fly personalized model creation and inference application enabled by this new paradigm. The code is available at https://github.com/steerapi/stitchnet.