KLAS: Using Similarity to Stitch Neural Networks for Improved Accuracy-Efficiency Tradeoffs
Authors: Debopam Sanyal
Anantharaman S. Iyer
Alind Khare
Trisha Jain
Akshay Jajoo
Myungjin Lee
James Clayton Kerce
Alexey Tumanov
Venue: ICLR 2026 Poster (2026)
Keywords: representation learning
neural networks
deep learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: KLAS is a stitch selection algorithm that automatically improves accuracy-efficiency curves by leveraging KL divergence to identify stitching points between pretrained models.

Abstract: Given the wide range of deployment targets, flexible model selection is essential for optimizing performance within a given compute budget.
Recent work demonstrates that stitching pretrained models within a model family enables cost-effective interpolation of the accuracy-efficiency tradeoff space.
Stitching transforms intermediate activations from one pretrained model into another, producing a new interpolated stitched network.
Such networks provide a pool of deployment options along the accuracy-efficiency spectrum.
However, existing stitching approaches often yield suboptimal tradeoffs and lack generalizability, as they primarily rely on heuristics to select stitch configurations.
We argue that constructing improved accuracy-efficiency tradeoffs requires explicitly capturing and leveraging the similarity between pretrained models being stitched.
To this end, we introduce KLAS, a novel stitch selection framework that automates and generalizes stitch selection across model families by leveraging KL divergence between intermediate representations.
KLAS identifies the most promising binary stitches from the $\mathcal{O}(k^2n^2)$ possibilities for $k$ pretrained models of depth $n$.
Through comprehensive experiments, we demonstrate that KLAS improves the accuracy-efficiency curve of stitched models at the same finetuning cost as baselines.
KLAS achieves up to $1.21\%$ higher ImageNet-1K top-1 accuracy at the same computational cost, or maintains accuracy with a $1.33\times$ reduction in FLOPs.

Published as a conference paper at ICLR 2026
KLAS: USING SIMILARITY TO STITCH NEURAL NETWORKS FOR IMPROVED ECCURACY-EFFICIENC
TRADEOFFS
Dehopam Sanyal', Anantharaman Iyer', Alind Khare?, Trisha Jain', Akshay Jajoo", Myungjin Lee", Clayton Kerce", Alexey Tumanov'
Georgia Institute of Technology. "Microsoft M365 Research, *Cisco Research, 'Georgia Tech Research Institute debopam. sanyal @gatech.
ABSTRACT
Given the wide range of deployment targets, flexible model selection essential for optimizing performance within given compute budget. Recent work demon- strates that stitching pretrained models within a model family enables cost-effective interpolation of the accuracy-efficiency tradeoff space. Stitching transforms in- termediate activations from one pretrained model into another, producing a new interpolated stitched network. Such networks provide a pool of deployment options along the securacy-effciinng spectrum. However, existing stitching approaches often yield suboptimal tradeoffs and lack generalizability, as they primarily rely on heuristics to select stitch configurations. We argue that constructing improved accuracy-efiic -fffirccceff tradeoff's requires explicitly capturing and leveraging the simi- larity between pretrained models being stitched. To this end, we introduce KLAS, novel stitch selection framework that automates and generalizes stitch selec- tion across model families by leveraging KL divergence between intermediate representations. KLAS identifies the most promising binary stitches from the O(k2n²) possibilities for k pretrained models of depth Through comprehensive experiments, we demonstrate that KL.AS improves the accuracy-effcienny curve of stitched models the same finetuning cost baselines. KLAS achieves up to 1.21% higher ImageNet- IK top-1 accuracy at the same computational cost, or maintains accuracy with 1.33x reduction in FLOPs.
INTRODUCTION
Large pretrained vision transformers (ViTs) and convolutional neural networks (CNNs) have trans- formed computer vision |Devlin et al. 2 Brown et al. (2020); Peters al. (018); Srivastava et al.. (2025), powering applications from self-driving cars Ouyang et al. (2019) to security systems Chen ctt al. (202I); Ji et al. (202I) and social media platforms Bai et al. (2018). Trained on vast datasets and distributed via repositories such as HuggingFace Wolf 2019). Py TorchHub Paszke et al. (2019), and TensorFlow Model Zoo Abadi et al. (2016), these models achieve strong performance across tasks including image classification, object detection, semantic segmentation, and face recognition Liu et al.) (2022): Ravi et al. (2024); y
An et al. (022). Yet, deployment environments vary dramatically from powerful cloud servers Hazelwood et al. (2018) to resource-constrained edge devices He al. (2018): David et al. (202I). This poses fundamental challenge: how can we efficiently deploy pretrained models under diverse computational budgets without sacrificing accuracy?
Neural architecture search (NAS) Cai et al. (2 Zoph & Le (016) aims to address this challenge by discovering architectures that balance accuracy and efficiency, but conventional one-to-one NAS requires prohibitively high GPU cost.
paradigms, such as one-shot NAS Liu al. (018): Guo et al. (2020) mitigate this burden by training a single supernetwork, while zero-shot NAS Li et al. (2023: Lin et al. (202T) avoids costly training altogether by ranking candidate architectures using gradient-based proxies. Despite these advances, both approaches remain limited: they confined to single model design space and cannot exploit the rich pool of pretrained models already available. This limitation prevents one-to-man NAS from satisfying the three key requirements of effective architecture search: (1) strong accuracy-efficient tradeoffs, (2) generalizability across diverse model architectures, and (3) low computational cost.
Many- tonyyyy NAS with model stitching Pan et al. (2023); y Bansal al. (2021) Yang et al. (2022b); Pan et al. (2024) offers a promising solution by combining blocks from different pretrained "anchor" models to form stitched networks that interpolate accuracy-efficiency tradeoffs. For example, SN-

---

Published as a conference paper at ICLR 2026
Pan et al. (2023), finetunes lightweight stitching layers (simple linear transformations) between anchors to drastically reduce cost. A key challenge, however, lies in selecting stitch configurations (i.e.. choosing which anchors and blocks to stitch). Even with only two anchor models of depth there are (nn²) possible stitch configurations, making exhaustive search computationally infeasible. Yet, SN-Net employs naive heuristic- hased selection (e.g., stitching adjacent anchors and blocks), which limits both the accuracy-efficiency tradeoff' and generalizability across model families.
We propose novel framework, KLAS (KL diver- gence based Anchor Stitching), which replaces heuristic-basee based stitching with principled similarity- driven approach. As illustrated in Fig. [ KLAS au- tomatically selects anchors and block pairs by eval- uating the compatibility of their activations using Kullback-Leibler (KL) divergence (1999). Un-
O(n') stach cantig choioas
like alternatives such mean square error Jadon ct al. (2024), cross-entrope Mao et al. (2023), or
CKA Kornblith et al. 2019). KL divergence effec- tively captures distributional differences between de-
Stithhing (SN-Net)
KLAS Algoritmm (Ours)
eision boundaries, enabling more accurate and effi- cient stitching with minimal finetuning. In doing so,
Devise houristics
Similarity using KL- divergence
KLAS addresses two key stitching questions: (1) can network B produce similar outputs when driven by transformed internal representations from network A? and (2) can such transformations be effective with minimal finetuning?
Select fixed anchors and blocks
Automatically select anchors blocks
Swiin
DeiT
ResNet
Swin
DeiT
ResNet
We evaluate KLAS on ImageNet-Il IK Deng et al. (2009) and CIFAR-100 Krizhevsky al. (2009) across multiple ViT-based and CNN-based model families and show that it consistently improves the accuracy-efficiency tradeoff over SN-Net. KLAS achieves up to 1.21% higher top-1 accuracy at the same FLOPs, alternatively, a 1.33x reduction in FLOPs at equivalent accuracy. We further show bene- fits of KLAS on large language models (LLMs) using TruthfulQA Lin et al. (2022). Our key contributions are: (1) identifying the limitations of heuristic-based many-to-many NAS approaches such as SN-Net, (2) introducing KL divergence as a principled similarity metric for selecting stitch configurations without instantiating training a stitch layer, (3) presenting the KLAS framework for efficient, generalizable, and low-cost stitching, and (4) demonstrating consistent improvements over SN-Net across model families and datasets.
Infordor accuncy-affiioonny
Superior acourap--fffiiency
(radeoffs
tradectfs
Figure Comparison of stitching strategics. Ex- isting many-to-many NAS approaches such as SN- Net rely on heuristic-based stitch selection, fix- ing anchors and blocks and yielding suboptimal accuracy-effiiienc tradeoffs. In contrast, our pro- posed KLAS framework leverages KL divergence to measure similarity between intermediate rep- resentations and automatically identify promising stitch configurations, producing superior accuracy- efficiency tradeoffs across diverse model families.
RELATED WORK
Model Stitching:
Model stitching was originally introduced as a tool for analyzing representation
netn & Vedaldi (2015) showed that early layers of one network could be connected to later layers of another via simple linear transformations, revealing functional similarities between disparate architectures. Building on this, [Bansal al. (202T) demonstrated that stitching remains effective even across models with different architectures or training protocols. While these studies established stitching as a powerful analytical tool, they did not investigate its potential for practical deployment optimization.
cquivalence across neural networks. [Lnce
Recent efforts have begun bridging this gap. SN-Net Pan et al. (2023) operationalized stitching for neural architecture search by combining pretrained anchor models from the same family, reducing training costs by 22x compared to traditional NAS. However, SN-Net relies on heuristic stitching strategies, such as connecting adjacent anchors ('nearest stitching") or assuming block similarity ("paired unpaired stitching"), " that ignore actual compatibility between blocks. ESTA He et al. (2024) extended SN-Net to large language models by applying the same heuristic, cever--ount-based stitching strategy when combining pretrained LLM blocks. StitchLLM et al. (2025) further adopts this identical heuristic within systems framework focused on resource-efficiet serving, but does not.
Pan
2

---

Published as a conference paper at ICLR 2026
propose a new stitch selection criterion. Concurrently, Csiszárik Csiszrii al. (2021 systematically analyzed stitching performance across block combinations, finding no correlation between common similarity metrics (e.g., CKA) and stitched model accuracy. Similarly, h & Jelasity (2025) showed that direct similarity measures often fail to reflect task performance, advocating task loss matching as a more reliable criterion. Together, these findings underscore the inadequacy of existing approaches for automatically identifying effective stitch points, a challenge that KLAS addresses through activation distribution analysis using KL divergence.
Neural Architecture Search: Traditional NAS methods Cai et al. (2020); Sahni et al. (2021): Pham al. (2118); Zoph & Le (2016); Wang al.] (202I) incur prohibitive computational costs, often requiring thousands of GPU hours to evaluate candidate architectures. One-shot NAS et al. (2018) mitigates this burden by training single supernetwork with shared weights across subnetworks, but still demands substantial resources for supernetwork training and remains restricted to a single architectural family. Zero-shot NAS et al. (202T) further reduces cost by predicting subnetwork performance from architectural indicators, yet the final models still require full training from scratch. Cascaded inference has long been explored as a means to reduce compute by exiting early when predictions are confident Viola & Jones (2001). Approaches include big / little-style two-stage models Park et al. (2015). multi-exit deep cascades with learned entropy thresholds Wang et al. (2017): Guan al. (2017); BBolubaasi et al. (2017), and exit-by-confidence ensemble routing Streeter (0188) Inoue (019). Recent work Wang et al. (2020); [Lezovizee et al.) (2023) extends cascades to four-way model ensembles using softmax-based confidence and input resolution scaling.
Stitching circumvents these limitations by leveraging pretrained model zoos rather than training networks from the ground up. Unlike DeRy Yang et al. (2022a). which reassembles pretrained components for individual resource constraints, KLAS systematically analyzes activation distributions to generate a continuous Pareto front of deployable models. This approach preserves cost efficiency by requiring only stitch-layer finetuning, while avoiding architectural constraints by enabling seamless integration of transformers, convnets, and hybrid models.
The rise of ViTs Dosovitskiy et al. (2020) has introduced new challenges for efficient deployment. Prior efforts optimized individual ViT's through token gruning[Pan) e
Transformers and ConvNets:
al. 202I), quantization Li et al. (2022), or dynamic inference Rao al. (202I), but these methods remain architecture--seecffc and require retraining under different resource constraints. Cross- architectural stitching approaches such SN-Net Pan et al. (2023) explored combining ViTs, yet were limited to same-family stitching due to reliance on weak heuristics.
KLAS breaks this barrier by introducing KL divergence between block activations intuitive simi- larity measure, enabling stitching across diverse architectural families. Whereas prior work Kornblith et al, (2019) found standard similarity metrics ineffective for predicting stitching performance, KLAS directly measures how well the outputs of one block can serve as the inputs to another, regardless of architecture. This capability allows KLAS to produce viable hybrids spanning hierarchical Swin Transformers and residual CNNs.
3 BACKGROUND AND MOTIVATION
3.1 NOTATIONS FOR MODEL STITCHING
Let ʃ X 3/ be feedforward neural network with m blocks (i.e., layers): ʃ fm .... o Si that of block i. By definition,
where fi Af_1 + A! maps the activation space of block Af X (i.e.. the input samples). For model stitching, we introduce the notations fsi fo ..... fi and f>i = fm ....o fi+1. Given two frozen networks f and g. and one block from each network. i E and g. the goal of stitching is to find out if 9>j. which we will refer to the target, can achieve its function using the representation of fsi, which we call the source. In examining this, we attempt to find transformation layer T Af- > A3 such that the stitched network I>j o T f's is functionally similar g. Here, we term (i.j) stitch configuration, meaning that the i-th layer of f is stitched to the j-th layer of g. Throughout this paper, we refer to T as the stitching layer. For an input x E X, we denote (<)(2) Transformer blocks with the same hidden dimensions are often grouped into stages. We omit the stage index in our notations for simplicity stitching only takes place between same stage blocks.
its source representation and (x)) as its target representation.
3

---

Published as a conference paper at ICLR 2026
3.2 SIMILARITY OF NEURAL NETWORKS
A central theme in analyzing neural networks is understanding the representations learned by internal layers, since these insights crucial for both interpreting and improving deep learning systems. A fundamental question in this context is: when are two representations similar? Prior work [Bansal et al. (202I); Kormblith al. 2 Yosinski et al. (2014) extends this inquiry by asking complementary question: in what way are two representations similar, once similarity is established? Together. these perspectives give rise to two notions of similarity: (1) representational similarity, which characterizes structural relationships between learned features, and (2) functional similarity, which evaluates their interchangeability in downstream tasks. We explain below why both are essential.
Representational similarity captures the alignment of geometric structures between intermediate activations, enabling the stitching layer (7) to learn a simple mapping (e.g., affine transformation) from the source the target. For instance, convolutional layers with comparable channel counts and spatial dimensions often permit lightweight 1 1x1 convolutions as effective stitching layers. Conversely, mismatched representations may require complex transformations, which can degrade performance or even cause the stitched network to fail. A widely used metric for representational similarity is Centered Kernel Alignment (CKA)) Kornblith et al. (2019), which leverages kernel functions to embed data into a higher-dimensional space where linear relationships between representations become more discernible.
Functional similarity evaluates whether the stitched network preserves the target's original output behavior, typically measured via task accuracy. High functional similarity implies that the source representations contain sufficient task-relevant information for the target, even if their geometric structures differ. Two common approaches include task loss matching (TLM), which optimizes the stitching layer by minimizing downstream task loss (e.g., mean squared error (MSE) or cross- pross-entropy (CE)), and direct matching (DM), which aligns source activations after the stitch layer and target activations through linear regression BBllogh & Jelasity (2025)
3.3
SHORTCOMINGS OF CURRENT SIMILARITY MEASURES
Metric
Type
Overlap
Top-1 Accuracy 186)
Swin B (SN-Net)
CKA
X (Unsupervised)
5.5%
83-
Swin-S S ISN-Nett Net)
MSE
/ (Supervised)
27.8%
82-
V (Supervised)
22.2%
DM
V (Supervised)
33.3%
Swi--T (FS)
SN-Net
X (Unsupervised)
61.1%
80
*
Nearest Stitching (SN Net)
Swin-Ti (SN-Net)
Far Stitching (FS)
KL-DivOwry
V (Supervised)
88.9%
B
10
12
14
16
FLOPs (G)
Table 1: Effectiveness of similarity metrics for recovering Swin Ti-B stitching, measured by overlap with the Ti-B stitch configurations. Existing metrics (CKA, MSE, CE. DM) and SN-Net heuristics achieve at most 61.1% overlap, while our KI. divergence based ap- proach achieves 88,9% demonstrating its re- liability as a similarity measure.
Figure 2: Comparison of nearest stitching (SN-Net) and far stitching (FS) on the Swin family. Far stitching hetween Ti and B models (green curve) achieves superior accuracy- efficiency tradeoff. surpassing nearest stitching (pink curve) by up 0.9% accuracy at equivalent FLOPs. Stars mark pretrained anchors (Ti, S, B).
In our empirical study of similarity metrics for model stitching, we identify key limitation of the naive heuristics adopted by SN-Net. Specifically, SN-Net assumes that combining nearest stitching for anchors with puired/ / unpaired stitching for blocks yields optimal results across model families. Nearest stitching restricts anchor connections to models of comparable complexity or performance, such the Ti-S-B stitching style in Swin transformers. However, experiments reveal that direct far stitching between Ti and B models (i.c., Ti-B stitching) achieves a superior accuracy-efficiene tradeoff, improving accuracy by up to 0.9% equivalent FLOPs, shown in Fig. 2 For this study, we apply paired unpaired block stitching with Twin-TT/S/BBBB models pretrained on ImageNet-22K, followed by finetuning the stitched supernetwork for 50 epochs ImageNet-1K:
We further examine whether existing similarity measures can automatically recover the Swin Ti-B stitching by evaluating the percentage overlap of stitch configurations. As shown in Tab. [ commonly
4

---

Published as a conference paper at ICLR 2026
used metrics such MSE, CE, CKA, and DM fail do so, while KL divergence is successful. The gap is particularly pronounced for CKA, which recovers only 5.5% of the configurations. We provide explanations in $.1 and further studies in $ and App. D
PROPOSED APPROACH
We present KLAS, generalizable stitch selection technique that identifies promising configurations for favorable sccuracy-efficien tradeoffs, at no additional finetuning cost over baselines. Central to KLAS is the observation that stitching success hinges activation similarity, which determines whether two networks can be integrated through a linear transformation layer. When stitching the source network's 's early layers (fez) the target network's s later layers (g>ş), the stitching layer (T) must map the source activations (Af) to the target activations (Af) while keeping finetuning costs low and preserving task performance. The first requirement corresponds representational similarity, and the second to functional similarity. At its core, KLAS employs KI. divergence, which show in 1 satisfies both objectives, unlike prior metrics that capture only one. This dual alignment ensures stitching viability: representational alignment simplifies stitch-layer learning, while functional alignment guarantees good downstream task performance.
4.1 KL DIVERGENCE AS A SIMILARITY METRIC
KL divergence satisfies the dual objectives of measuring both representational and functional simi- larity by directly quantifying alignment between activation distributions and their downstream task performance. As a representational similarity metric, KL divergence captures distributional differ- ences between intermediate activations, indicating whether two blocks generate patterns that can be mapped via lightweight transformations. At the same time, it serves as a functional similarity metric by assessing how well one block's s outputs preserve the target model's 's decision boundaries: low KL value implies that the source block provides task-relevant information that the target can leverage with minimal finetuning. This dual capability arises from KI. divergence's grounding in information theory, it evaluates both statistical distance between representations (representational alignment) and the retention of class- diisssiicriiiiative information (functional alignment). By minimizing KL divergence between intermediate activations, we ensure that the stitching layer requires only minor adjustments to align representations while preserving the target's 's original performance: an advantage over metrics that emphasize either representational alignment or task loss matching.
13
(a): Swin
(b): DefT
(c): LeViT
Figure 3: Probe training convergence architectures. Top-1 accuracy of linear probes trained for four epochs shown for each block in (a) Swin. (b) DeiT, (c) LeVIT families. Probes converge rapidly, demonstrating that meaningful intermediate representations can be extracted with negligible, one-time probe training cost.
Linear Probe Classifiers: Formally, KL divergence measures the statistical distance from probabil- ity distribution to a target distribution y, denoted as DKL (pll4) 4). We use linear classifier probes Alain Bengio (2018). placed after different blocks, to estimate the KL divergence between the softmax probability distributions of activations. Each block in an anchor model is equipped with a simple 1 x convolutional probe, for instance, Swin-B requires 24 probes, while DeiT-S requires 12. Training all probes independently would be very expensive. Instead, we introduce unified architecture, ProbeNet, which jointly trains all probes by activating one at a time within a single forward-backwer pass. This design adds only a negligible one-time cost, probes converge rapidly in setup (by epoch 4; Fig. 3). For example, training all 24 probes in Swin-B required just 0.25 GPU days in total, rather than 24 0.25 GPU days if trained separately. The pseudocode for ProbeNet is provided in Alg. Iin App. A By avoiding redundant inference and backpropagation, ProbeNet makes large-scale linear probing both efficient and practical.
5

---

Published as a conference paper at ICLR 2026
Let P. (z) denote the softmax probability distribution produced by the probe inserted after block of anchor model f on imput z E De. where D, is the validation split. Similarly, let P. (z) represent the distribution from the probe after block j of anchor model g. After training the linear probes on the training split (Dt), we compute the similarity score é as defined in Eq. D We outline the procedure for obtaining KI. divergence based similarity scores for all source-target block pairs using classifier probes in Alg. in App. A In this work, we fix the direction of stitching such that the source anchor always has lower complexity than the target anchor. Experiments in Pan et al. (2023) indicate that this choice yields more stable training and consistently better performance compared to stitching in the reverse direction.
crOSs anchor activation distance (Q)
Σ ZED
DKL(P Pf
@(P! .P))
O(P,.PP)
[(i,))
(2)
[Del
(I)
intro-anchor block capacity (I)
4.2 THE KLAS FRAMEWORK: KL DIVERGENCE BASED ANCHOR STITCHING
Anchor Selection using Last Block KL Divergence: The first step in stitching is to identify suitable anchors (i.e.. pretrained models). We achieve this by computing the KI. divergence between the softmax probability distributions produced by the last block of each anchor. A low KI. divergence indicates that two models have similar decision boundaries and class-contidence distributions, making them more compatible for stitching. This provides principled anchor selection strategy, in contrast to the heuristic nearest stitching used in SN-Net.
Block Selection using KI. Divergence: Once the anchors chosen, we select stitchable block pairs using a stitch score I, defined in Eq. 2 and demonstrated in Fig. 4 Unlike SN-Net's 's naive paired / unpaired stitching heuristic for block selection, our method evaluates the compatibility of block pairs without instantiating or training a stitch layer. This design is central to KLAS's s efficiency: candidate stitched models are assessed without incurring any finetuning cost. When stitching a source block i from anchor f to target block in anchor g, the resulting stitched model タ>j ' fe; interpolates between f and g in both accuracy and FLOPs. Thus, the target anchor's 's accuracy serves as an upper bound for any stitched network derived from it, i..., Age(g) is the maximum achievable accuracy of 9>j o T o fsi. Vi.j.
(1) Cross-anchor activation distance [Ω = O(Pp P9) quan- tifies the transformation a hypothetical stitch layer (Z) would need to apply to the source activations (A!) in order to match the target activations (A.). In other words, it measures how much the stitched model Р>j T o fs: would have to modify its intermediate representations to approximate the behavior of the target model g. A smaller l implies that the activations after layer i in f are already close to the activations after layer j in g, making the block pair (i, more favorable for stitching. is important to emphasize that this is a purely hypothetical measure: since the models are not stitched yet, the (j + 1) block in does not actually process activations from f. Rather, Ω reflects the expected burden on the stitch layer alone if such a connection were to be instantiated.
G
Σ
Figure 4: Ω measures how closely the activations of source block i from align with those of target block j in reflecting the transformation burden on the hypothetical stitch layer. Σ mea- sures how much block j+ 1 in the target anchor y transforms its inputs relative to block j. indicating block (i + 1)'s ability to absorb mismatched represen- tations from the source.
(2) Intra-anchor block capacity === O(P. Po+1)l captures how much a block in pretrained anchor transforms its input representations, measured by comparing the outputs of con- secutive blocks within the same anchor. A large Σ indicates substantial shift in activations from block j to block j + 1 in model g. Since g achieves any stitched model's 's upper-bound accuracy, this implies that block j + has high learning capacity, i.e., it plays a critical role in refining representations for task performance. High-capacity blocks are especially valuable when constructing stitched models, they can absorb distributional differences from the source while maintaining task-relevant transformations. Thus, incorporating such blocks increases the likelihood of achieving strong downstream accuracy without high training cost.

---

Published as a conference paper at ICLR 2026
KLAS algorithm formulation: We employ probe classifiers project the activation spaces of intermediate blocks onto the output space y. Since stitching is restricted to occur within the same stage, omit explicit stage indices for notational simplicity (see $3.1). During the stitch finetuning phase, we select a set of candidate stitch configurations (S) using a tunable threshold (r). These candidates drawn from set of FLOPs buckets (B), which partition the computational space between f and g. The bucket granularity controls the density of the resulting eccuracy-efficiene tradeoff curve. To ensure coverage, choose B such that each block the target anchor g has at least one stitch configuration represented in S. Formally, the candidate set S is defined as:
V..)) ((O-))) ST, v(Li) E E cO (1 (1,)") arg min ((..)) (b.))eb
S Ri; Ri DEB
RJ;
where j) denotes the stitch score for block pair (i.j). This formulation in Eq. 3 balances computational efficiency with stitch quality by ensuring both bucket-level coverage and threshold- based filtering. The complete KLAS algorithm presented in Alg. 3 in App. A
EXPERIMENTS AND EVALUATION
Please refer to App. B for details on the experimental setup. All results in this section are reported after finetuning the stitched network for 50 epochs with identical settings following SN-Net Pan et al. (2023). The end-to-end overhead of KLAS (ProbeNet + Stitched network finctuning) on 8xA40 NVIDIA GPUs is 16 hours for the Swin model family.
5.1 MEASURING SIMILARITY USING KI. DIVERGENCE
Tab. shows that KLAS significantly outperforms existing similarity measures, achieving area un- der accuracy-efficiene curve (AUC) of 0.8950, com- pared to 0.8345 for SN-Net and much lower values for MSE, CE, and DM. Notably, CKA, widely used for representational similarity, lags behind KLAS by over 8%. These results along with further experi- ments App- DepDDDia highlight the effectiveness of KL divergence in capturing both representational alignment and functional alignment, making it a superior choice for guiding stitch configuration selection.
SN-
Metric
CKA
MSE
CE
DM
KLAS (aurs)
AUC
0.8124
0.7564
0.8023
0.7642
0.8345
0.8950
Table 2: Comparison of similarity metries the Swin family. Reported values AUC scores of stitch configuration selection.
5.2 ANCHOR SELECTION USING KLAS
We demonstrate that last block KL divergence pro- vides a reliable criterion for selecting pretrained an- chors that are more compatible for stitching and more likely yield better accuracy-efficience tradeoff's (line 2 in Alg, 3). Since the softmax distribution is already computed in the final blocks of the anchors, additional probes and training are not required at this stage. As shown Tab. and Fig. (5) for both model families with more than two anchors (i.e., Swin and DeiT), last block KL divergence successfully identi- fies the anchor pairs that should be stitched together for improved performance.
Family
Anchors
KL Divergence
Ti-S
4.12 10-
DeiT
S-B
2.48 10-
Ti-B
5.62 x 10
Ti-S
2.70 3
Swin
S-B
x 10-3
Ti-B
2.38 10
Table 3: Last block KL divergence for anchor
lection in DeiT and Swin families (pretrained on ImageNet-22K)) Lower KL divergence values in- dicate greater compatibility between anchors for stitching (in blue).
In the Swin family, stitching Ti-B yields higher AUC compared to Ti-S and S-B stitches chosen by SN-Net. Conversely, within the DeiT family, the Ti-S and S-B stitches selected by SN-Net (via the nearest stitching heuristic) offer a good scurrcy-efffciencc tradeoff. These trends, as shown in Fig. [ demonstrate that the nearest stitching heuristic used in SN-Net is universally optimal across model families. In contrast, selecting anchors based on last block KL divergence offers a principled and generalizable strategy that consistently identifies the most compatible anchor pairs.
7

---

Published as a conference paper at ICLR 2026
82
Dola
Accuracy 185
80
DOAT5
Accuracy (S)
84
83
DelT
Nearast Stitching Swin TI-5-8)
Deil (5-8]
Far Stitching (Swin 11-9)
72
DelT (T-5)
80
Nearest Stitching (Swin 5-8)
70
Swin
Stitching
2.5
5.0
10.0
15.0
17.5
10 FLOPs (G)
FLOPs (G)
(a): DeiT
(b): Swin
Figure 5: Accuracy-efficiency tradeoff's anchor selection DeiT and Swin families. (a) For DeiT, stitching Ti-S and S-B (nearest stitching heuristic from SN-Net) yields the optimal accuracy-efficiency curve, while stitching Ti-B underperforms.
(b) Conversely, for Swin, stiching Ti-B (far stitching) provides superior tradeoff compared to Ti-S-B, which is chosen by SN-Net. These results show that the nearest stitching hcuristic is not universally optimal, whereas last block KL divergence (sce Tab. consistently identifies the most compatible anchors across families.
5.3 BLOCK SELECTION USING KLAS
After selecting the anchors, the next step is to identify block pairs within these anchors for the stitched network finetuning phase. As an initial experiment, we evaluate whether KLAS can automatically recover the paired and unpaired stitching heuristics proposed SN-Net. To this end, we restrict ourselves to the same anchors used in SN-Net, but perform block selection following lines 4-12 of Alg. 3 Tab. 4 shows that the block selection component of KLAS automatically recovers many of the stitch configurations identified by SN-Net's heuristics, though not perfectly. Crucially, when finetune stitched networks using the stitches selected by KLAS, our method outperforms SN-Net under this constrained setting (i.e., using the same anchors SN-Net), as shown in Fig. 6 A key goal of our KL divergence based approach is to automatically discover stitch configurations that heuristies have shown be effective prior work, while also identifying additional promising candidates. The results in Tab. 4 validate KL divergence as a viable similarity metric that can guide stitch discovery without reliance on ad-hoc heuristics.
Family
Method
Total Stitches
AUC
SN Net T1.56
TI4-8
SN-Net KLAS
68
0.731
DeiT
83-
SMi
(40 overlap)
0.743
Tos Accuracy
SN-Net
38 (28 overlap)
0.835
Swin
KLAS
0.842
Table 4: Comparison of SN-Net and KLAS when con- strained the nearest stitching heuristic (Swin: Ti-S-B; DeiT: Ti-S-B). Despite this restriction, KLAS achieves higher AUC by leveraging its KL divergence based block selection, while also recovering substantial portion of the heuristic stitches used by SN-Net.
30
12
14
16
FLOPS (G)
Figure 6: We fix the anchors as the ones selected in the SN-Net (Ti-S-B). We see even then KLAS returns better accuracy-officiency curve.
For our main experiments, we follow the setup described in App. B We use three DeiT anchors Tou- vron et al. (202T) (DeiT-Ti, DeiT-S, and DeiT-B) and three Swin anchors Liu et al. (2021) (Swin- Ti. Swin-S, and Swin-B), Each DeiT anchor has 12 blocks of equal spatial resolution, resulting in 432 possible stitches across the iccuracy-efficienee curve. For Swin, stitching is applied only within stages (as SN-Net), since spatial resolution preserved within a stage. Swin-Ti has 12 blocks, while Swin-S and Swin-B have 24 blocks each, leading to 1 1122 possible stitches in total. We attach probes at the end of every block, yielding 36 probe locations in DeiT and 60 in Swin. KL divergence then computed for all possible stitch pairs by treating the softmax output the stitch start point ImageNet-1 K as the source distribution and the corresponding softmax at the stitch end point as the target distribution.
Fig. 7 shows that our algorithm achieves a superior accuracy ecfuracy-efficiency curve across all model families in terms of AUC scores. Even in settings where the nearest stitching heuristic performs
8

---

Published as a conference paper at ICLR 2026
reasonably well (e.g., DeiT), KLAS performs better. We observe that while algorithm produces a sparser set of stitches at lower FLOPs in Swin, it still closely matches SN-Net in accuracy. The key advantage of our approach is that it does not impose fixed mapping between source and target blocks; instead, it adapts based on the pretrained anchor weights. Consequently, if the anchor weights change, the selected source-target block pairs may also change. Tab. highlighht the overlap between stitches selected by KLAS and those chosen heuristically by SN-Net, showing that KLAS successfully recovers many effective stitches while avoiding rigid reliance on fixed heuristies. key advancement is that KLAS also discovers new stitch configurations, most notably stitches from Swin- Ti to Swin-B anchors, that were not captured by heuristic approaches in the baseline. We provide further ablation studies on KLAS by comparing it to a naive Min-KL approach in App. C We discuss KI divergence correlations in App. D and block-level KI. divergence variations in App.
compare KLAS to cascades in App. E and present KI. divergence heatmaps
31.0
LOW 11200 35-VEE 18.04 EUGAN URC LATTTINN
RLASIN
25
5.
7.5
10.0
12.5
35.0
37.5
0.8
0.0
2.0
05
(a): Swin, AAUC +0.06
(b): DeiT, AAUC +0.012
(c): LeViT, AAUC +0.006
24
23
12
(d): ResNet, AAUC +0.014
(c): ResNet-SSin, (f): Swin CIFAR-100, (g): LeVIT CIFAR-100.
AAUC: +0.014
:AUC++0.00
Figure 7: Accuracy-effiienny tradeoff curves comparing KLAS with SN-Net across multiple model families and datasets: (a) Swin, (b) DeiT, (c) LeViT, (d) ResNet, (e) ResNet-Swin cross-architcctur,, (f) Swin CIFAR-100. and (g) LeViT on CIFAR- 100. Each point corresponds
stitched model after finetuning, plotted by FLOPs (x axis) and Top-1 accuracy (y-axis). Across all settings, KLAS achieves higher accuracy at comparable compute, yielding consistently positive AAUC. A positive AAUC C' indicates that KLAS outperforms SN-Net over the full FL.PPs space.
5.4 RESULTS ON DENSE PREDICTION TASKS
We further report final semantic segmentation results for Mask2Former Cheng et al. (2022) models with Swin-T and Swin-B backbones on ADE20K (hou et al. (2019). We assess the stitched models selected by SN-Net and KLAS. All models evaluated using standard mloU the ADE20K validation set; FLOPPs are measured per-image with input size 512 x 512. The Swin backbones are not frozen in the Mask2Former archi- tecture, while the rest of the model is frozen during finetuning. For each selection strategy, we divide the FLOPs space into three buckets of equal width. The entire stitched network is finetuned for 160k steps under the same schedule and batch size as in the original Mask2Former setup. We use the same steps outlined in $5.1 $5.3 run KLASS for the semantic segmen- tation task. KLAS consistently outperforms SN-Net across the FLOPs space as shown in Tab. For instance, KLAS achieves up to +0.9% mIoU, with comparable FLOPs to the SN-Net counterparts. In general, SN-Net selects lower quality stitch buckets. These results confirm that KLAS enables more effective iccuracy-efficiees tradeoffs than SN-Net for dense prediction tasks like semantic segmentation. We provide more details on the implementation in App.
Model
FLOPs(G)
(l)
SN-Net-1
152
29.4
KLAS-1
145
29.8
SN-Net-2
274
32.6
KLAS-2
277
33.5
SN-Net-3
327
37.7
KLAS-3
316
37.8
Table 5: Semantic segmentation results on ADF20K using Mask2Former with Swin-T and Swin-B. KL.AS-selcctdd stitches outperform SN-Nee baselines at similar FLOPs budgets.
configurations in the same FLOPs

---

Published as a conference paper at ICLR 2026
5.5 APPLICABILITY TO LARGE LANGUAGE MODELS (LLMS)
To evaluate whether KLAS generalizes to instruction-tuned LLMs, we replicate stitch-
Model
Method
ROUGE-1
ROUGE-2
Stitched LLaMa 1.6
ESTA
0.576
0.304
ing experiment based on ESTA He et al. (2024),
Stitched LLMMa 1.4 B
KLAS
0.593
0.337
comparing it against KLASS-ased selection
Stitched LLaMa 2.7 B
ESTA
0.631
0.353
ing TruthfulQ Lin et al. (2022). We use Llama 3.2 IB and 3B Grattafiori al. (2024) anchor models. Stitched models are created by com- bining source blocks from Llama 1B target blocks from Llama 3B. Stitching is done using two strategies: (1) ESTA, which selects posi-
Stitched LLaMa 2.6 B
KLAS
0.645
0.379
Table 6: Results on LLMs using Llama 1B and 3B. KLAS-sclected stitehed models improve over ESTA- selected stitched models, with fewer parameters.
tions based on layer count (e.g.. 1B first & layers 3B remaining), and (2) KLAS, which selects based on minimum stitch score using per-layer next-token distributions using linear probes. Evaluation is performed in the generative setting on TruthfulQA, and we report ROUGE-1 and ROUGE-2 scores against reference completions after finetuning for 10 epochs in both strategies. Tab. 6 shows that KLAS consistently identifies more similarly aligned stitch points (i.e., blocks) than ESTA's heuristic- based selection. On TruthfulQA, KLAS-stitches models outperform similar-sized ESTA-stitched models by up 0.017 ROUGE-1 and 0.033 ROUGE-2, thus offering approach to designing stitched LLMSs. We provide more experimental details in App.
principled and effective
5.6 ABLATION STUDIES ON KLAS HYPERPARAMETERS
Threshold (r). KLAS selects stitch configurations by thresh- olding the stitch score [i relative to the minimum score within each FLOPs bucket. We sweep 7 from 1% to 10% more than this per-bucket minimum and report the resulting average accuracy and AUC Tab. 7 A clear tradeoff emerges: very high thresholds (e.g., 7 10%) admit noisy, low-quality stitch configurations, while low thresholds (e.g.. 7 = 1%) become overly selective and reduce the number of selected stitches, yielding sparse Pareto fronts. The default = 5% provides the best balance between robustness and diversity. We use Swin-Ti, Swin-S and Swin-B with 20 buckets for these experiments.
Threshold (t)
Avg Top-1(%)
AUC
1%
83.72
0.8934
3%
83.74
0.8942
5%
83.76
0,8950
10%
83.69
0.8931
Table 7: Effect of threshold r on KLAS stitch selection. High admits noisy pairs; low causes sparsity.
Bucket Granularity. To understand the effect of discretiz- ing the FLOPs space, we vary the number of FLOPs buckets used during stitch selection. The block granularity is inversely proportional to the number of buckets. As shown in Tab. using finer buckets (i.e., lower bucket granularity) increases the density of Pareto front, but KLAS maintains strong per- formance even with coarser buckets (i.c., higher bucket gran- ularity). Overall sensitivity remains low for any reasonable bucket count (10-20). We use Swin-Ti, Swin-S and Swin-B with 5% for these experiments.
#Buckets
Avg Top-1(%)
AUC
10
83.65
0,8902
15
83.75
0.8947
20
83.76
0.8950
Table 8: Ablation FLOPs bucket granularity. KLAAS remains robust across bucket granularity choices.
DISCUSSION AND CONCLUSION
This paper introduces KLAS, a novel, efficient and automatic framework for improving the accuracy- efficiency tradeoff produced by stitching pretrained models. The method leverages KL divergence to identify stitched configurations that are likely to perform well after finetuning the stitched super- network. Empirical evaluations show that KLAS consistently outperforms heuristic-based baselines by providing principled, similarity- selection strategy that generalizes across diverse model architectures, datasets, tasks and modalities, while incurring negligible additional computational cost. Beyond accuracy alone, generative workloads impose distinct system constraints across the prefill and decode phases (e.g., toie-tt-ffistttoeee and time- time-between-e between-tokenss. This creates a new optimization problem: balancing accuracy with two latency budgets (and often memory throughput constraints) across both phases. natural future direction is to extend KLAS to this setting by enforcing KV-eache compatibility across stitched blocks and redefining similarity measures accordingly.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This material is based upon work partially supported by the National Science Foundation (NSF) under Grant Number CNS-2420977 as well sponsored research awards from Cisco Research and Georgia Tech Research Institute (GTRI. This work was further supported in part by CoCoSys, one of seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA. We would also like to express sincere gratitude the reviewers and the AC panel for their insightful comments and thoughtful considcration. Disclaimer: Any opinions, findings, and conclusions recommendations expressed in this material those of the authors and do not necessarily reflect the views of the National Science Foundation.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. {TensorFlow)) J: a system for (Large- } machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265-283, 2016.
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes, 2018.
Xiang An, Jiankang Deng, Jia Guo, Ziyong Feng, XuHan Zhu, Jing Yang, and Tongliang Liu. Killing two birds with one stone: Efficient and robust training of face recognition cnns by partial fc. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4042 4051, 2022.
Aditya Annavaijala, Alind Khare, Animesh Agrawal, Igor Fedorov, Hugo Latapie, Myungjin Lee, and Alexey Tumanov. ('epsilon) ps: Delayed ("eepsioo))-shit -shrinking for faster once-for-all training. arXiv preprint arXiv:2407.06167 2024.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arxiv. arXiv preprint arXiv:/803.01271, 10, 2018.
András Balogh and Márk Jelasity. How not to stitch representations to measure similarity: Task loss matching versus direct matching. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 15472-15480, 2025.
Yamini Bansal, Preetum Nakkiran, and Boaz Barak. Revisiting model stitching to compare neural representations. Advances neural information processing systems, 34:225-236, 2021.
Tolga Bolukhasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In International conference on machine learning, pp. 527-536. PMLR, 2017.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once for all: Train one network and specialize it for efficient deployment. In International Conference on Learning Representations, 2020. htLp:t /BLL/tttps//aavvv.oorrt///////// 0ttps//pd//110.09991...
Chien-Yint Chen, Debopam Sanyal, and Sibin Mohan. Indistinguishability prevents scheduler side channels in real-time systems. In Proceedings of the 2021 ACM SIGSAC Conference on Computer and Communications Security, pp. 666-684, 2021.
Bowen Cheng, Ishan Misra, Alexander Schwing, Alexander Kirillov, and Rohit Girdhar. Masked- attention mask transformer for universal image segmentation. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 1290-1299 1299, 2022.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
11

---

Published as a conference paper at ICLR 2026
Adrián Csiszárik, Péter Korösi-Szabó, Akos K Matszangosz, Gergely Papp, and Dániel Varga. Similar- ity and matching of neural network representations, 2021. URL https//arxiv. org/abs/2110.14633, 2021.
Robert David, Jared Duke, Advait Jain, Vijay Janapa Reddi, Nat Jeffries, Jian Li, Nick Kreeger, lan Nappier, Meghma Natraj, Tiezhen Wang. et al. Tensorflow lite micro: Embedded machine learning for tinyml systems. Proceedings of machine learning and systems, 3:800-811, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp.
leee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidi- rectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1929,, 2020.
Ben Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, and Matthijs Douze. Levit: a vision transformer in convnet's clothing for faster inference, 2021.
Aaron Grattafiori et al. The Ilama herd of models. arXiv preprint arXiv:2407.21788 2024. 20244
Jiaqi Guan, Yang Liu, Qiang Liu, and Jian Peng. Energy-efficient amortized inference with cascaded deep classifiers. arXiv preprint 1rXiv:1710.03..0 2017.
Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun. Single path one-shot neural architecture search with uniform sampling. In European conference computer vision, pp. 544-560. Springer, 2020.
Kim Hazelwood, Sarah Bird, David Brooks, Soumith Chintala, Utku Diril, Dmytro Dzhulgakov, Mohamed Fawzy, Bill Jia, Yangqing Jia, Aditya Kalro, et al. Applied machine learning at facebook: A datacenter infrastructure perspective. In 2018 IEEE international symposium on high performance computer architecture (HPCA), pp. 620 629. IEEE, 2018.
Haoyu He, Zizheng Pan, Jing Liu, Jianfei Cai, and Bohan Zhuang. Efficient stitchable task adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
28555 28565, 2024.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Ame: Automl for model compression and acceleration on mobile devices. In Proceedings of the European conference on computer vision (ECCV), pp. 784 400, 2018
Bodun Hu, Shuozhe Li, Saurabh Agarwal, Myungjin Lee, Akshay Jajoo, Jiamin Li, Le Xu, Geon-Woo Kim, Donghyun Kim, Hong Xu, al. Stitchllm: Serving llms, one block at time. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pp. 26887-26903, 2025.
Hiroshi Inoue. Adaptive ensemble prediction for deep neural networks based on confidence level. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 1284-1299 PMLR, 2019.
Aryan Jadon, Avinash Patil, and Shruti Jadon. A comprehensive survey of regression- based loss functions for time series forecasting. In International Conference on Data Management, Analytics & Innovation, pp. 117-147. Springer, 2024.
12

---

Published as a conference paper at ICLR 2026
Xiaoyu Ji, Yushi Cheng, Yuepeng Zhang, Kai Wang, Chen Yan, Wenyuan Xu, and Kevin Fu. Poltergeist: Acoustic adversarial machine learning against cameras and computer vision. In 2021 IEEE symposium on security and privacy (SP), pp. 160-175. IEEE, 2021.
Alind Khare, Animesh Agrawal, Aditya Annavajjala, Payman Behnam, Myungjin Lee, Hugo Latapie, and Alexey Tumanov. Superfednas: Cost-efficient federated neural architecture search for on- device inference. In European Conference Computer Vision, pp. 161-179. Springer, 2024.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019.
Alex Krizhevsky. Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Department of Computer Science, University of I'Toronto, 2009.
Luzian Lebovitz, Lukas Cavigelli, Michele Magno, and Lorenz K Muller. Efficient inference with model cascades. Transactions on Machine Learning Research, 2023.
Karel Lenc and Andrea Vedaldi. Understanding iegp represseeeeeetssssttttttossss by measuring their equivariance and equivalence, 2015. URL httpppsss/aaa.c...oor/aaa///////
Guihong Li. Yuedong Yang. Kartikeya Bhardwaj, and Radu Marculescu. Zico: Zero-shot nas via inverse coefficient of variation on gradients. arXiv preprint arXiv:2301.11300 2023.
Yanjing Li, Sheng Xu, Baochang Zhang, Xianbin Cao, Peng Gao, and Guodong Guo. Q-vit: Accurate and fully quantized low-bit bit vision transformer. Advances in neural information processing systems, 35:34451-34455 2022.
Ming Lin, Pichao Wang, Zhenhong Sun, Hesen Chen, Xiuyu Sun, Qi Qian, Hao Li, and Rong Jin. Zen-nas: A zero-shot nas for -eg--perforance image recognition. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 347-356, 2021.
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 3214-3252. Association for Computational Linguistics, 2022.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search, arXiv preprint arXiv:/806.0000 09055, 2018.
Ze Liu, Yutong Lin, Yue Cao. Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and paltern recognition, pp. 11976-1 1986, 2022.
Anqi Mao, Mehryar Mohri, and Yutao Zhong. Cross-entropy loss functions: Theoretical analysis and applications. In International conference on Machine learning, pp.
PMLR, 2023.
Zhenchao Ouyang, Jianwei Niu, Yu Liu, and Mohsen Guizani. Deep cnn-based real-time traffic light detector for self-driving vehicles. IEEE transactions on Mobile Computing, 19(2):300-315 2019.
David Oygenblik, Abhinav Vemulapalli, Animesh Agrawal, Debopam Sanyal, Alexey Tumanov, and Brendan Saltaformaggio. Villainnet: Targeted poisoning attacks against supernets along the accuracy-latency pareto frontier. In Proceedings of the 2025 ACM SIGSAC Conference on Computer and Communications Security, pp. 2189-2200 2025.
Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai. Scalable vision transformers with hierarchical pooling. In Proceedings of the IEEE/evf international conference on computer vision, pp. 377-386, 2021.
Zizheng Pan, Jianfei Cai, and Bohan Zhuang. Stitchable neural networks. In CVPR, 2023.
Zizheng Pan, Jing Liu, Haoyu He, Jianfei Cai, and Bohan Zhuang. Stitched vits are flexible vision backbones. In European Conference on Computer Vision, pp. 258-274. Springer, 2024.
13

---

Published as a conference paper at ICLR 2026
Eunhyeok Park, Dongyoung Kim, Soobeom Kim, Yong- Deok Kim, Gunhee Kim, Sungroh Yoon, and Sungjoo Yoo. Big/little deep neural network for ultra low power inference. In 2015 international conference hardware/software codesign and system synthesis (codes+ isss), pp. 124-132. IEEE, 2015.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Matthew E. Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. ArXiv, abs/1802.05365 2018.
Hieu Pham, Melody Guan, Barret Zoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International conference on machine learning, pp. 4095-4104, PMLR, 2018.
Yongming Rao. Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 34:13937-13949, 2021.
Nikhila Ravi, Valentin Gabeur, Yuan-Ting Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, et al. Sam 2: Segment anything in images and videos. arXiv preprint arXiv:2408. 00714. 2024.
Joseph Redmon, Santosh Divvala, Ross Girshick, and Ali Farhadi. You only look once: Unified, real-time object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 779-788, 2016.
Manas Sahni, Shreya Varshini, Alind Khare, and Alexey Tumanov. Comp/ofal compound once- for-all networks for faster multi-platton deployment. International Conference on Learning Representutions, 2021.
/nne/fforume /ttp//orrrmmmmiIII1k1k181--
Debopam Sanyal, Jui-Tse Hung. Manav Agrawal, Prahlad Jasti, Shahab Nikkhoo, Somesh Jha, Tianhao Wang, Sibin Mohan, and Alexey Tumanov. Pareto-secure machine learning (psml): Fingerprinting and securing inference serving systems. arXiv preprint arXiv:2307.01299 2023.
Divyansh Srivastava, Akshay Mehra, Pranav Maneriker, Debopam Sanyal, Vishnu Raj, Vijay Ka- marshi, Fan Du, and Joshua Kimball. Dpar: Dynamic patchification for efficient autoregressive visual generation. arXiv preprint arXiv:2512.21867, 2025.
Matthew Streeter. Approximation algorithms for cascading prediction models. In International conference on machine learning, pp.
PMLR, 2018.
Hugo Touyron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention, 2021.
Paul Viola and Michael Jones. Rapid object detection using a boosted cascade of simple features. In Proceedings of the 2001 IEEE computer society conference on computer vision and pattern recognition. CVPR 2001, volume pp. I-I. Ieee, 2001.
Dilin Wang, Meng Li, Chengyue Gong, and Vikas Chandra. Attentivenas: Improving neural architecture search via attentive sampling, 2021.
Xiaofang Wang, Dan Kondratyuk, Eric Christiansen, Kris M Kitani, Yair Alon, and Elad Eban. Wisdom of committces: An overlooked approach to faster and more accurate models. arXiv preprint arXiv:2012.01988, 2020.
Xin Wang, Yujia Luo, Daniel Crankshaw, Alexey Tumanov, Fisher Yu, and Joseph E Gonzalez, Idk cascades: Fast deep learning by learning not to overthink. arXiv preprint arXiv:/706 1776600088. 2017.
Wolf. Huggingface's 's transformers: State-of oftte-of-theaan the-art natural language processing, arXiv preprint 190.731910 2019.
14

---

Published as a conference paper at ICLR 2026
Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly. Advances in neural information processing systems, 015-25739-257555 2022a.
Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly, In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 25739-25753. Curran Associates, Inc., 2022b. URL ELLAGS11ccostTTaaaaea888221121e0
File 
pdf
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances neural information processing systems, 27, 2014.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision, 122(3):02--321, 2019.
ZipLab. Sn-net: Stitchable neural networks. http:s//ghthu..cmmmz.paaaaa -et 2023. Accessed: 2025-05-16.
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv:1611.0157,, 2016.
15

---

Published as a conference paper at ICLR 2026
A ALGORITHMS
Algorithm ProbeNet Pseudocode (see $4.1)
1: Class ProbeNet
2: Init(model, stage_block block info, output_ddi dim 1000):
3: Store model as anchor (frozen)
4: for each stage in stage_ _lloc_innooo
5: for each block_din in stage:
6:
Create probe: Lincar(block dim + output_dimy + Softmax Store probe in self lel..layys(sttgellloockk
7:
8:
Set current features + None
9:
Freeze anchor parameters; keep probes trainable
10:
Method extract_block (fatuc__eaturee(()::
12:
Run catbooreextrat_blcccce feretuf_fatre(())) with no_grad Store result in current features
14:
Method clear block_ features(): Set current_featuree +- None
16:
18:
Method iter_ blockslock():
19:
yield (stage_id, block_id) for each stage/block
20:
21:
Method forward(x, stage id, block id, get all probe losses, criterion, target): if wwpppooooossseeee
22:
23:
return if current_feature is None:
self. layer, criterion, target, dims)
24:
25:
extract cxtract_block__eeee_
Get block feature + reshape pass through probe
27:
return output
Algorithm 2 Get Similarity Scores (see $4.2)
Inputs: Set of all anchor models M
Output: Set of all similarity scores K.
1: K. + ()
2: for each model m in M do
3:
for each block in m do
4:
Add linear classifier probe P.
5: Train pm on train split (D,)
6: end for
7:
end for
8: for each model pair (nn, n) in M, with m having lower complexity than do
9:
for each block anchor m do
10:
for each block j in anchor n do - ),
11:
12: end for
13: end for
end for
15: Return KC
B EXPERIMENT SETUP
Baselines. We first compare KLAS with existing stitching methods for ofate-of-the-art image classification models w.r.t. accuracy-efficiency tradeoffs. Specifically, compare KLAS with the
16

---

Published as a conference paper at ICLR 2026
Algorithm 3 KLAS: KL divergence based Anchor Stitching (see $4.2)
Inputs: Anchor models M, similarity scores K. from Alg. threshold 7, buckets B
Output: Set of selected stitch configurations S
1: S +-0
2: Pick anchors g E M based on their final layer KL divergence
3: Sort all the stitches based on FLOPs
4: for each bucket E B do
5: Rs ()
6: for each stitch configuration (i.)) do
7:
Compute I ( 5) (see Eq. 2
8: R₂ + R2 (1,
9: end for
10: Compute set of stitch configs RT using Rs (see Eq.
S S-SU R2
12: end for
13: Return
primary baseline SN-Net Pan et al. (2023), which introduced heuristic stitching techniques including nearest stitching and paired unpaired stitching. The finetuning costs of all techniques are 50 epochs and we use all the same parameters and hyperparameters for finetuning as in Pan et al. (2023) to enable fair comparison. We run all experiments on 8x A40 NVIDIA GPUs.
Success Metric. KLAS is compared against the SN-Net baseline on the following success metric, accuracy-efficiene, tradeoff': accuracy of selected stitched models as function FLOPs. To compare tradeoff spaces obtained from different baselines, we use the area under the curve (AUC) of accuracy and normalized FLOPs. The higher the AUC, the better. Additionally, this metric measures the stitch configuration quality by evaluating how effectively each method selects stitch points that are likely perform well.
Datasets. We evaluate all methods on CIFAR-100 and ImageNet- 1 datasets. The complexity of datasets increases from CIFAR-I00 to ImageNet- 1k. The datasets vary in the number of classes, image resolution, and number of train / test samples. The stitching models and the probes were trained on the train splits and the KL divergence values in Alg. 2 were obtained using the validation splits of the datasets. The train and validation splits were obtained from the official train set of ImageNet- K. We report the final accuracies on the official validation set of ImageNet- FLOPs proxy for the efficiency of the deployment point (i.e., stitched model), as is widely used in NAS works farks/Oygenblik et al. (0255) y Annavajjala et al. (024): Khare et al. (2024): Sanyal et al.
We use
(0233.
Model Architecture Families. All experiments are conducted using pretrained models from several different architecture families: transformers including DeiT Touvron et (202I), Swin Liu al. (2021), LeViT Graham et al. (2021), CNNs including ResNet He et al. (2016). and LLMs including Llama Grattationi al. (2024). ). All anchor models are pretrained on variant of the dataset being used for finetuning the stitched network using standard training protocols outlined in Pan et al. (2023). We use the ResNet and Swin architectures for our experiments on cross- architectee stitching. To ensure fair comparison, KIAAS and all baselines use identical pretrained weights for the corresponding anchor models, and stitching configurations are evaluated using the same batch sizes as Pan ct al. 2023).
Hyperparameters. The hyperparameters for each model were set as defined in their respective papers (Liu et al. 2021 Touvron et al. 2021 Graham et al. 202 He et al. 2016). The stitched network finetuning hyperparameters were adapted from the SN- codebase ZipLab (202)) We set the threshold (r) to 5% more than the minimum stitch score (I') in the bucket. We assign the bucket granularity used for creating the set of buckets (B3) such that every block in the target anchor has at least representative stitch configuration in S. This gives 20 buckets for the Swin model family.
17

---

Published as a conference paper at ICLR 2026
C MIN-KL APPROACH VS. KLAS
While KL divergence is useful metric for identifying promising stitches, using Min-KL strategy of directly selecting the stitches with the lowest KL values and training them suboptimal. The key issue is that KL divergence is not symmetric, and stitches with different target blocks are not directly comparable. To address this, stitches must be evaluated addresses this by selecting the stitch with the lowest similarity score I' (see Eq. ( for each bucket, ensuring both faimess across buckets target blocks and balanced coverage of stitched models across the FLOPs range. As shown in Fig. this principled strategy outperforms the naive Min-KL. baseline, which underperforms even compared to SN-Net on DeiT' and Swin.
target-award aware manner. KLAS
82
80
Swi swinn-Tl
eglo Swin-5
SoiinBB
Parants (N)
40
80
DaITS
60
78
60
80
Top 1(%)
76
40
20
72
SN-Net TI--8
50
70
2.5
5.0
7.5
10.0
12.5
15.0
17.5
5.0
10.0
12.5 FLOPS (G)
15.0
17.5
20.0
FLOPs (G)
(a): DeiT
(b): Swin
Figure 8: Min-KL algorithm suboptimal because naively choosing stitch configurations with lowest KL divergence values is not target-aware across the FLOPs space. KL divergence is not a symmetric metric and hence all candidate stitch configurations across the FLOPs space cannot be compared to each other directly.
D A DEEPER DIVE INTO KL DIVERGENCE
D.1 WHY KI. OUTPERFORMS MSE, CE, CKA, AND DM FOR STITCH SELECTION
We attribute KI. divergence's superior stitch configuration recovery (Tab. T shows it recovers 88.9% of viable Swin Ti-B stitch configs vs. 22-33% for MSE/CE/DM and only 5.5% for CKA) its ability capture both representational and functional alignment. Let (z((2) and (<)(x) be intermediate features from source and target models. Let p/ (x) and (x(x) be softmax probability distributions from linear probes. The similarity metrics are defined follows:
MSE:
1
Σ
(155((z) (x)1
Dr
CE:
Σ
log P/ (x)[p(>)
DM:
min T
Σ
IT (s) (x)
G
CKA:
FT
F
F
Metric
Pearson (to AAec)
Spearman
Stitch score (T)
0.84
0.81
MSE
0.42
0.37
CKA
0.28
0.21
Table 9: Pearson and Spearman correlations between candidate similarity metrics and stitched model accuracy drop (AAcc). The stitch score (based on KL) exhibits the strongest predictive correlation, confirming its utility for selecting performant configurations (evaluated Swin family).
18

---

Published as a conference paper at ICLR 2026
To measure how well different similarity metrics prediet stitch configuration quality, we compute how each metric score relates to the final stitched model's ' accuracy. For every stitch pair, we calculate the accuracy drop compared to the target model as AAcc Accuracy
Accuracysstchhed We then compute two correlation scores between the similarity metric and AAcc: Pearson correlation (r) measures linear relationship, and Spearman correlation (p) measures whether ranks agree (i.e., monotonic trend). Results shown in Tab. 9 divergence between probe distributions) has the strongest correlation with accuracy drop. indicating that stitch candidates with low stitch score likely perform better after finetuning. In contrast, MSE and CKA show much weaker correlation. This helps us conclude that I' shows a near-monotonic increase: as the score rises, so does the accuracy drop, while CKA and MSE exhibit noisier and less predictive patterns.
The stitch score (T from Eq. (2) (computed using KL
D.2 KL DIVERGENCE IS BOTH REPRESENTATIONAL AND FUNCTIONAL
To show that the KL divergence based stitch score (T) captures both representational and functional similarity, we run two additional experiments:
Shuffle-label Stitch Score: We retrain the probes with randomly shuffled labels. This removes class semantics while preserving feature structure. The resulting scores lose correlation with task performance, similar to unsupervised metrics.
2. Class- Conditional CKA: We compute CKA separately per class, then average across classes to inject label information. This improves correlation modestly but still underperforms w.r.t. stitch score. The stitch score outperforms other metrics because it reflects both representational structure and functional behavior, as shown in Tab. I0 Supervision is essential; without it, metric utility drops the level of unsupervised baselines.
Variant
Pearson
Spearman P
Stitch score (T)
0.84
0.81
Shuffled-label score
0.19
0.16
CKA (vanilla)
0.28
0.21
CKA (class-conditiona))
0.51
0.46
Table 10: Correlation of stitch performance with label-aware and label-free metrics. Removing supervision (shuffled labels) weakens KL-based stitch score, while adding class-conditioning to CKA improves its stitch ranking. These results confirm that supervision critical for good stitch selection.
E KLAS AND SN-NET VS. MODEL CASCADES
We compare KLAS and SN-Net to early-exit cas-
caded inference baselines adapted from Lebovitz et al. (2023). Since in our work we only con- sider binary stitches, in our setup, we have just two starting models. Inference begins with small model (Swin-S), and proceeds to a larger model (Swin-B) only if the small model's 'ss confi- dence (maximum softmax) falls below thresh- old. Thresholds are tuned to optimize the accuracy-FLOPs Pareto front, and the average inference cost is measured in FLOPs over the full ImageNet-IK validation set. The cascaded models are not finetuned and share no parame- ters. For our experiment, we use their method with the following settings: (i) input resolution is 224 x 224 and (ii) routing is based on the max softmax score. We sweep thresholds in the range [0.6. 0.95) and report the best value per FLOP bin. For example, thresholds of 0.85, 0.90, and 0.95 correspond to mean budgets of 10.1G, 11.3G, and 12.6G FLOPs, respectively.
Swi
Top-1 Azsuaay
84.0
83.6
SN Net (Swin S B)
(Swh 5-8)
83.0
Swhs
Coscado
FLOPs (B)
Figure 9: Accuracy-FLOPP tradeoff on ImageNet-1K K for Swin-S and Swin-B. KLAAS and SN-Net stitched models outperform cascaded inference baselines that use softmax-threshold early exits. FLOPs for cas- cades are average inference cost over the validation set. FLOPPs for KLAS are deterministic by design.
19

---

Published as a conference paper at ICLR 2026
As shown in Fig. both KLAS and SN-Net produce stitched models that are consistently superior to cascaded inference across the full FLOPs spectrum. For instance, at > GFLOPs, KLAS improves over the best cascade by >0.4% Top-1 accuracy. This trend holds from low high FLOPs regimes. Despite the adaptíve nature of cascades, KLASS-ased stitched models yield stronger Pareto front performance with two starting models for reasons: (1) stitching enables model restructuring by stitching blocks with similar learned representations into a unified network, rather than deferring work across the smaller and larger models; and (2) the well-known problem with cascades that samples are routed based on softmax confidence, which is poorly calibrated in low-entropy regions, leading to suboptimal exits Wang al. (2020); Lebovitz al. (2023). In contrast, stitching builds single, end-to-end architecture that avoids runtime uncertainty, yielding higher accuracy per FLOP even in the average-case, as shown in Fig. 9 is easy to see that in the worst-case scenario, stitching will perform better than cascades for this particular setting (i.e., with two models).
We further evaluate a cross-family pair consisting of a CNN and a Transformer: ResNet-18 (1.8 GFLOPs) and Swin-B (15.4 GFLOPs). We construct three KLAS stitched configurations whose FLOPs are close to those of three corresponding early-exit cascades, tuned with the same softmax- threshold procedure as in Lebovitz et al. (003). Tab. 11 shows that even when the two anchors are architecturally different, KLAS consistently outperforms cascaded inference at similar or lower FLOPs, improving Top-1 accuracy by 2-3% across all three operating points.
Config
Method
FLOPs (G)
Top-1 (%)
KLAS
3,409
79.60
Cascade
3.455
76.13
KLAS
15.043 15.197
84.52
2
Cascade
82.46
KLAS
17.127
84.74
3
Cascade
17.183
82.79
Table 11: KLAS vs. carly-exit cascades for a cross-family pair (ResNet-18 - Swin-B) on ImageNet- IK. KLAS outperforms cascaded inference Lebovitz al. (2023) at similar or slightly lower FLOPs across all three operating points. FLOPs for cascades are average inference cost over the validation set. FLOPs for KLAS are deterministic by design.
We would like to further draw attention to the fact that the FLOP numbers reported for cascades in Fig. 9 and Tab. [T correspond to average cost over the validation set. The compute a cascade not deterministic: it depends on how often examples fall below the confidence threshold and are forwarded to the second model. In practice, many inputs remain below the threshold, so cascades execute all the blocks of the first anchor and all the blocks of the second one. This often yields only modest accuracy gains while spending extra FLOPs on unnecessary blocks. By contrast, once KLAS stitched configuration is selected and finetuned, the model executes a fixed sequence of blocks for every input, leading to deterministic and easily budgetable FL.OP cost for all samples.
The KL. divergence heatmaps in Fig. further clarify why stitching selects different configurations than cascades. A pure cascade that always runs the full small model and then the large model corresponds to connecting the last block of the small anchor to the first block of the large anchor, i.e., the upper-left corner in each heatmap. These block pairs have among the highest KL divergence values, indicating poor representational and functional compatibility. KLAS explicitly searches for low KL divergence pairs and therefore prefers connections in the darker regions near the diagonal. As a result, such "top-left" cascade-style configurations are never chosen: they are dominated by lower KL divergence stitches that achieve better accuracy- FL.Ps tradeoffs than cascades.
F KL DIVERGENCE HEATMAPS
KL Divergence heat maps provide a visual representation of the similarity between intermediate blocks of two anchor models. Each entry in the heat map corresponds to the KL divergence between the softmax distributions of a source block and a target block, as measured by linear probes. Low values (darker regions) indicate block pairs with higher compatibility for stitching, while higher
20

---

Published as a conference paper at ICLR 2026
Stage
3230
stage?
80000
Tiny Gaie Stage
3331
3000
0000
(a): Swin Ti-S
(b): Swin S-B
(c): Swin Ti-B
Dell Timp Sersal HL Diverganece
Sacak
Biceryann
Deet They Sanee Divergance
2011
X033
X025
3000
X009
3007
3033
X03
2004
2802
XCE3
(d): DeiT Ti-S
(c): DeiT S-B
(f) DeiT Ti-B
Figure 10: Heatmaps of KI. divergence values across different block pairs Swin Stage (top row) and DeiT (bottom row) anchor combinations. Each cell represents the divergence between the output probability distributions of source block (x-axis) and target block (y-axis). Lower values (copler colors) indicate higher compatibility between block pairs, making them more promising candidates for stitching.
values suggest poor alignment. These heat maps allow us to identify promising stitch configurations glance.
In Fig. 10 we observe that the KL divergence values are generally lower when stitching between blocks of similar depth across anchors, as shown by the darker regions near the diagonals in most heatmaps. This indicates that earlier blocks align better with earlier blocks and later blocks with later blocks. However, this does not exactly give us paired stitching, heuristic used in SN-Net. Additionally, in both Swin and DeiT, larger anchors (e.g.. Base) tend to introduce higher divergence when paired with earlier blocks of smaller anchors, highlighting their greater representational shift.
Fig. 11 visualizes block-level KI. divergence across different anchor pairs in Swin Ti-B Stage clear trend emerges: divergence is lowest along regions where blocks of similar depth are compared, suggesting that carly-to-early and late-to-late block pairings naturally more compatible for stitching. In contrast, pairing shallow blocks from smaller anchors with deeper blocks from larger anchors results in higher divergence, reflecting greater representational mismatch. These structured patterns validate that KL divergence effectively captures cross cross-anchor anchor compatibility and provides a principled signal for selecting good block pairs.
G DETAILS ON THE SEMANTIC SEGMENTATION EXPERIMENT
A key difficulty arises from the architectural expectations of dense prediction models, which typically rely on backbones such ViT's and ResNets. These models are integrated with additional components (e.g., pixel decoders, transformer decoders, detection heads) that require activations of specific dimensions Cheng et al. (2022): Ravi et al. (2024); Redmon et al. (2016). contrast, SN-Net introduces stitched backbones and anchors with varying intermediate activation sizes, which creates mismatch in feature dimensions and prevents straightforward training. We address this activation- dimension inconsistency using stage-specific linear maps. Promising directions include the design of flexible mapping layers or adaptive projection mechanisms that can align heterogeneous feature representations across stitch configurations. Another avenue to explore more efficient training
21

---

Published as a conference paper at ICLR 2026
Differences By Layer, ti-b, Stage 2, input
KL Differences By Layer, ti-b, Stage 2, input
60000
30219
6002
32829
9407
0M0
183
(a): Block 0
(h): Block
KL Differences By Layer, ti-b, Stage input 2
KL Differences By Layer, ti-b, Stage 2, input
40035
80020
Adsta
1005
1000
K00
(c): Block
(d): Block
KL Differences Layer, ti-b, Stage 2, input
00.00
(e): Block 4
Figure 1:: KI. divergence values across different target layers when stitching from Swin-Ti to Swin-B in Stage Each sub-figure shows the KI. divergence distribution for target block inputs. The plots highlight that divergence varies systematically across layers, with certain source-target block pairs yielding significantly lower divergence values, indicating better stitching compatibility.
hlock (Blocks 0-4) with different
strategies that minimize additional computational overhead while maintaining compatibility with downstream dense prediction heads.
The original M2F architecture uses hierarchical transformer- based backbone, where feature maps from each stage are fed into a pixel decoder and a transformer decoder. When replacing the backbone with stitched Swin backbone, the M2F architecture expects feature maps with consistent dimensions. However, since stitched networks combine layers from multiple Swin variants (e.g., Tiny, Small, Base), their intermediate feature dimensions may vary across stages in the models. This necessitates careful design modifications to ensure compatibility. Our stitching and finetuning procedure is as follows:
22

---

Published as a conference paper at ICLR 2026
1. Stitch layers initialized using inverse based source & target feature dimensions.
2. Four additional linear transformation layers (randomly initialized) are inserted (one per stage) to project stitched outputs to match expected shape of Swin-S or Swin-B each stage, since only these two can be target anchors we fix the direction in stitching.
3. The M2F head (pixel + transformer decoder) is kept frozen during training to minimize overhead.
4. On each mini- random stitch configuration is sampled.
5. Full forward pass through the SN-Net backbone and M2F head.
6. Compute segmentation loss (cross entropy + auxiliary losses).
7. Update backbone + linear mapping weights based on loss gradients.
There a key challenge in enabling dimensionally consistent feature flow from stitched backbone to decoders. This is traced to the fact that stitch configurations do not always end with the same final feature dimensionality (e.g., 384, 768, or 1024 channels depending on the final layer in the sampled path). The decoders expect fixed final feature dimensions. We solve this using stage-specific linear maps, which are randomly initialized.
H DETAILS ON THE LLM STITCHING EXPERIMENT
For all LLM stitching experiments in $5.5 we adopt an identical finetuning setting derived from the ESTA setup He et al. (2024) to ensure that the two methods are comparable. We stitch between Llama 3.2 instruction suttted 1B and 3B anchors Grattafiori et al. (0024) and evaluate performance in the generative setting TruthfulQA benchmark Lin al. (2022). All stitched models in both ESTA and KLAS are trained under identical hyperparameters: 10 epochs finetuning of next-token prediction with cross-entropy loss, batch size 64, learning rate 5 10-5 with cosine decay, 1-epoch warmup period, Adam W optimizer, and a 128-token sequence length.
The training and test datasets are not identical. We use the generative setting of TruthfulQA benchmark and apply 60/40 split: (i) finetuning (training) is performed on 60% of TruthfulQA, and (ii) evaluation (test) is performed exclusively on the held-out 40%. All stitched models (both ESTA and KLAS) are finetuned only the 60% portion and are never exposed to the remaining 40% used for evaluation. Thus, all stitched LLMs are evaluated strictly on unseen TruthfulQA questions. We also note that the Llama 3.2 instruction-tuned IB and 3B anchor models Grattafiori et al. (2024) are pretrained on a large mixture web-scale English corpora, code, and cestruction-follovin datasets, and do not include TruthfulQA in their pretraining mix.
ESTA applics fixed stitching points (e.g., using the first layers of the IB model and the remaining layers from the 3B model), whereas KLAS selects stitch locations using KL divergence computed over linear probe outputs to select stitch points that have similar next-token prediction distributions. Tab. oshows that across comparable stitched model sizes, KLAS consistently yields higher ROUGE-1 and ROUGE-2 scores than ESTA, demonstrating better stitch configuration selection and hence, an improved Pareto front of stitched models.
I USE OF LLMS IN OUR SUBMISSION
Yes, used Large Language Models (LLMs) aid in polishing the writing. Specifically, LLMs were used for grammar checking, improving readability, and refining phrasing. All technical content, experiments, and analyses were conducted and written by the authors.
23


Reviews:
[{"review_id": "PclOY4kk3B", "reviewer": "Reviewer_B36v", "summary": "Stitching pretrained models offers a cost-effective way to explore accuracy-efficiency tradeoffs. Current stitching methods rely on heuristics, leading to sub-optimal results. This paper improves this by using KL-divergence to automate stitch selection and generalize across architectures. KLAS outperforms baselines, achieving higher accuracy at the same computational cost or reducing FLOPs while maintaining accuracy.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Technical contribution: This paper addresses a fundamental challenge in current model stitching methods: these methods rely on heuristic-based stitch selection, fix anchors and blocks, and thus yield suboptimal accuracy-efficiency tradeoffs. Specifically, it proposes a coarse-grained anchor selection strategy that leverages the KL divergence of the last block for anchor identification, and employs block-level KL divergence for fine-grained selection. The metric for candidate set selection is well-justified, as it considers both sampling coverage during fine-tuning and the quality of anchor filtering. \n\n* Experiments effectively demonstrate the advantages of the proposed KL-divergence-based anchor stitching approach.", "weaknesses": "* Why MSE, CE, CKA, DM can have significantly lower percentage of stitch configurations than KL divergence? Mathematically, it’s not that straightforward.\n\n* The method’s applicability to large language models and multimodal LLMs has yet to be explored—extending it to these models would further enhance the paper’s impact.\n\n* What performance can KLAS achieve on dense prediction tasks, including object detection, semantic segmentation, and depth estimation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iEOFwDCvhn", "reviewer": "Reviewer_XiFF", "summary": "The paper proposes stitching blocks of neural network by computing similarity between blocks by taking the KL-Divergence of linear probe probability distributions. The paper suggests that \"far-stitching\" where models with significant difference in complexity or performance may also be worthy of stitching. The paper proposes to automatically identify block pairs to stich opposed to SN-Net where they use defined constraints to choose the pairs. The paper compares the proposed KLAS stitching framework with SN-Net on DeiT and Swin model architectures trained on CIFAR-100 and imagenet-1k datasets. The results show marginal gain in performance compared to SN-Net.", "questions": "1.  For the stitch finetuning step is the number of training steps per stitched model constant or does it vary depending on when the stitch converges?\n2. I did not get if the the finetuned stitched model is same for both SN-NET and KLAS if the stitch pairs overlap.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 1, "presentation": 2, "contribution": 3, "strengths": "1. The paper explores functional similarity along with representation similarity as a metric while choosing to stitch networks which is logical. If two networks produce similar same functional output, then the models are better suited for selecting stitching pairs, when considered with representation similarity.\n2. The paper also shows that far stitching where two networks may not have similar performance still can be stitched together is an important contribution.\n3. The experimentation on DeiT and Swin family of models to compare with SN-Net shows similar performance but seemingly less finetuning cost.", "weaknesses": "1. In line 232 the paper states \"As a representational similarity metric, KL divergence captures distributional differences between intermediate activations, indicating whether two blocks generate patterns that\ncan be mapped via lightweight transformations.\" which is not supported in my opinion as KL divergence compares the probabilistic score distribution not the representation itself. (Think different features can be used to reach same conclusion with similar confidence).\n\n2. There is only marginal improvement when comparing stitched models of similar flops compared to SN-Net, while individual anchor-block stitches give an edge to KLAS, collectively the gain is minimal.\n\n3. The lack of representational similarity is a concern (as mentioned in point 1 KL divergence doesnt directly compare representations), if there is difference in representational similarity there would be more finetuning cost.\n \n4. The details on stitched model finetuning could be added for better assessment. This is important to assess the improvement in AUC.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UDB6IaFEk2", "reviewer": "Reviewer_cTKy", "summary": "This paper proposes the use of inserted linear probes and KL divergence between them to find layers compatible for stitching between different networks.", "questions": "In what scenarios is stitching preferrable to a model cascade? Why?\n\nDo linear probes work for regression problems? For very large classifiers (think LLMs)? (the discussion mentions this very shallowly, there is no need to show improvement both for prefill and decode, just pick the easier case and show some improvement)", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- clear language and well written\n\n- the idea of using linear probes + KL divergence for stitching is sound and well supported by experiments", "weaknesses": "- The declared goal of the KLAS approach is to use models from a pre-trained zoo to construct models that provide new accuracy / cost trade-offs. This is presented as an alternative to NAS (a somewhat unfair comparison as NAS enriches to base-model pool). However, the standard method of addressing this problem are model cascades / committees (e.g. \"WISDOM OF COMMITTEES: AN OVERLOOKED APPROACH TO FASTER AND MORE ACCURATE MODELS\" or \"Efficient Inference With Model Cascades\"). Model cascades are not discussed at all and not compared against in this paper. A comparison against a strong cascade baseline is essential to be able to claim utility of the proposed method for the suggested purpose. In addition to the direct comparison to a model cascade, please also explain whether there are situations where stitching is preferrable over cascading for structural reasons (maybe average case complexity vs. worst case complexity?). A 1) worst-case and 2) average-case acc/cost tradeoff curve comparing stitching and cascading would be meaningful.\n\n- \"KL divergence uniquely satisfies the dual objectives\" this claim is not substantiated. Probably the authors mean ~ 'uniquely among the few measures we consider here'; if yes, please rephrase, if not, please provide a proof that the KL divergence is indeed unique in this regard.\n\n- it is unclear whether the method is specific to image classifiers or works in other domains as well. The impact of the paper would be much broader, if there was some indication that it works for regression tasks (--> how to substitute the KL div??) and for non-vision classification (e.g. also token prediction in LLMs). Currently the method is of limited interest as it only applies to a niche task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2Qf5jCv6vF", "reviewer": "Reviewer_8KnS", "summary": "This paper proposes KL-divergence–based Anchor Stitching (KLAS), a framework for selecting where and how to “stitch” pretrained networks so as to interpolate accuracy–efficiency trade-offs more effectively than heuristic stitching (e.g., SN-Net). KLAS (i) chooses anchor pairs by the last-block KL divergence of their predictive distributions and (ii) ranks block pairs with a stitch score \\\\(\\\\Gamma(i,j)\\\\) that combines cross-anchor activation distance and the target block’s “capacity” (consecutive-block KL). To obtain the distributions, the authors train lightweight linear probes (ProbeNet) on each block; probes converge in ≈4 epochs and add a small one-time cost. Across DeiT, Swin, LeViT, and ResNet (ImageNet-1K, CIFAR-100/10), KLAS improves the AUC of the accuracy–FLOPs trade-off curves over SN-Net; gains include up to +1.21% top-1 accuracy at equal FLOPs or 1.33× FLOPs reduction at equal accuracy.", "questions": "1. **Un/SSL compatibility.** Can KLAS be made label-free, e.g., by using self-supervised probes (DINO-style heads) or pseudo-labels? How does anchor/block selection change when anchors are self-supervised?\n2. **Calibration sensitivity.** How sensitive are \\\\(\\\\Theta\\\\) and \\\\(\\\\Gamma\\\\) to softmax temperature, label smoothing, and class imbalance? Could you report ablations and perhaps use temperature-scaled KL?\n3. **Search cost accounting.** What is the end-to-end overhead (wall-clock/GPU-days) of ProbeNet + KLAS vs. SN-Net’s heuristics at equal stitch fine-tuning budgets? Please include variability across families.\n4. **Ranking validity.** Beyond the Min-KL vs. \\(\\Gamma\\) comparison, can you provide rank correlation between KLAS scores and post-fine-tuning accuracy across candidates (per bucket and overall)?\n5. **Cross-family stitches.** For cases with very small ΔAUC, what failure modes arise (capacity mismatch, optimization instability, probe mis-ranking)? Any diagnostics (e.g., KL heatmaps) that predict such cases?\n6. **Dense tasks.** Could you include full experiments on detection/segmentation (COCO/ADE20K) where stitching locations affect multi-scale features, and compare to dynamic routing/pruning baselines?\n\n**If the author can address my questions, I am willing to improve my rating.**", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Principled selection vs. heuristics.** The paper clearly articulates why nearest/paired stitching can be suboptimal and replaces it with a similarity-driven criterion grounded in KL divergence, with explicit formulas for \\\\(\\\\Theta\\\\) and \\\\(\\\\Gamma\\\\).\n2. **Dual notion of similarity.** The KL criterion is argued (and operationalized) to reflect both epresentational alignment and functional compatibility, addressing shortcomings of CKA/MSE/CE/DM for choosing stitch points.\n3. **Efficient probing.** ProbeNet trains one set of blockwise probes efficiently (≈0.25 GPU-days for Swin-B) and shows fast convergence, keeping selection overhead modest.\n4. **Broad empirical coverage.** Results span multiple families (ViTs, CNNs) and even cross-architecture stitches (e.g., ResNet↔Swin), consistently improving trade-off AUC vs. SN-Net.\n5. **Anchor selection that adapts by family.** Last-block KL correctly prefers far stitching (Ti→B) for Swin but nearest (Ti↔S, S↔B) for DeiT, demonstrating generality beyond a fixed heuristic.", "weaknesses": "1. **Supervision dependence for “similarity.”** KL is computed on softmax outputs of supervised probes; thus, selection intrinsically depends on labels and probe training. This limits claims for unsupervised/self-supervised representation learning and may bias choices toward the probe’s training distribution.\n2. **Asymmetry & calibration sensitivity.** KL’s asymmetry and dependence on calibration/temperature may distort distances across blocks/models; the method normalizes by an intra-anchor term, but robustness to temperature, class imbalance, or label smoothing is underexplored.\n3. **Metric/selection design choices.** The final selection uses bucketized FLOPs and a threshold \\\\(\\\\tau\\\\) (5% of the minimum in each bucket). The AUC metric and bucket granularity can influence conclusions; more sensitivity studies would strengthen the case.\n4. **Scale of gains in some regimes is small.** Several settings show marginal improvements (e.g., cross-architecture ΔAUC≈+0.002), raising questions about practical significance across all families.\n5. **Representation-learning scope.** Experiments are largely supervised classification; there is no evaluation with self-supervised anchors (e.g., DINO/MAE) nor transfer via frozen-backbone linear probes on diverse tasks.\n6. **Limited evidence for dense prediction/generalization.** Dense-task adaptation is left as future work with only preliminary results; rigorous detection/segmentation studies are missing.\n7. **Compute accounting.** While probe training is “negligible,” it is still an added search cost versus purely heuristic SN-Net; a wall-clock comparison for the full pipeline (probes + stitch fine-tuning) would help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)