跳转至

DiBO: 用扩散语言模型做离线黑盒优化(DNA + 机器人形态)

会议: ICML 2026
arXiv: 2603.17919
代码: 论文页提供(here 链接)
领域: 黑盒优化 / 扩散语言模型 / 设计生成
关键词: 离线 BBO, 扩散 LLM, 双向建模, 域适应, 离线 RL

一句话总结

DiBO 把扩散语言模型 LLaDA-8B 适配到离线黑盒优化场景,用 delimiter token 统一 prompt/design/label 三类异构信号,再走「域适应 → masked-response SFT → label-improvement RL」三段后训练,让模型能在 500 条标注样本下学到 Design-Bench 多个任务的 SOTA(DNA 任务上 +8% 归一化分),单 H100 1.5 小时就能跑完一个离散任务。

研究背景与动机

领域现状:黑盒优化(BBO)在 DNA 序列、机器人形态、材料发现等领域很关键,但实验标注昂贵,没法做在线优化。离线 BBO 假设给定 \(\mathcal{D} = \{(\bm{x}_i, y_i)\}\) 静态数据集,想找出比数据集更优的新设计 \(\bm{x}^*\)。传统做法有两条路线——(a) 学个代理 \(f_\theta(\bm{x})\) 然后梯度上升(COMs、ICT、MATCH-OPT),但 OOD 时代理梯度不可靠;(b) 学逆向生成模型(CbAS、MIN、BONET、DDOM)直接 sample 高分设计。

现有痛点:(1) 自回归 LLM(OPRO、UniSO-T)虽然能直接生成 token 设计但是单向的——DNA 等设计任务里每个位点同时受前缀和后缀约束,左到右生成抓不住双向依赖;(2) 现有扩散 BBO 方法(DDOM、GTG)多数是任务特化架构在连续空间跑,没法自然融入文本任务描述;(3) 现有的离线 BBO 方法在 small-data setting(\(N \approx 500\))下严重过拟合,缺乏大模型 prior 救济。

核心矛盾:要兼顾双向建模(适合 DNA/形态)+ 文本任务描述融合(适合通用 BBO)+ 利用 LLM pretrain prior,单一架构很难做到。扩散 LLM 天生双向,但是它在自然文本上预训练,跟「设计 token + 数值 label」这种异构信号有 domain gap。

本文目标:把扩散 LLM 适配到 BBO,让它既能保留双向建模优势,又能在 small-data 下学到「prompt→更优 design」的映射,并且能用 RL 信号做精细对齐。

切入角度:用 delimiter token + 统一 prompt-response 语料解决「设计/标签 vs 自然文本」的语义角色冲突;用「先 masked joint reconstruction,再 SFT,再 RL」三段后训练逐步对齐。

核心 idea:把扩散 LLM 的「同时预测被 mask 的 prompt token + response token」当作 BBO 域适应目标,把「response label 大于 prompt 全部 label」的样本对当 SFT 数据,再把「response label - prompt label」当 reward 走 one-step log-prob RL,三段串行就能让 8B 扩散 LLM 在 500 个样本上把 BBO 做到 SOTA。

方法详解

整体框架

输入:(1) 一段自然语言任务描述(含设计语义、格式、优化目标)+ (2) 一组 few-shot 的(设计,标签)对 + (3) 一个让模型「生成更优设计」的指令。输出:一段被 delimiter 包起来的设计 + 标签 token 序列。

DiBO 在冻结扩散 LLM 之上加四件事:(a) tokenizer 扩展两组 delimiter |design-start|/|design-end||label-start|/|label-end|;(b) 域适应 DA 阶段在统一 prompt-response 语料上联合预测 prompt 和 response 的 masked token;(c) SFT 阶段只预测 response 的 masked token,以「prompt 之外的最高标签设计」作为 target;(d) RL 阶段用 label improvement 作 reward,one-step unmask 近似 log prob。推理时用一次性 greedy unmask 填满 masked response。

关键设计

  1. Delimiter token + 统一 prompt-response 语料的域适应:

    • 功能:把异构信号(prompt 文本、设计 token、标签数值)统一编码,让扩散 LLM 学会 delimiter 的语义角色。
    • 核心思路:扩展 tokenizer 加 4 个 delimiter token;每条训练样本写成 [prompt 文本][few-shot (design, label) 对][生成更优设计的指令] + [response design][response label] 的统一序列,每个 design/label 都被 delimiter 包起来。域适应目标是 \(\mathcal{L}_{\mathrm{DA}} = -\mathbb{E}[\frac{1}{t} \sum_{i=1}^{L} \mathbf{1}[q_t^i=[M], o_t^i=[M]] \log p_\theta(q_0^i, o_0^i | q_t, o_t)]\),对 prompt 和 response 的 masked token 同时做重建。few-shot 上下文 design 用 kernel 相似度从 offline pool 里挑跟 response design 相近的,避免模型学到「prompt 和 response 毫无关系」的退化映射。
    • 设计动机:直接把数值标签当普通文本会让扩散 LLM 把它们当 noise;显式 delimiter + 联合重建让模型显式学到 segmentation 边界和角色。比起 task-specific 架构(DDOM 等),这套语料方案能直接复用扩散 LLM 的预训练 prior 和注意力机制。
  2. 两段后训练:masked-response SFT + label-improvement RL:

    • 功能:让扩散 LLM 在「给定 prompt 上下文」时生成超过上下文最大值的高标签设计,并把粗粒度(SFT)和细粒度(RL)信号串起来。
    • 核心思路:SFT 阶段冻结 prompt 部分、只在 response 上做 masked reconstruction,loss 为 \(\mathcal{L}_{\mathrm{SFT}} = -\mathbb{E}[\frac{1}{t} \sum_i \mathbf{1}[o_t^i=[M]] \log p_\theta(o_0^i | q_0, o_t)]\),target response 满足 \(y(o) > \max y(\text{prompt})\),给模型一个「response 必须更优」的硬归纳偏置。RL 阶段构造 \(\mathcal{D}_{rl}\),reward \(r(q, o) = y(o) - y(q)\)(既有正也有负),不像 SFT 那样要求 strict improvement,loss 为 \(\mathcal{L}_{\mathrm{RL}} = -\mathbb{E}[\frac{1}{|o|} \sum_k p_\theta(o_k | q, o_{\text{fullmask}}) \cdot \frac{r(q,o)}{\sigma}]\),用 one-step unmask 近似 token-wise log prob,行为策略假设均匀,省掉 ratio clipping 和 KL 正则。
    • 设计动机:SFT 给「方向」(应该更优),RL 给「幅度」(更优多少);one-step unmask 是关键 trick——传统扩散 log-prob 要 iterative denoising 才稳,但 one-step 在 BBO 这种短序列上够用且节省 50× 算力,让 RL 在 H100 单卡可跑。
  3. kernel-similarity 上下文选择 + 单步 greedy 推理:

    • 功能:保证 prompt 例子和 target response 在同分布内,避免「prompt 全是垃圾、response 突变高分」的不合理学习信号;推理时单步出 candidate 提升采样效率。
    • 核心思路:训练数据构造时,给定 target response design \(o\),从 offline pool 用 kernel similarity 算 \(k(o, x_i)\) 选 top-7 作 few-shot prompt,保证局部连续性。推理时用 nfew=7 个 in-context examples,对 masked response 做一次 forward 然后 greedy fill,丢弃重复输出直到 128 个 unique candidate。
    • 设计动机:BBO 的「improve over prompt」假设在 prompt 全在低分区时崩溃——模型很难一步跳到高分区。kernel similarity 让 prompt 例子和 target 在 design space 局部接近,相当于让模型只学「小步改进」的归纳偏置。单步 greedy 推理是扩散 LLM 比 AR LLM 在 BBO 上的额外好处——AR 要 K 步生成 K-token design,扩散一次性出。

损失函数 / 训练策略

三段串行:DA 1024 步(连续 2048 步)→ SFT 1024 步 → RL 128 步。所有阶段都用 PagedAdamW8bit + Bfloat16 + 100 步 linear warmup + 常数学习率。学习率:DA 和 SFT 用 \(2 \times 10^{-5}\),RL 用 \(1 \times 10^{-6}\)(防止 RL 把 SFT 学到的「先验」拉散)。每任务有 500 样本的 offline pool,nfew=7 个 in-context 例子。

实验关键数据

主实验:Design-Bench(100th percentile 归一化分,8 seed 平均)

方法 Ant Morphology D'Kitty Morphology TF Bind 8 TF Bind 10 平均 Rank Mean ↓
\(\mathcal{D}\)(best) 0.565 0.884 0.439 0.511
Grad-mean 0.709 0.920 0.843 0.736 0.802 4.25
COMs 0.647 0.934 0.843 0.709 0.783 4.5
ExPT 0.929 0.950 0.810 0.703 0.848 4.0
OPRO(AR LLM) 0.517 0.856 0.758 0.500 0.658 14.0
DDOM(扩散) 0.590 0.929 0.739 0.497 0.689 11.25
MCTS-transfer 0.648 0.910 0.857 0.628 0.761 7.25
DiBO (ours) 0.932 0.912 0.946 0.741 0.883 3.5

DiBO 在 4 个任务中 3 个夺冠(Ant、TF8、TF10),TF Bind 8 上比最强基线领先 0.082;D'Kitty 上略输 ExPT 0.038 但仍属第一梯队;Rank Mean 3.5 + Rank Median 1.0 全场最优。

关键消融:扩散 vs 自回归 backbone(同样的 DA→SFT→RL 流程)

任务 阶段 自回归(LLaMA-3.1-8B-Instruct) DiBO(扩散) 提升
TF Bind 8 DA 0.803 0.883 +0.080
TF Bind 8 SFT 0.875 0.939 +0.064
TF Bind 8 RL 0.915 0.946 +0.031
TF Bind 10 DA 0.623 0.644 +0.021
TF Bind 10 SFT 0.633 0.704 +0.071
TF Bind 10 RL 0.682 0.741 +0.059
Ant RL 0.930 0.932 +0.002
D'Kitty RL 0.912 0.912 0.000

离散 DNA 任务(TF8/TF10)上扩散 backbone 在所有三个阶段都显著领先 AR backbone,验证「双向建模对 DNA 等强双向依赖任务有真实帮助」;连续机器人任务上差距收敛到 0,说明 6D/60D 连续设计本身没那么强的双向依赖,扩散优势主要来自离散+双向场景。

关键发现

  • 小数据 + 大模型 prior 的甜区:在 \(N=500\) 样本上做出 SOTA 是论文亮点——传统 BBO 在这个数据量下严重过拟合,扩散 LLM 8B 参数的 pretrain prior 充当了正则。
  • 三段后训练的累积收益:DA → SFT → RL 每一段都涨分(TF Bind 8 从 0.883 → 0.939 → 0.946),证明它们提供互补信号——DA 教格式,SFT 教方向,RL 教幅度。
  • OPRO 的失败说明 prompting 不够:同样的 LLM-for-BBO 思路,OPRO 只 prompt 不微调,平均 0.658 远低于 DiBO 0.883,说明 LLM 必须经过域内适应才能真正利用 pretrain prior。
  • 训练成本:单 H100、1.5 GPU 小时跑完 TF Bind 8 全流程,对 8B 模型来说极其便宜,说明扩散 LLM + 短设计序列在算力上完全可行。

亮点与洞察

  • delimiter token + 统一语料是简单但有效的桥:相比专门设计架构,扩展 tokenizer 几个 token 就把异构信号问题解决了,方法上极轻量且可迁移到其他「文本 + 数值 + 结构」混合的场景(如金融、化学)。
  • 三段后训练对应「prior → 方向 → 幅度」的认知层级:DA 让模型识别格式,SFT 给「应该更优」的硬约束,RL 用 reward 细调。这种「先粗后细」的策略对所有 LLM 适配新任务(医学问答、代码补全、科学发现)都有借鉴价值。
  • 扩散 LLM 在离散设计任务上对 AR LLM 的真正优势:消融表清晰显示在 DNA(强双向)上扩散胜出明显、在机器人形态(连续、弱双向)上扩散和 AR 几乎平手。这给「什么时候该用扩散 LLM」一个非常具体的判别准则。
  • one-step unmask 近似 log prob:把扩散模型的 RL 训练成本从 N 步 denoising 降到 1 步,对所有扩散类 RL 工作都有方法论意义。
  • 小数据 BBO 上 8B 模型可行:开源 BBO 工具普遍用 100M 以下小模型,本文证明 8B 模型在 500 样本上也能稳定训练,给 LLM-for-science 路线注入信心。

局限与展望

  • 设计空间有限:实验只覆盖了 8/10 长 DNA 和 60D/56D 连续机器人形态,对几百维以上的设计空间(蛋白质序列、电路设计)能否扩展未验证。
  • 行为策略均匀假设的合理性:RL 阶段假设老策略是均匀分布,省掉 IS ratio——但当 offline 数据有明显非均匀分布时这个近似可能引入 bias,作者承认是简化。
  • kernel similarity 选 prompt 例子的脆弱性:选 top-7 接近 target 的例子,相当于偷偷把 target 信息泄漏给 prompt(虽然 design 还是要模型生成);是否构成「数据泄漏」需要更严格讨论。
  • TF Bind 10 上的相对优势小:DiBO 在 TF10 上的 +0.005 比 TF8 上的 +0.082 小很多,可能跟 ddG label 翻转 + 任务噪声有关,论文没深究。
  • 没和 inference-time MCTS 类方法(dLLM)正面比:MCTS-transfer + dLLM 这类纯推理时方法和 DiBO 的 train-time 方法的算力-性能 tradeoff 没系统讨论。

相关工作与启发

  • vs OPRO(AR LLM 直接 prompting):他们只用 prompt 不微调;DiBO 证明微调是必须的,并给出三段式 recipe。
  • vs DDOM / GTG(扩散 BBO):那些方法用任务特化扩散在连续空间跑,没法接文本指令;DiBO 用通用扩散 LLM 把文本+设计+标签放一起。
  • vs dLLM(Yuan et al. 2026,prompt 扩散 LLM + MCTS):他们冻结扩散 LLM 走 inference-time search;DiBO 直接微调,证明 train-time 适配能拿更优结果,省 inference 算力。
  • vs PV-Tuning / DiscQuant:思路上都属于「post-training 把基础模型适配到约束子空间」,可见这种范式正在跨领域普及。
  • 启发:对任何「需要 LLM 处理结构化非文本信号」的领域(电路、分子、电力调度),「delimiter token + 统一语料 + 三段后训练」是一个通用 recipe;对 LLM-for-science,small-data + pretrain prior 的甜区是一个值得长期投入的方向。

评分

  • 新颖性: ⭐⭐⭐⭐ 把扩散 LLM 系统地适配到 BBO 是新的研究范式,delimiter + 三段后训练的设计实用且 elegant。
  • 实验充分度: ⭐⭐⭐⭐ Design-Bench 4 任务 + 10+ 基线 + 扩散/AR backbone 严格对照消融 + Top-K/RNA 等附录扩展,覆盖全面。
  • 写作质量: ⭐⭐⭐⭐ 流程图 1 简明易懂,loss 数学表达清晰,appendix 详细。略遗憾的是「为什么扩散 LLM 在连续 ant/d'kitty 上没有相对 AR 的优势」缺乏深入分析。
  • 价值: ⭐⭐⭐⭐⭐ small-data BBO 实用场景(药物、材料、机器人),1.5 H100 小时跑完单任务 + 代码开源,对工业 R&D 直接可用。