Technical Summary: KLAS - Using Similarity to Stitch Neural Networks for Improved Accuracy-Efficiency Tradeoffs

1. Problem Statement

The deployment of large pretrained vision transformers (ViTs) and convolutional neural networks (CNNs) faces a fundamental challenge: diverse environments range from powerful cloud servers to resource-constrained edge devices, requiring flexible model selection within specific compute budgets. While Neural Architecture Search (NAS) attempts to balance accuracy and efficiency, traditional one-to-one NAS is computationally prohibitive, and one-to-many or zero-shot NAS approaches are often limited to a single model design space, failing to leverage the rich pool of existing pretrained models.

Model stitching offers a promising alternative by combining blocks from different pretrained "anchor" models to create interpolated networks. However, existing many-to-many stitching approaches, such as SN-Net, rely on naive heuristics (e.g., "nearest stitching" of adjacent anchors or fixed block pairing) to select stitch configurations. These heuristics often yield suboptimal accuracy-efficiency tradeoffs and lack generalizability across diverse model families. Furthermore, the search space for stitch configurations is vast ($O(k^2n^2)$ for $k$ models of depth $n$), making exhaustive search infeasible. The core problem is the lack of a principled, automated method to identify compatible stitch configurations that maximize performance without incurring high training costs.

2. Methodology: The KLAS Framework

The authors propose KLAS (KL divergence based Anchor Stitching), a framework that automates stitch selection by leveraging the similarity between intermediate representations of pretrained models.

Core Concept: KL Divergence as a Similarity Metric

KLAS posits that successful stitching depends on both representational similarity (geometric alignment of activations) and functional similarity (preservation of task-relevant information and decision boundaries). The authors argue that Kullback-Leibler (KL) divergence is uniquely suited to capture both:

• Representational Alignment: It quantifies the statistical distance between activation distributions, indicating if a lightweight transformation can map source activations to target activations.
• Functional Alignment: It assesses how well a source block's outputs preserve the target model's decision boundaries, ensuring the stitched network maintains downstream task performance with minimal finetuning.

Unlike metrics such as Centered Kernel Alignment (CKA), Mean Squared Error (MSE), or Cross-Entropy (CE), which often fail to predict stitching success, KL divergence directly measures the compatibility of probability distributions generated by linear probes.

The KLAS Algorithm

The framework operates in two primary stages:

1. Anchor Selection:

   • KLAS computes the KL divergence between the softmax probability distributions of the last blocks of candidate pretrained models.
   • Models with low last-block KL divergence are deemed compatible anchors, as they share similar decision boundaries and class-confidence distributions. This replaces the heuristic "nearest stitching" used in prior work.

2. Block Selection:

   • Once anchors are selected, KLAS evaluates block pairs using a stitch score $\Gamma(i, j)$ defined as:
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     Where:
     • $\Theta(P^f_i, P^g_j)$ is the cross-anchor activation distance (KL divergence between source block $i$ and target block $j$).
     • $\Theta(P^g_j, P^g_{j+1})$ is the intra-anchor block capacity (KL divergence between consecutive blocks in the target anchor).
   • This score balances the transformation burden on the hypothetical stitch layer against the target block's ability to absorb distributional differences.
   • ProbeNet: To efficiently compute these scores without training full stitch layers, KLAS employs a unified architecture called ProbeNet. It attaches linear classifier probes to every block of the anchor models and trains them jointly in a single forward-backward pass. This incurs negligible one-time cost (e.g., 0.25 GPU days for Swin-B) compared to training individual probes.

3. Configuration Selection:

   • Candidate stitch configurations are filtered based on a tunable threshold $\tau$ relative to the minimum score within FLOPs buckets. This ensures a dense Pareto front of deployable models across the accuracy-efficiency spectrum.

3. Key Contributions

1. Identification of Heuristic Limitations: The paper demonstrates that naive heuristics (like those in SN-Net) are not universally optimal, often missing superior "far stitching" configurations (e.g., connecting Swin-Ti directly to Swin-B) that heuristics overlook.
2. Principled Similarity Metric: The introduction of KL divergence as a metric for stitch selection that captures both representational and functional alignment without requiring the instantiation or training of stitch layers.
3. KLAS Framework: A generalizable, low-cost framework that automates anchor and block selection across diverse model families (ViTs, CNNs, and LLMs).
4. Empirical Validation: Comprehensive experiments showing consistent improvements over baselines across multiple datasets (ImageNet-1K, CIFAR-100, ADE20K, TruthfulQA) and architectures.

4. Experimental Results

The authors evaluated KLAS against SN-Net and other baselines under identical finetuning conditions (50 epochs for vision, 10 epochs for LLMs).

• Accuracy-Efficiency Tradeoffs:

  • ImageNet-1K: KLAS achieves up to 1.21% higher Top-1 accuracy at the same computational cost or a 1.33× reduction in FLOPs at equivalent accuracy compared to SN-Net.
  • AUC Scores: On the Swin family, KLAS achieved an Area Under the Curve (AUC) of 0.8950, significantly outperforming SN-Net (0.8345) and other similarity metrics like CKA (0.8124).
  • Cross-Architecture: KLAS successfully stitched ResNet-18 with Swin-B, outperforming cascaded inference baselines by 2-3% Top-1 accuracy at similar FLOPs.

• Dense Prediction (Semantic Segmentation):

  • On ADE20K using Mask2Former, KLAS-selected stitches achieved up to +0.9% mIoU over SN-Net baselines at comparable FLOPs budgets.

• Large Language Models (LLMs):

  • Using Llama 1B and 3B anchors on the TruthfulQA benchmark, KLAS-stitched models outperformed ESTA (a heuristic-based LLM stitching method) by up to 0.017 ROUGE-1 and 0.033 ROUGE-2, despite having fewer parameters in some configurations.

• Metric Effectiveness:

  • KL divergence successfully recovered 88.9% of viable Swin Ti-B stitch configurations, whereas CKA recovered only 5.5% and SN-Net heuristics 61.1%.

5. Significance and Claims

The paper claims that KLAS addresses the critical gap in model stitching by replacing ad-hoc heuristics with a principled, similarity-driven approach. Its significance lies in:

• Generalizability: It works across diverse architectural families (Transformers, CNNs, Hybrids) and modalities (Vision, Language), overcoming the limitations of methods restricted to single families.
• Cost-Efficiency: It generates a continuous Pareto front of deployable models with negligible additional computational overhead (only probe training and stitch-layer finetuning), avoiding the prohibitive costs of traditional NAS.
• Systematic Optimization: By explicitly capturing distributional differences via KL divergence, KLAS enables the discovery of non-intuitive, high-performing stitch configurations (e.g., "far stitching") that heuristic methods miss.

The authors conclude that KLAS provides a robust, automated solution for optimizing the accuracy-efficiency tradeoff in deployment scenarios, effectively leveraging the existing ecosystem of pretrained models without requiring training from scratch.