GRASP: Guided Residual Adapters with Sample-wise Partitioning¶

会议: NeurIPS 2026
arXiv: 2512.01675
领域: 医学图像
关键词: 流匹配, 长尾分布, 残差适配器, 医学影像生成, 梯度一致性

一句话总结¶

GRASP 通过基于条件空间（频次分布与健康先验或文本语义聚类）的静态确定性划分，将样本解耦路由至扩散 Transformer 前馈层中专有的残差适配器，从优化底层消除长尾分布下头尾类别的破坏性梯度干涉，在保持流匹配目标与采样器零侵入的前提下大幅提升罕见病症生成的保真度与多样性。

研究背景与动机¶

文本到图像流匹配模型（Text-to-Image Flow Matching Transformers，如 FLUX.1）极大地推进了高质量视觉生成的前沿，但在真实世界的严重长尾分布下却暴露出严重的退化问题。在医学影像（如胸部 X 光 MIMIC-CXR-LT）等核心场景中，常见病症和无异常样本占据绝大多数，而具有高度临床诊断价值的罕见病症样本极度匮乏。现有的生成模型在微调时往往被高频头部类别主导，导致尾部罕见类别在保真度与多样性上发生模式崩溃（Mode Collapse），生成的罕见病症图像结构单一甚至丢失关键病灶特征，严重制约了合成数据在下游少样本病种增强中的临床应用潜力。

现有缓解长尾生成崩溃的方法主要分为两派，但均未能触及问题的核心症结。后验采样方法（如各类无分类器引导 CFG 变体或少数类加权采样）主要在推理阶段驱使采样轨迹偏向低密度区域，这虽然提升了搜索范围，却无法改变已训练模型的真实条件分布，因而无法挽回训练期即已发生的条件模式崩溃；而修改目标函数或条件路径的训练期方法（如条件-无条件对比对齐、病灶掩码可控合成或因果解耦）改变了模型学习的具体内容，却未能解决优化底层的病理机制。梯度动力学研究表明，长尾崩溃的根源在于头部与尾部样本的梯度在范数与方向上存在巨大失配，导致参数更新在全局平均时发生毁灭性的方向抵消；更关键的是，如果在长尾设定下直接引入动态路由的混合专家系统（MoE），门控网络自身的梯度同样会被海量头部样本淹没，使得门控网络根本没有足够的动力将尾部样本分流给专用专家。

本文从条件流匹配（Conditional Flow Matching）的本质几何结构切入：条件流匹配并不学习单一的速度场，而是学习由条件索引的一族概率路径集合，条件取值天然对应着速度场的局部路径几何。核心 idea：将条件变量作为流匹配概率路径几何与梯度方向的结构代理，基于外部粗粒度先验（频率层级/健康分布或文本语义聚类）对条件空间实施静态确定性划分，并在冻结主干的 Transformer 前馈层中为各分区指派独立的残差适配器，在保持流匹配目标与采样器零侵入的同时从优化底层阻断头尾梯度相互干扰。

方法详解¶

整体框架¶

GRASP 是一种非侵入式的架构干预方案，它完全保留了基础扩散 Transformer 的流匹配损失函数与推断期采样流程，仅通过显式约束参数梯度的吸收范围来重塑优化轨迹。整个流程以冻结权重的预训练流匹配骨干网络（FLUX.1-dev，包含 19 个双流文本-图像 Block 与 38 个单流图像 Block）为基础，首先将输入条件映射为预设的静态分区索引，随后在双流 Block 的前馈网络（FeedForward）中挂载并激活对应专区的残差适配器，最终将共享骨干输出与专区残差修正相叠加以合成目标速度场。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["多模态输入<br/>图像潜变量 z + 条件 c"] --> Route["静态条件分区 π(c)<br/>确定性查表 / 离线聚类"]
    In --> Backbone["冻结主干网络 Fl(hl)<br/>FLUX.1 扩散 Transformer"]
    Route --> Adapters["引导式残差适配器 Ak,l<br/>专属于分区 k = π(c)"]
    Backbone --> Merge["残差融合<br/>hl+1 = Fl(hl) + Ak,l(hl)"]
    Adapters --> Merge
    Merge --> Out["条件速度场预测 vθ(z, t, c)<br/>无侵入式流匹配 ODE 采样"]

关键设计¶

1. 条件空间确定性静态分区：以路径几何代理优化组内梯度一致性

长尾优化的理想目标是最大化参数更新的梯度一致性（Gradient Consistency, GC），即样本子集梯度与全集梯度之间的余弦相似度：

\[GC_\theta(X') = \frac{g_\theta(X) \cdot g_\theta(X')}{\|g_\theta(X)\| \|g_\theta(X')\|}\]

然而在大型 Transformer 上在线计算与聚类样本梯度在计算上不可行。GRASP 指出，在条件流匹配中，条件变量 \(c\) 直接索引了连接基础分布与目标数据分布 \(p_\text{data}(\cdot | c)\) 的概率路径族。相同或相似条件的样本处于同一路径簇，其对速度场所诱导的梯度更新天然具有相容的局部几何方向；相反，不相干条件的样本更新则会引发破坏性抵消。因此，条件变量本身就是最优的梯度对齐代理。

为了彻底避免学习型门控在长尾分布下被头部梯度主导而饿死尾部专家的缺陷，GRASP 采用静态映射函数 \(\pi: \mathcal{C} \to \{1, \dots, K\}\)： - 频次与健康先验分区（针对医学场景）：在 MIMIC-CXR-LT 上，将样本按类别频率划分为头部（Head）、中部（Medium）、尾部（Tail）三类，并结合胸片特有的健康/病理二元先验单独划分出一个健康（Healthy / No Finding）集群，共构成 \(K=4\) 个专家分区。这强制将占绝对多数的健康样本与病理样本隔离开来。 - 文本特征离线聚类分区（针对通用场景）：在类别繁多或缺乏显式先验的 ImageNet-LT 上，直接提取冻结文本编码器的特征表示 \(f(x_i) = \psi(t_i)\)，通过无监督的二分 K-means（Bisecting K-means）聚类构建 \(K=4\) 个静态语义分区。

静态分区保证了每一个稀有尾部样本都必然完整且确定地更新其指派的专家参数，杜绝了动态路由中尾部专家退化的问题。

2. 引导式残差适配器：全局表征共享与局部动态解耦

为了在参数高效的同时为长尾类别提供充足的适应容量，GRASP 在 Transformer 模块中插入并行残差分支。对于第 \(l\) 个 Transformer 层的隐状态 \(h_l \in \mathbb{R}^d\)，冻结基座变换记为 \(F_l(h_l)\)，专属于分区 \(k = \pi(c_i)\) 的适配器输出为 \(A_{k,l}(h_l)\)。层间更新遵循残差加法：

\[h_{l+1} = F_l(h_l) + A_{\pi(c_i), l}(h_l)\]

在全局流匹配速度场层面上，该设计将目标速度场严谨地解耦为全局共享基础速度场与专家专有残差修正之和：

\[\hat{v}_\theta(z, t, c) = s_\phi(z, t, c) + a_{\pi(c), \psi_{\pi(c)}}(z, t, c)\]

其中 \(s_\phi\) 由预训练基座提供，承载跨条件的通用视觉与解剖结构；\(a_{\pi(c), \psi_{\pi(c)}}\) 则专注于拟合特定频段或病理特征的局部修正。每个适配器采用双层 MLP 结构：\(A_{k,l}(h) = W^{(2)}_{k,l} \sigma(W^{(1)}_{k,l} h)\)。与常见低秩自适应（LoRA）采用极低瓶颈维度（\(r \ll d\)）不同，GRASP 的中间隐藏维度设为 \(d' \approx d\) 甚至 \(d' > d\)（如 756d，或高维变体），避免了在跨越巨大域间差距时的表示容量瓶颈。参数更新机制保证了第 \(k\) 个专家的参数 \(\psi_k\) 仅接收指派至该分区的样本梯度：\(\nabla_{\psi_k} \mathcal{L} = \mathbb{E}[\mathbf{1}\{\pi(C) = k\} \nabla_{\psi_k} \ell]\)，从根本上隔离了头部梯度对尾部专家的冲刷。

3. 零侵入架构与推断期少数类引导正交组合

由于 GRASP 完全未触碰扩散模型的训练目标与数值采样器，它与推断期的长尾增强技术天然正交、可即插即用协同工作。在推断阶段，GRASP 可以直接叠加自引导少数类采样（Minority Guidance, MG）：GRASP 在训练期奠定了坚实的高质量条件分布边界，而 MG 则在推断阶段重新调整采样算力、引导流轨迹朝向极端低密度子流形演进。二者叠加不仅未产生负面干扰，反而在所有多样性指标上取得了超越任一单独方法的最佳表现。

一个完整示例¶

以合成极罕见病症“气腹症（Pneumoperitoneum）”（MIMIC-CXR-LT 训练集中仅有 32 例真实样本）为例： 1. 静态分区映射：系统根据病理标签查表，将气腹症判定为尾部病理类别，确定性映射至尾部专家 \(E_0\)（\(\pi(c) = 0\)）；而高频的“无异常（No Finding）”（训练集包含 53,260 例）则映射至健康专家 \(E_3\)（\(\pi(c) = 3\)）。 2. 前向推理计算：输入潜变量与文本嵌入送入 FLUX.1 骨干网络。在前 19 个双流 Block 的前馈层中，气腹症样本仅经过激活的适配器 \(A_{0,l}\)，而其余适配器（如 \(A_{1,l}, A_{2,l}, A_{3,l}\)）被严格旁路。 3. 反向梯度传播：气腹症样本产生的损失梯度只反向传播至 \(A_{0,l}\) 的投影矩阵；占主导地位的 5.3 万例健康样本梯度则全部导入 \(A_{3,l}\)。尾部专家的参数免受来自健康大类的方向冲刷，成功保留了膈下游离气体这一微小但关键的放射学特征。

损失函数 / 训练策略¶

训练采用标准的条件流匹配目标，未引入任何辅助正规化项或门控平衡损失：

\[\mathcal{L}_{\text{CFM}} = \mathbb{E}_{t, x_0, x_1, c} \left[ \| v_\theta(x_t, t, c) - u_t(x_1 | x_0) \|^2 \right]\]

其中 \(x_0 \sim \mathcal{N}(0, I)\) 为高斯噪声，\(x_1 \sim p_{\text{data}}(x|c)\) 为目标图像潜变量，线性概率路径为 \(x_t = (1-t)x_0 + t x_1\)，目标速度场为 \(u_t(x_1 | x_0) = x_1 - x_0\)。 - 微调配置：以 FLUX.1-dev 为基座，CFG Scale 设为 5，使用 8 张 NVIDIA H100 GPU，批大小为 8，总计微调 10,000 步。全量微调耗时约 30 小时（240 GPU-hours/seed），而 GRASP 适配器微调仅需 7-12 小时（56-96 GPU-hours/seed）。 - 轮询重采样策略（可选）：在每个训练 step 强制从各分区抽取 1 个样本，使专家利用率分布更加平滑；实验表明即使不启用该重采样，GRASP 依然保持高度稳健。

实验关键数据¶

主实验¶

在 MIMIC-CXR-LT 上的胸片特征空间（DenseNet-121 特征）评估中，GRASP 在图像质量（FID）和覆盖多样性（Coverage、IRS）上均显著超越基线。

方法	全标签 Cov ↑	全标签 IRS ↑	全标签 FID ↓	宏平均 Cov ↑	宏平均 IRS ↑	宏平均 FID ↓
全参数微调 Full FT (class)	.55±.01	.11±.0	.06±.01	.62±.16	.17±.16	.14±.11
全参数微调 Full FT (text+class)	.65±.01	.16±.01	.05±.01	.70±.13	.20±.15	.12±.10
MoCLE [14] (动态门控 LoRA)	.84±.01	.32±.01	.04±.0	.85±.12	.35±.17	.09±.07
单一顺序 Adapter (756d)	.86±.01	.33±.02	.03±.0	.85±.12	.37±.17	.07±.07
MoE-GRASP (学习型门控 LowDim)	.938±.003	.609±.013	.013±.001	.907±.005	.456±.003	.054±.002
GRASP (756d)	.90±.01	.35±.01	.02±.0	.92±.10	.53±.17	.05±.06
GRASP-LowDim	.94±.0	.50±.0	.01±.0	.92±.09	.56±.12	.04±.05
GRASP-HighDim	.92±.01	.41±.01	.02±.0	.91±.10	.52±.10	.04±.05
Vanilla + MG [43] (仅推断引导)	.62±.02	.15±.01	.05±.01	.67±.02	.19±.02	.13±.01
GRASP-LowDim + MG [43]	.94±.0	.63±.02	.01±.0	.92±.01	.60±.01	.05±.0

在合成数据支持下游 DenseNet-121 疾病分类任务中（NIH-CXR-LT 测试集，13 个病症类别），GRASP 合成数据展现出极佳的泛化支持能力：

训练数据来源	非零 F1 类别数 (#nz / 13) ↑	宏平均 Macro F1 ↑	无异常 No Finding F1	相对 GRASP-HD 显著性检验 (p)
真实数据 Real MIMIC-CXR-LT (上限)	3	.079±.007	.881±.002	–
全参数微调 FT (class)	1	.068±.000	.882±.000	.004
全参数微调 FT (text+class)	1	.068±.000	.882±.000	.004
MoCLE [14]	5	.069±.002	.881±.001	.004
单一顺序 Adapter (756d)	1	.068±.000	.882±.000	.031
MoE-GRASP (动态门控)	7	.073±.004	.873±.010	.027
GRASP (756d)	9	.076±.004	.872±.006	.203
GRASP-LowDim	9	.077±.005	.868±.010	.359
GRASP-HighDim	8	.080±.006	.848±.020	–

消融实验¶

针对条件划分机制（频次标签 vs 文本特征嵌入聚类）及样本重采样（Resample）的消融对比（DenseNet 特征空间）：

划分依据 (Label / Embed)	轮询重采样 (Resample)	覆盖率 Cov ↑	检索分数 IRS ↑	FID ↓	核心结论说明
Label (频次+健康)	关 (−)	.91±.10	.52±.10	.04±.05	基准静态标签分区，表现稳健
Label (频次+健康)	开 (✓)	.92±.09	.52±.12	.04±.05	重采样平滑专家利用率，微弱提升多样性
Embed (文本嵌入聚类)	开 (✓)	.91±.10	.52±.16	.05±.06	无人工标签介入下表现逼近标签分区

各划分方案下 MIMIC-CXR-LT 的专家实际计算负载占比： - Label 静态分区（无重采样）：E0: 13.01%, E1: 10.34%, E2: 15.44%, E3 (健康): 61.20%（呈现出高度偏态利用率，但由于条件纯度极高，生成质量最佳）。 - Label + 重采样：E0: 24.03%, E1: 19.78%, E2: 17.08%, E3: 39.11%。 - Embed 聚类 + 重采样：E0: 25.64%, E1: 26.46%, E2: 25.81%, E3: 22.09%（各专家负载极均衡）。

关键发现¶

静态路由击败动态门控：在长尾极端场景下，MoE-GRASP（动态门控）虽然在头部类别上展现出高覆盖率，但在尾部类别显著退化；而静态 GRASP 凭借确定的梯度分配，在尾部病症上始终维持高梯度一致性（前 1,000 步优化梯度一致性显著高于 MoCLE 和单适配器）。
负载失衡不等于性能劣势：实验证明，Label 划分中尽管健康专家占用了 61.2% 的样本，但其病理专家的条件纯度更高、梯度方向更统一，生成效果优于强行均匀划分的方案，强有力地佐证了“条件空间几何纯度是梯度代理核心”的理论假定。
下游分类覆盖面大幅拓展：在极度困难的下游分类中，全参数微调仅能在高频类（No Finding）上取得非零 F1，而 GRASP 在 13 个病症中激活了 9 个类别的有效分类能力（非零 F1），甚至持平了真实训练数据的表现。

亮点与洞察¶

将生成条件确立为流匹配梯度的几何代理：深刻洞察了条件流匹配的速度场路径结构，揭示了条件聚类直接等价于概率路径切空间梯度的对齐，避免了昂贵的动态梯度计算。
反直觉的“静态优于动态”长尾设计：常规 MoE 盲目推崇可学习门控，本文明确指出动态门控在极端长尾中会被头部样本“梯度绑架”，提出静态分区反而是保护尾部更新的确定性护城河。
卓越的可组合性与零侵入工程哲学：不魔改 Flow Matching 损失、不破坏 ODE 求解器，使得架构优化与推断阶段引导（如 Minority Guidance）无缝融合，实现了 \(1+1 > 2\) 的性能倍增。

局限与展望¶

依赖先验分区的可用性与颗粒度：静态划分函数在训练前即被冻结，若缺乏领域专业知识（如医学健康分类）且文本嵌入表征粗糙，分区的最优性将受限，无法在训练中动态纠错。
极端长尾类别样本绝对量不足：对于样本数极低（例如只有数例）的极端病症（如主动脉钙化、气纵隔等），即使梯度完全隔离，适配器依然缺乏足够的数据以形成稳健特征，其下游分类 F1 仍为零。
未来方向：探索半静态路由机制（即大类静态锚定 + 细粒度自适应微调），并将此范式推广到三维医学体数据（CT/MRI）的多模态流匹配生成中。

评分¶

新颖性: ⭐⭐⭐⭐☆ 首次从流匹配路径几何角度论证条件空间作为梯度对齐代理，提出长尾下静态分区胜过动态门控的深刻论断。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖胸部 X 光（MIMIC-CXR-LT、NIH-CXR-LT）与通用 ImageNet-LT，覆盖多特征空间评估、梯度动力学校验及下游分类下游全流程检验。
写作质量: ⭐⭐⭐⭐⭐ 论证脉络极其清晰，从梯度失配的优化病理切入，逻辑闭环，图表详实。
价值: ⭐⭐⭐⭐⭐ 为扩散 Transformer 解决长尾生成缺陷提供了轻量、非侵入且高度实用的全新工业与学术范式。