可微分高效算子搜索

# 可微高效算子搜索
来源：https://arxiv.org/html/2606.05232
裴晓欢¹　张纪元²　郭远帆²　冯卫国²　黄涛³　Cho-Jui Hsieh⁴　徐畅¹
¹悉尼大学　²字节跳动　³上海交通大学　⁴加州大学洛杉矶分校
[email protected]

###### 摘要

高效模型在很大程度上依赖于人工设计的缩减算子，例如剪枝、合并、池化和自适应重加权，这些操作可在任意时刻和任意位置进行。我们表明，这些看似不同的方法可以统一为单一共享算子空间中的不同运行模式。基于这一观察，我们引入了**高效算子搜索**，这是一个统一框架，其中连续参数控制令牌信息是被移除、锐利合并、均匀池化还是软重分配。我们无需手工设计算子组合，而是定义了一个*高效搜索空间*：对层激活、保留预算和算子模式进行可微参数化，以及一个*高效搜索策略*：在单边预算和成本约束下最小化任务损失的期望。该流程导致从**高效设计问题**到**算子搜索问题**的更广泛转变，为未来高效建模提出了新范式。有趣的是，我们发现大多数先前主流基线可以视为这个共享算子空间的特殊情况，并进一步揭示了跨不同基准的一致算子模式。即使作为这一新范式的早期探索，我们提出的方法也在各种基准上取得了有竞争力的结果。我们相信，它为先前方法提供了统一视角，为理解当前差异提供了原则性透镜，并为未来高效建模提供了通用基础。网页：EOS (https://www.terrypei.com/eos)。

## 1 引言

扩展多模态基础模型在视觉推理、多模态理解、指令遵循和长上下文多模态推理方面带来了显著的性能提升[12 (https://arxiv.org/html/2606.05232#bib.bib12),10 (https://arxiv.org/html/2606.05232#bib.bib10),26 (https://arxiv.org/html/2606.05232#bib.bib26),27 (https://arxiv.org/html/2606.05232#bib.bib27),6 (https://arxiv.org/html/2606.05232#bib.bib6),28 (https://arxiv.org/html/2606.05232#bib.bib28),9 (https://arxiv.org/html/2606.05232#bib.bib9)]。然而，这些增益通常伴随着高昂的推理成本，因为密集的视觉令牌会被大语言骨干网络反复处理，尤其是在高分辨率、多图像和长上下文场景中[31 (https://arxiv.org/html/2606.05232#bib.bib31),13 (https://arxiv.org/html/2606.05232#bib.bib13),23 (https://arxiv.org/html/2606.05232#bib.bib23),24 (https://arxiv.org/html/2606.05232#bib.bib24),30 (https://arxiv.org/html/2606.05232#bib.bib30),22 (https://arxiv.org/html/2606.05232#bib.bib22),20 (https://arxiv.org/html/2606.05232#bib.bib20)]。现有工作通过人工设计的令牌缩减算子（包括剪枝、合并、池化和自适应重加权）来提高基础模型的效率。例如，SparseVLM[32 (https://arxiv.org/html/2606.05232#bib.bib32)] 根据注意力分数移除视觉上不重要的令牌，而 ToMe[1 (https://arxiv.org/html/2606.05232#bib.bib1)] 合并冗余令牌以缩短输入序列同时保留信息。尽管这些方法展示了有前景的精度-效率权衡，但它们通常作为独立的压缩方案开发，并采用手动选择的层、预算和算子形式，使得这些角落算子背后的共享结构很大程度上未被探索。

<figure>
<p>参见说明 (a) 手动设计 vs. 自动搜索。</p>
<p>参见说明 (b) 统一算子空间与性能。</p>
</figure>

图1：高效算子搜索概览。(a) EOS 将手动设计的缩减方案替换为自动算子搜索。(b) 搜索得到的混合算子位于统一算子空间内，并在相同令牌预算下提升了性能。尽管实现方式不同，现有的缩减行为可以解释为丢弃、转移或重新分配令牌信息的不同方式。这种解释启发了一种共享的公式，其中代表性高效算子由连续参数控制。具体来说，信息传递门 \(\gamma\) 决定丢弃的令牌信息是移除还是保留，而分配温度 \(\tau\) 控制保留的信息是锐利合并、均匀池化还是软重分配。在这种视角下，纯剪枝、硬合并、平均池化和自适应重加权自然成为同一算子空间的不同运行模式（图1 (https://arxiv.org/html/2606.05232#S1.F1)）。基于这一观察，我们提出 **高效算子搜索**，这是一个将高效模型从**算子设计问题**重新定义为**算子搜索问题**的统一框架。我们的方法不是手动构造另一个固定的压缩规则，而是搜索在哪里压缩、保留多少令牌以及应用哪种缩减行为。这实现了在精度-效率约束下灵活组合多种压缩基元。我们在代表性多模态基准上进行了大量实验。结果表明，我们的框架能够恢复强人工设计基线，并进一步发现更好的算子配置，实现改进的性能-效率权衡。

我们的贡献可从三个方面总结：(1) 我们引入了一个统一算子空间，将现有令牌缩减方法重新解释为单一共享公式的不同模式。(2) 我们提出了高效算子搜索，它联合搜索层选择、令牌预算和连续算子参数，而不是依赖手动设计的方案。(3) 我们展示了所提框架能够恢复现有基线，并为高效多模态基础模型发现更强的混合配置。

## 2 方法

<figure>
<p>参见说明</p>
</figure>

图2：高效算子搜索概览。给定一个冻结的多模态基础模型，EOS 在每个解码器层通过三个耦合组件参数化令牌缩减：层激活 \(g_l\)、保留预算 \(c_l\) 和算子模式 \(\Omega_l = (\gamma_l, \tau_l, \theta_l, \rho_l, \nu_l)\)。在每个活跃层，重要的视觉令牌被保留为锚点，而剩余候选者由一个统一缩减算子处理。根据学习到的算子参数，丢弃的信息可以被移除、锐利合并、均匀池化或软重分配给锚点。我们的目标是将高效基础模型从手动设计的压缩方案转变为可搜索的算子空间。给定一个具有 \(L\) 个解码器层的冻结多模态基础模型，现有的令牌缩减方法通常手动指定*在哪里*减少视觉令牌、*保留多少*令牌，以及*应用哪个*算子。相比之下，我们的**高效算子搜索** (EOS) 将这些选择参数化为可学习的搜索变量，并在任务和效率约束下通过可微松弛联合优化它们。如图2 (https://arxiv.org/html/2606.05232#S2.F2) 所示，EOS 在每个活跃缩减层选择重要的视觉令牌锚点，然后应用统一算子来决定剩余令牌信息是被丢弃、转移、池化还是重加权。

### 2.1 预备知识

令 \(\mathbf{X}^{(l)} \in \mathbb{R}^{N_l \times d}\) 表示解码器层 \(l\) 处的视觉隐藏状态，其中 \(N_l\) 是视觉令牌数量，\(d\) 是隐藏维度。令牌缩减操作将 \(\mathbf{X}^{(l)}\) 划分为保留的锚点集和缩减候选集：

\[\mathbf{X}^{(l)} = \mathbf{A}^{(l)} \cup \mathbf{D}^{(l)}, \quad \mathbf{A}^{(l)} \in \mathbb{R}^{K_l \times d}, \quad \mathbf{D}^{(l)} \in \mathbb{R}^{M_l \times d}, \tag{1}\]

其中 \(K_l\) 是保留的锚点数量，\(M_l = N_l - K_l\) 是缩减候选数量，且 \(\mathbf{A}^{(l)} \cap \mathbf{D}^{(l)} = \emptyset\)。锚点根据重要性分数 \(\mathbf{q}^{(l)} \in \mathbb{R}^{N_l}\) 选择，该分数可以根据先前高效推理方法，基于注意力或相似性的令牌重要性来实例化。剩下的问题是如何在锚点选择后处理 \(\mathbf{D}^{(l)}\)。纯剪枝直接移除 \(\mathbf{D}^{(l)}\)，令牌合并将每个丢弃的令牌转移到相似的锚点，池化均匀分配丢弃的信息，而自适应重加权根据丢弃令牌的重要性重新缩放锚点。EOS 将这些行为视为单一共享算子族的不同运行模式。

### 2.2 统一缩减算子

为简单起见，当上下文清晰时我们省略层索引 \(l\)。给定锚点 \(\mathbf{A} \in \mathbb{R}^{K \times d}\) 和缩减候选 \(\mathbf{D} \in \mathbb{R}^{M \times d}\)，我们首先计算归一化相似度矩阵：

\[\mathbf{S} = \bar{\mathbf{D}} \bar{\mathbf{A}}^\top \in \mathbb{R}^{M \times K}, \qquad \bar{\mathbf{x}}_i = \frac{\mathbf{x}_i}{\|\mathbf{x}_i\|_2 + \epsilon}, \tag{2}\]

其中 \(\mathbf{x}_i\) 表示来自 \(\mathbf{A}\) 或 \(\mathbf{D}\) 的行令牌特征，归一化按令牌进行，\(\epsilon > 0\) 是为数值稳定性使用的小常数。每个条目 \(S_{ij}\) 衡量缩减候选 \(\mathbf{d}_i\) 和锚点令牌 \(\mathbf{a}_j\) 之间的余弦相似度。从缩减候选到锚点的分配由下式给出：

\[\mathbf{W} = \mathrm{softmax}_{\mathrm{anchor}}\left(\frac{\mathbf{S}}{{\color[rgb]{0,0.390625,0.78515625}\tau_l}}\right) \in \mathbb{R}^{M \times K}, \tag{3}\]

其中 \(\mathrm{softmax}_{\mathrm{anchor}}(\cdot)\) 对每行在锚点维度上进行归一化，\(\tau_l > 0\) 是分配温度。较小的 \(\tau_l\) 导致更锐利的最近锚点分配，而较大的 \(\tau_l\) 产生跨锚点的更平滑分配。

我们进一步定义每个令牌的传递门：

\[m_i = \sigma\left(\beta(\max_j S_{ij} - {\color[rgb]{0,0.390625,0.78515625}\theta_l})\right), \qquad \mathbf{m} = [m_1, \ldots, m_M]^\top, \tag{4}\]

其中 \(\theta_l\) 是相似度阈值，\(\beta\) 控制门的锐利程度。统一的传递操作然后由下式给出：

\[\tilde{\mathbf{A}} = \mathbf{A} + {\color[rgb]{0,0.390625,0.78515625}\gamma_l}\,\mathbf{W}^\top(\mathbf{D} \odot \mathbf{m}), \tag{5}\]

其中 \(\mathbf{m}\) 沿特征维度广播，\(\gamma_l \in [0,1]\) 是信息传递门。当 \(\gamma_l = 0\) 时，缩减候选没有显式特征传递，实际上被剪枝。当 \(\gamma_l > 0\) 时，它们的信息根据分配矩阵 \(\mathbf{W}\) 转移到保留的锚点。

为了包含自适应重加权并稳定输出尺度，我们进一步应用：

\[\hat{\mathbf{A}} = \tilde{\mathbf{A}} \odot \left(1 + {\color[rgb]{0,0.390625,0.78515625}\rho_l}\,\mathbf{s}_{\mathrm{imp}}\right), \qquad \hat{\mathbf{A}} \leftarrow (1 - {\color[rgb]{0,0.390625,0.78515625}\nu_l})\hat{\mathbf{A}} + {\color[rgb]{0,0.390625,0.78515625}\nu_l}\frac{\|\mathbf{A}\|_F}{\|\hat{\mathbf{A}}\|_F + \epsilon}\hat{\mathbf{A}}, \tag{6}\]

其中 \(\rho_l\) 表示锚点重加权强度，\(\nu_l\) 表示范数保持系数，\(\mathbf{s}_{\mathrm{imp}} \in \mathbb{R}^K\) 表示由缩减候选导致的锚点级重要性。这里，\(\gamma_l\) 控制显式特征传递，而 \(\rho_l\) 控制间接锚点重新缩放。在实践中，\(\mathbf{s}_{\mathrm{imp}}\) 由下式给出：

\[\mathbf{s}_{\mathrm{imp}} = \mathrm{softmax}_{\mathrm{anchor}}(\mathbf{S})^\top \mathrm{softmax}_{\mathrm{candidate}}(\mathbf{q}_\mathcal{D}), \tag{7}\]

其中 \(\mathbf{q}_\mathcal{D}\) 表示缩减候选的重要性分数，\(\mathrm{softmax}_{\mathrm{candidate}}(\cdot)\) 在缩减候选维度上进行归一化。

###### 定义 2.1 (统一高效算子)。

在层 \(l\) 处，高效缩减算子由参数元组定义：

\[\Omega_l = \left({\color[rgb]{0,0.390625,0.78515625}\gamma_l}, {\color[rgb]{0,0.390625,0.78515625}\tau_l}, {\color[rgb]{0,0.390625,0.78515625}\theta_l}, {\color[rgb]{0,0.390625,0.78515625}\rho_l}, {\color[rgb]{0,0.390625,0.78515625}\nu_l}\right), \tag{8}\]

并通过公式 (2)–(6) 将视觉隐藏状态从 \(\mathbf{X}^{(l)} \in \mathbb{R}^{N_l \times d}\) 映射到缩减表示 \(\hat{\mathbf{A}}^{(l)} \in \mathbb{R}^{K_l \times d}\)。参数 \(\gamma_l\) 和 \(\tau_l\) 决定主导缩减模式，而 \(\theta_l\)、\(\rho_l\) 和 \(\nu_l\) 提供令牌级门控、锚点重加权和尺度保持。

参数元组 \(\Omega_l\) 定义了单层的缩减行为。EOS 进一步将此算子模式与层激活和令牌预算耦合，形成一个可搜索的空间，覆盖在哪里压缩、压缩多少以及如何处理缩减的令牌。

### 2.3 高效搜索空间

EOS 搜索三个耦合维度：**层激活**、**令牌预算**和**算子模式**。对于所有解码器层，完整的搜索空间由下式给出：

\[\boxed{\bm{\Theta} = \Bigl\{\underbrace{{\color[rgb]{0.78515625,0.1953125,0}g_l}}_{\text{哪些层}},\;\underbrace{{\color[rgb]{0,0.58984375,0.1953125}c_l}}_{\text{多少}},\;\underbrace{{\color[rgb]{0,0.390625,0.78515625}\gamma_l, \tau_l, \theta_l, \rho_l, \nu_l}}_{\text{哪个算子}}\Bigr\}_{l=0}^{L-1}}\]