🤖 machine learning

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

本文介绍了 KLAS，这是一个新颖的框架，它利用 KL 散度识别最相似的中间表示，从而自动化并优化预训练神经网络的拼接，与现有的基于启发式的方法相比，显著改善了精度与效率之间的权衡。

原作者： Debopam Sanyal, Anantharaman Iyer, Alind Khare, Trisha Jain, Akshay Jajoo, Myungjin Lee, Clayton Kerce, Alexey Tumanov

发布于 2026-05-29

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Debopam Sanyal, Anantharaman Iyer, Alind Khare, Trisha Jain, Akshay Jajoo, Myungjin Lee, Clayton Kerce, Alexey Tumanov

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个由预构建机器人组成的庞大图书馆。其中一些机器人小巧快速但不够智能（如同玩具机器人）；另一些则庞大、缓慢却极其聪明（如同超级计算机）。通常，如果你需要一个“恰到好处”的机器人——既足够智能以胜任工作，又足够快速以便在电池上运行——你就必须从零开始打造一台全新的机器人。这需要耗费大量的时间、金钱和精力。

几年前，科学家们发现了一条捷径：拼接（Stitching）。与其从头构建新机器人，不如将小机器人的“眼睛”与大机器人的“大脑”缝合在一起。这样，你就能得到一台介于两者之间的混合机器人。

然而，这种拼接方法存在一个重大问题。负责缝合的人们使用的是糟糕的经验法则。他们只会说：“把小机器人的第二只眼睛缝到大机器人的第二个大脑上”，或者“只连接大小相似的机器人”。他们并没有真正检查这些部件是否真的契合。有时，部件之间会发生冲突，导致新机器人无法良好运行。

KLAS 登场了。

这篇论文的作者 KLAS 就像一位不靠猜测的大师裁缝。相反，他们使用一种名为**KL 散度（KL Divergence）**的特殊工具（将其想象为“兼容性扫描仪”），来精确衡量一个机器人部件的“思维”与另一个部件的相似程度。

以下是 KLAS 的工作原理，分解为简单步骤：

1. “最近邻”拼接的问题

想象你正在尝试配对袜子。旧的方法（称为 SN-Net）会说：“只需匹配货架上彼此最靠近的袜子。”

缺陷： 有时，一只鲜红色的袜子（一个小模型）实际上与一只深蓝色的袜子（一个大模型）完美匹配，因为它们由相同的布料制成，尽管它们在货架上相距甚远。旧方法错过了这些完美匹配，因为它只关注距离，而非质量。

2. KLAS 的“兼容性扫描仪”

KLAS 使用一个名为KL 散度的数学概念。想象有两个人在交谈。

旧方法： 检查他们是否站在一起。
KLAS： 倾听他们在说什么。它会问：“如果 A 说话，B 是否能完美理解其语气和含义？”
如果答案是“是”（低 KL 散度），KLAS 就知道这两个部件配合良好。如果答案是“否”（高 KL 散度），KLAS 就知道它们会发生冲突。

3. "ProbeNet"的魔力

为了在不实际先缝合机器人（这会既慢又贵）的情况下检查这种兼容性，KLAS 使用了一个名为ProbeNet的巧妙技巧。

想象你想知道汽车引擎是否适配新的底盘。与其将它们螺栓固定在一起并尝试驾驶，不如将一个小巧的临时传感器连接到引擎上。该传感器会快速检查引擎的输出是否符合底盘的预期。
KLAS 通过“线性探针”（linear probes）来做这件事。它在预训练模型上放置一个微小的临时测试层，以观察它们如何“思考”。这几乎不花时间，也几乎不花钱。

4. 结果：完美契合

通过使用这种扫描仪，KLAS 能够找到最佳的组合，而不仅仅是“最接近”的组合。

主张： 在他们的测试中，KLAS 发现的组合在准确性（智能程度）上优于旧方法，同时使用了相同数量的计算能力。
或者： 它发现的组合智能程度相当，但使用了更少的算力（节省电池和资金）。

论文中的现实世界示例

作者在以下领域测试了该方法：

图像识别： 让计算机更好地识别猫、狗和汽车。他们拼接了不同版本的"Swin"和"DeiT"模型。KLAS 发现，有时将最小模型直接连接到最大模型（跳过中间模型）比旧的“逐步”方法效果更好。
医疗/地图任务： 他们甚至将其应用于“语义分割”（Semantic Segmentation，即像地图一样为照片的每个像素着色以显示其内容）。KLAS 在相同的计算预算下生成了更优质的地图。
语言模型： 他们在大型语言模型（如撰写文本的模型）上进行了测试。KLAS 发现了混合小语言模型与大语言模型的更好方法，使得结果比旧的“层数计数”方法更智能、更真实。

核心结论

该论文认为，我们无需从头构建新的 AI 模型即可获得速度与智能的完美平衡。我们只需停止猜测哪些部件应该缝合在一起。通过使用“兼容性扫描仪”（KL 散度）来倾听部件如何思考，我们可以自动拼接出适合任何任务的完美混合 AI，从而节省时间、金钱和精力。

简而言之： KLAS 是一位在裁剪前先测量布料的大师裁缝，确保最终制成的西装完美合身；而旧方法则只是根据布料挂在衣柜里的位置进行猜测。

技术摘要：KLAS——利用相似性拼接神经网络以优化精度与效率的权衡

1. 问题陈述

部署大型预训练视觉 Transformer（ViTs）和卷积神经网络（CNNs）面临一个根本性挑战：多样化的环境从强大的云服务器到资源受限的边缘设备不等，这要求在特定的计算预算内进行灵活的模型选择。虽然神经架构搜索（NAS）试图平衡精度与效率，但传统的一对一 NAS 计算成本过高，而一对多或零样本 NAS 方法通常局限于单一模型设计空间，未能利用丰富的现有预训练模型库。

模型拼接（Model Stitching）提供了一种有前景的替代方案，通过组合不同预训练“锚点”模型的块来创建插值网络。然而，现有的多对多拼接方法（如 SN-Net）依赖简单的启发式规则（例如相邻锚点的“最近邻拼接”或固定的块配对）来选择拼接配置。这些启发式方法往往产生次优的精度 - 效率权衡，且缺乏跨不同模型家族的泛化能力。此外，拼接配置的搜索空间巨大（对于深度为 $n$ 的 $k$ 个模型，复杂度为 $O(k^2n^2)$ ），使得穷举搜索不可行。核心问题在于缺乏一种原则性的、自动化的方法，以在不产生高昂训练成本的情况下识别出能最大化性能的兼容拼接配置。

2. 方法论：KLAS 框架

作者提出了KLAS（基于 KL 散度的锚点拼接），这是一个通过利用预训练模型中间表示的相似性来自动化拼接选择的框架。

核心概念：KL 散度作为相似性度量

KLAS 认为，成功的拼接取决于表示相似性（激活的几何对齐）和功能相似性（任务相关信息和决策边界的保留）。作者指出，Kullback-Leibler（KL）散度独特地适合捕捉这两者：

表示对齐：它量化了激活分布之间的统计距离，表明轻量级变换是否可以将源激活映射到目标激活。
功能对齐：它评估源块的输出在多大程度上保留了目标模型的决策边界，确保拼接后的网络在微调后能维持下游任务性能。

与中心核对齐（CKA）、均方误差（MSE）或交叉熵（CE）等往往无法预测拼接成功率的指标不同，KL 散度直接测量由线性探针生成的概率分布的兼容性。

KLAS 算法

该框架主要在两个阶段运行：

锚点选择：
- KLAS 计算候选预训练模型最后块的 softmax 概率分布之间的 KL 散度。
- 最后块 KL 散度较低的模型被视为兼容的锚点，因为它们共享相似的决策边界和类别置信度分布。这取代了先前工作中使用的启发式“最近邻拼接”。
块选择：
- 一旦选定锚点，KLAS 使用定义的拼接分数 $\Gamma(i, j)$ 来评估块对：
  $\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$
  其中：
  - $\Theta(P^f_i, P^g_j)$ 是跨锚点激活距离（源块 $i$ 与目标块 $j$ 之间的 KL 散度）。
  - $\Theta(P^g_j, P^g_{j+1})$ 是锚点内块容量（目标锚点中连续块之间的 KL 散度）。
- 该分数平衡了假设拼接层上的变换负担与目标块吸收分布差异的能力。
- ProbeNet：为了在不训练完整拼接层的情况下高效计算这些分数，KLAS 采用了一种名为 ProbeNet 的统一架构。它将线性分类器探针连接到锚点模型的每个块，并在单次前向 - 后向传递中联合训练它们。与训练单独探针相比，这产生的额外一次性成本微乎其微（例如，Swin-B 仅需 0.25 GPU 天）。
配置选择：
- 基于 FLOPs 桶内最小分数的可调节阈值 $\tau$ 过滤候选拼接配置。这确保了在精度 - 效率谱系上生成密集的部署模型帕累托前沿。

3. 主要贡献

启发式局限性的识别：论文证明，简单的启发式方法（如 SN-Net 中的方法）并非普遍最优，往往会错过更优越的“远距离拼接”配置（例如，直接将 Swin-Ti 连接到 Swin-B），而这些配置常被启发式方法忽略。
原则性的相似性度量：引入 KL 散度作为拼接选择指标，无需实例化或训练拼接层即可捕捉表示和功能对齐。
KLAS 框架：一个可泛化、低成本的框架，可自动化跨不同模型家族（ViTs、CNNs 和 LLMs）的锚点和块选择。
实证验证：全面的实验表明，在多个数据集（ImageNet-1K、CIFAR-100、ADE20K、TruthfulQA）和架构上，KLAS 相比基线方法具有一致的改进。

4. 实验结果

作者在相同的微调条件下（视觉模型 50 个 epoch，LLMs 10 个 epoch）评估了 KLAS 与 SN-Net 及其他基线方法。

精度 - 效率权衡：
- ImageNet-1K：与 SN-Net 相比，KLAS 在相同计算成本下实现了高达1.21% 的 Top-1 准确率提升，或在同等准确率下实现了1.33 倍的 FLOPs 减少。
- AUC 分数：在 Swin 家族上，KLAS 实现了0.8950的曲线下面积（AUC），显著优于 SN-Net（0.8345）和其他相似性指标如 CKA（0.8124）。
- 跨架构：KLAS 成功拼接了 ResNet-18 与 Swin-B，在相似 FLOPs 下，其 Top-1 准确率比级联推理基线高出 2-3%。
密集预测（语义分割）：
- 在 ADE20K 上使用 Mask2Former，KLAS 选择的拼接在可比 FLOPs 预算下，相比 SN-Net 基线实现了高达**+0.9% 的 mIoU**提升。
大型语言模型（LLMs）：
- 在 TruthfulQA 基准测试上使用 Llama 1B 和 3B 锚点，尽管某些配置的参数量较少，KLAS 拼接的模型在 ROUGE-1 和 ROUGE-2 指标上分别比基于启发式的 LLM 拼接方法 ESTA 高出0.017和0.033。
指标有效性：
- KL 散度成功恢复了**88.9%**的可行 Swin Ti-B 拼接配置，而 CKA 仅恢复了 5.5%，SN-Net 启发式方法恢复了 61.1%。

5. 意义与主张

论文主张，KLAS 通过用原则性的、基于相似性的方法取代临时启发式规则，解决了模型拼接中的关键差距。其意义在于：

泛化性：它适用于多样化的架构家族（Transformer、CNN、混合架构）和模态（视觉、语言），克服了局限于单一家族的方法的局限性。
成本效益：它以可忽略的额外计算开销（仅探针训练和拼接层微调）生成连续的部署模型帕累托前沿，避免了传统 NAS 的高昂成本。
系统性优化：通过 KL 散度显式捕捉分布差异，KLAS 能够发现启发式方法所遗漏的非直观、高性能的拼接配置（例如“远距离拼接”）。

作者总结道，KLAS 为部署场景中的精度 - 效率权衡优化提供了一个鲁棒、自动化的解决方案，有效利用了现有的预训练模型生态系统，而无需从头开始训练。