跳转至

Aligning LLMs with Biomedical Knowledge using Balanced Fine-Tuning

会议: NeurIPS 2026
arXiv: 2511.21075
代码: https://github.com/TencentAILabHealthcare/BFT
领域: 医学图像
关键词: 生物医学大模型、平衡微调、认知不确定性、局部组归一化、策略初始化

一句话总结

针对生物医学文本中密集低置信度区间富集关键实体与因果链的特有不确定性结构,本文提出双尺度平衡微调(BFT),通过局部组归一化 token 重加权与序列级极值困难样本调度,在保持训练稳定性的同时充分保留领域知识,为下游稀疏奖励强化学习与生物表征构建提供了强鲁棒初始化。

研究背景与动机

大语言模型在通用领域的后训练(Post-training)极大依赖有监督微调(SFT)与强化学习(RL),但将其迁移到高度专业的生命科学与生物医学领域时,通用的置信度假设遭遇了严重的水土不服。在通用文本中,模型的低预测概率往往对应词汇表面多样性或句式风格变化等偶发不确定性(Aleatoric Uncertainty),降低这些低置信度 token 的损失权重能够过滤噪声并稳定优化过程(如动态微调 DFT)。然而,生物医学文本的知识分布呈现出本质不同的拓扑结构:低置信度区域通常并非零散分布的笔误或风格倒装,而是以长程连续片段的形式出现,高度富集了罕见生物实体名称、基因与突变位点、信号通路节点以及复杂的生化因果机制。

这就引发了现有微调范式的核心矛盾:标准 SFT 对所有 token 赋予均等权重,在反向传播中低置信度 token 产生过大梯度扰动,导致优化不稳定;而基于绝对置信度的加权方案(如 DFT)则无差别地压制所有低置信度 token,直接扼杀了模型学习罕见实体与严密因果推理链条的机会。进一步,在生物计算的下游任务(如基因扰动预测、细胞虚拟仿真)中,环境奖励通常具有极度稀疏的特性(仅在最终表型或扰动方向分类正确时给予二值结果),缺乏密集的步骤级反馈,若微调初始化阶段模型未能习得扎实的因果骨架,后续策略优化往往会迅速退化甚至陷入模式崩溃。

为了打破这一困境,研究团队通过对 PubMed Central 全文与 CRAFT 专家标注语料的系统诊断发现,生物医学文本中连续密集的低置信度片段具有极高的专家概念富集度(超过 60% 命中人工实体标注),其不确定性本质上反映了模型的认知盲区(Epistemic Uncertainty)。本文的核心 idea 是:将生物医学文本的不确定性解耦为局部偶发噪声与密集认知不确定性,提出双尺度平衡微调(BFT),在 token 尺度以局部滑动窗口均值为基准实施组归一化重加权,在序列尺度依据局部置信度极值向高知识密度困难样本重新倾斜优化预算。

方法详解

整体框架

BFT 的核心目标是在不引入额外判别模型或外部知识库的前提下,完全利用模型自身预测概率的时空分布结构实现自适应知识对齐。方法通过一个单一的超参数——滑动窗口大小 \(g\)(默认 \(g=256\)),在序列内部和序列之间两个空间尺度上协同分配学习率与梯度权重。输入为标准生物医学指令-回答样本对,模型前向传播计算各目标 token 的预测置信度序列;随后,中心化滑动窗口在序列上滑动,计算每个 token 邻域的局部平均置信度。

在 token 尺度,BFT 舍弃了全局绝对置信度门限,而是将当前 token 的绝对置信度除以其局部窗口平均置信度。若一个低置信度 token 身处高置信度上下文中(孤立异常点),其权重被压低以抑制风格扰动;若其身处同样低置信度的上下文之中(密集生化推理区),其相对比值将接近甚至达到上限 1.0,从而完整保留参数更新信号。在样本尺度,算法提取整个序列中局部窗口置信度的全局最小值,以此量化该样本包含的最弱知识瓶颈,构造有界样本加权系数 \(s_b\)。两者结合构成了无需奖励模型的双尺度加权交叉熵优化目标。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入:生物医学指令与参考文本序列"] --> Win["滑动窗口提取局部上下文置信度<br/>窗口大小 g=256 计算局部均值"]
    Win --> D1["局部组归一化 Token 级重加权<br/>局部相对置信度比值裁剪"]
    Win --> D2["序列级自适应困难样本重分配<br/>提取最弱局部区域极值赋权"]
    D1 --> Loss["双尺度 BFT 损失加权优化"]
    D2 --> Loss
    Loss --> D3["面向稀疏奖励强化的鲁棒策略初始化<br/>保留因果链与稀有实体支持下游 GRPO"]
    D3 --> Out["输出:强因果推理模型与生物表征"]

关键设计

1. 局部组归一化 Token 级重加权:解耦偶发噪声与密集认知不确定性

针对传统动态微调(DFT)无差别抑制低置信度 token 而抹除罕见专业实体的问题,BFT 提出了最简形式的组归一化权重。对于序列中第 \(t\) 个目标 token \(y_{b,t}^*\),给定模型预测概率 \(c_{b,t} = \pi_\theta(y_{b,t}^* \mid y_{b,<t}^*, x_b)\),首先定义以 \(t\) 为中心的对称局部窗口 \(W_{b,t}\)(包含自身及左右邻居至多 \(g\) 个有效 token),计算局部上下文平均置信度: $\(C_{b,t}^{\mathrm{loc}} = \frac{1}{|W_{b,t}|} \sum_{j \in W_{b,t}} c_{b,j}\)$ 进而定义 token 级归一化相对权重: $\(\rho_{b,t} = \mathrm{clip}\left(\frac{c_{b,t}}{C_{b,t}^{\mathrm{loc}} + \varepsilon}, 0, 1\right), \quad w_{b,t}^{\mathrm{BFT}} = \mathrm{sg}(\rho_{b,t})\)$ 其中 \(\mathrm{sg}(\cdot)\) 为梯度截断操作。该设计的巧妙之处在于:在稀疏低置信度区域(Group A,典型如通用过渡词),\(c_{b,t}\) 较低而 \(C_{b,t}^{\mathrm{loc}}\) 接近 1.0,比值 \(\rho_{b,t}\) 随之降低,行为退化为类似于 DFT 的抑制效应,有效避免模型拟合无规律的词汇变异;而在密集低置信度区域(Group B,典型如生化反应链路),局部窗口内的所有 token 均处于较低置信度,\(C_{b,t}^{\mathrm{loc}}\) 本身数值较小,使得相对比值 \(\rho_{b,t}\) 显著回升。例如当 \(c_{b,t}=0.10\) 时,若局部均值高达 0.87,权重仅为 0.115;但若局部均值同为 0.20,权重则跃升至 0.50,相较 DFT 恢复了数倍的梯度回传强度,精确保护了领域知识承载区。

2. 序列级自适应困难样本重分配:基于局部不确定性极值的样本预算调度

仅在序列内部调整 token 相对权重无法改变样本间的优化资源分配,常规训练容易被大批量平庸且冗长的简单文本主导梯度方向。BFT 复用了局部上下文置信度序列 \(\{C_{b,t}^{\mathrm{loc}}\}_t\),定位每个样本中最具挑战性的局域知识洼地,提取序列级极小值: $\(p_b^{\mathrm{conf}} = \min_{1 \le t \le |y_b^*|} C_{b,t}^{\mathrm{loc}}\)$ 基于该极小值构建有界样本级平衡系数: $\(s_b = \mathrm{sg}(1 - p_b^{\mathrm{conf}})\)$ 由于 \(p_b^{\mathrm{conf}} \in [0, 1]\),系数 \(s_b\) 自然严格受限在 \([0, 1]\) 之间。当某样本通篇表达顺畅、模型置信度极高时,\(p_b^{\mathrm{conf}} \approx 1\) 导致 \(s_b \approx 0\),模型不对该平庸样本耗费过多优化能力;相反,只要序列中包含任意一段密集的因果逻辑瓶颈,\(p_b^{\mathrm{conf}}\) 显著下降,即拉动 \(s_b\) 趋近 1.0。该设计与 token 级加权构成了完美的正交互补关系:\(w_{b,t}^{\mathrm{BFT}}\) 在序列内部根据局部相对难度再平衡计算预算,而 \(s_b\) 则在序列跨度上根据绝对专业难度再分配批次梯度。

3. 面向稀疏奖励强化的鲁棒策略初始化:为下游强化学习保留因果推理轨迹

在细胞微环境扰动预测与复杂生物过程推理中,环境验证多为稀疏二值奖励。若微调策略输出的推理轨迹充斥着幻觉或跳步,强化探索(如 GRPO)在长程采样中将很难命中正确的终点,导致策略熵崩溃或严重遗忘。BFT 通过前两项设计,在有监督阶段保留了完整的生化机理因果链条(相比 SFT,实体提及数提升 71%,因果连词数提升 84%,突变/药物提及率高达 95%)。该初始化使得策略采样不仅停留在最终答案分布上,而且其思维链(CoT)天然穿行于生理学上合理的生化通路上,使后续基于稀疏二值准确率奖励的强化学习探索能够在高密度有效支撑集上展开,避免了传统 SFT 模型在 RL 阶段性能反而下降的困境。

一个完整示例

以一段包含肺癌靶向治疗分子机制的回答片段为例,假设序列中包含两类典型低置信度 token: - 位置 A(句式修饰词):模型在生成过渡状语短语时遇到表面同义词选择,单 token 预测概率 \(c_{b,A} = 0.12\)。由于周围处于常规临床叙述句,滑动窗口内平均置信度较高(\(C_{b,A}^{\mathrm{loc}} = 0.85\))。计算比值 \(\rho_{b,A} = 0.12 / 0.85 = 0.141\)。该位置的梯度被大幅抑制,避免模型过度拟合特定语篇过渡习惯。 - 位置 B(关键变异与通路):模型在推导下游耐药途径时遇到“KRAS G12C 突变通过旁路激活 MEK 磷酸化导致 EGFR-TKI 耐药”,涉及连续罕见生化命名与机制推导,单 token 预测概率 \(c_{b,B} = 0.10\)。由于整个生化因果区均处于认知盲区,局部滑动窗口内平均置信度同样低迷(\(C_{b,B}^{\mathrm{loc}} = 0.18\))。计算比值 \(\rho_{b,B} = \min(0.10 / 0.18, 1.0) = 0.556\)。该 token 的权重是 DFT(仅 0.10)的 5.5 倍以上,参数更新得以充分保留。 - 序列尺度调度:遍历整条回答后,检测到该序列的最薄弱窗口平均置信度为 \(p_b^{\mathrm{conf}} = 0.18\),计算样本级权重 \(s_b = 1 - 0.18 = 0.82\)。整条富含专业推理的样本在 mini-batch 中获得了高权重的反向传播关注度。

损失函数 / 训练策略

结合上述两级权重,BFT 最小化带有序列有效 token 掩码 \(m_{b,t} \in \{0, 1\}\) 的批次归一化损失函数: $\(\mathcal{L}_{\mathrm{BFT}}(\theta) = \frac{1}{B} \sum_{b=1}^B s_b \frac{\sum_t m_{b,t} w_{b,t}^{\mathrm{BFT}} \ell_{b,t}}{\sum_t m_{b,t} + \varepsilon}\)$ 其中 \(\ell_{b,t} = -\log c_{b,t}\) 为标准交叉熵损失。基座模型选用 DeepSeek-R1-Distill 系列(涵盖 14B、32B 和 70B 版本),训练数据由剔除重叠泄露的 HealthBench Consensus 临床轨迹(GPT-5.1 合成)与 NCBI 衍生生物知识对(GPT-OSS-120B 合成)混合而成,规模为 48,503 条样本。使用 LlamaFactory 进行全参数微调,优化器为 AdamW,学习率 \(1 \times 10^{-5}\),采用线性预热(前 3% steps)与余弦衰减调度,序列截断长度为 4096,总训练轮数为 2 个 epoch(1,516 个优化步)。在 32 张 NVIDIA H20 GPU 上运行 ZeRO-3 与 bf16 混合精度,70B 模型耗时约 36 小时(~1,150 GPU 小时)。

随后的 GRPO 强化学习阶段基于 TRL 框架,在 Tahoe-100M 的 45 种细胞系扰动数据上微调 1 个 epoch,学习率设为 \(5 \times 10^{-7}\),KL 散度系数 \(\beta = 0.05\),采样组大小 \(G = 8\)。强化奖励采用极简形式:扰动方向(上调/下调/无显著变化)正确获得任务奖励 1.0,格式正确获得格式奖励 0.1,不引入任何密集进程奖励或长度偏置。

实验关键数据

主实验

在虚拟细胞扰动推理基准 VCWorld 上,评估模型在 5 种癌细胞系(C32、HepG2C3A、HOP62、Hs 766T、PANC-1)下预测化学扰动后基因表达上调、下调或无变化的联合准确率。对比闭源 SOTA 基准、基础基座模型及各后训练变体,并验证 GRPO 强化学习阶段前后的准确率变化:

方法 C32 HepG2C3A HOP62 Hs 766T PANC-1 平均准确率
VCWorld (Gemini-2.5-Flash) 0.69 ± .01 0.67 ± .02 0.72 ± .01 0.69 ± .02 0.62 ± .02 0.68
VCWorld (Base 70B) 0.42 ± .02 0.37 ± .02 0.41 ± .02 0.39 ± .02 0.40 ± .02 0.40
DFT 70B 0.54 ± .03 0.46 ± .02 0.49 ± .02 0.45 ± .03 0.51 ± .02 0.49
SFT 70B 0.58 ± .02 0.53 ± .02 0.58 ± .02 0.55 ± .01 0.57 ± .02 0.56
BFT-sample 70B (仅样本级) 0.61 ± .02 0.58 ± .02 0.62 ± .02 0.58 ± .02 0.59 ± .02 0.60
BFT-token 70B (仅token级) 0.68 ± .02 0.65 ± .02 0.67 ± .01 0.66 ± .02 0.64 ± .02 0.66
BFT 70B (完整方案) 0.72 ± .01 0.70 ± .01 0.71 ± .02 0.70 ± .01 0.65 ± .01 0.70 (0.696)
DFT 70B + GRPO 0.46 ± .02 0.39 ± .02 0.43 ± .03 0.38 ± .02 0.44 ± .02 0.42 (↓ 0.07)
SFT 70B + GRPO 0.51 ± .02 0.45 ± .02 0.51 ± .02 0.46 ± .02 0.50 ± .02 0.49 (↓ 0.08)
BFT-sample 70B + GRPO 0.65 ± .02 0.59 ± .02 0.65 ± .02 0.62 ± .02 0.61 ± .02 0.62 (↑ 0.03)
BFT-token 70B + GRPO 0.72 ± .02 0.70 ± .02 0.71 ± .02 0.69 ± .02 0.65 ± .02 0.69 (↑ 0.03)
BFT 70B + GRPO 0.76 ± .01 0.74 ± .01 0.75 ± .01 0.72 ± .01 0.71 ± .01 0.74 (↑ 0.04)

除文本生成与推理外,将微调后模型生成的实体描述文本输入至文本编码器(OpenAI text-embedding-ada-002 及领域专用 BioBERT),评估其在基因水平生物属性预测上的 AUC 表现:

编码器 对齐方法 / 基线 TF Range Dosage No-methy Lys4 GGI PPI
- scGPT (专有基础模型) 0.557 0.853 0.804 0.847 0.635 0.661
- GenePT (NCBI原生摘要) 0.585 0.891 0.910 0.939 0.806 0.681
ada-002 DFT 70B 0.519 0.808 0.761 0.793 0.598 0.604
ada-002 SFT 70B 0.534 0.838 0.769 0.812 0.616 0.622
ada-002 BFT 70B (本文) 0.629 0.913 0.927 0.910 0.823 0.769
BioBERT DFT 70B 0.495 0.806 0.755 0.777 0.566 0.583
BioBERT SFT 70B 0.505 0.836 0.763 0.798 0.584 0.596
BioBERT BFT 70B (本文) 0.591 0.872 0.879 0.864 0.741 0.707

消融实验

针对样本加权系数的潜在偏差(如长序列更易获得极低极小值)与梯度尺度的控制,在 70B 模型与统一 VCWorld 设置下进行了严谨的机制消融与替代目标对比:

机制配置 VCWorld 准确率 相对说明
完整 BFT 70B (原始极值系数) 0.696 完整双尺度联合方案
仅 Token 级加权 (BFT-token, \(s_b \equiv 1\)) 0.660 去除样本级调度,性能下降 3.6%
仅样本级加权 (BFT-sample, \(w_{b,t} \equiv 1\)) 0.600 缺乏局部 token 归一化,性能下降 9.6%
Batch 归一化样本系数 \(\widehat{s}_b\) 0.700 消除 batch 均值波动,准确率微幅提升
固定窗口位置采样系数 \(s_b^{(K)} (K=8)\) 0.690 消除序列长度带来的极值偏置,仍保留 95.5% 增益
DFT + BFT 样本系数 \(s_b\) 0.528 证明若 token 级直接使用绝对置信度,样本级增强无法挽回损失
标准 SFT 70B 0.562 基线监督微调
替代加权基线 ProFit 0.618 基于概率筛选的高价值 token 加权
替代加权基线 Anchored SFT 0.607 基于参考策略锚定的加权微调

关键发现

  • 双尺度模块互补性显著:token 级局部组归一化贡献了最大的单一跃升(从 SFT 的 0.562 提升至 BFT-token 的 0.660),而样本级调度进一步将性能推升至 0.696;若单独在 DFT 上施加样本系数仅得 0.528,证明局域相对置信度解耦才是知识保全的核心前提。
  • 强化学习初始化的分水岭:在相同的 GRPO 稀疏奖励设置下,SFT 和 DFT 的准确率分别严重滑坡 0.08 和 0.07,而 BFT 各变体均实现正向提升(+0.03 到 +0.04)。对生成轨迹的深度分析显示,BFT 初始化保留了深层因果链(平均深度 3.6 vs SFT 2.3)与罕见生化实体(召回率 0.63 vs SFT 0.41),有效支撑了稀疏探索。
  • 通用能力保留特性:在 MMLU 与 CMMLU 通用评测中,DFT 遗忘最少,BFT-token 与完整 BFT 紧随其后,而 SFT 遗忘最重,验证了 token 级抑制偶发变异对减缓灾难性遗忘的正面价值。

亮点与洞察

  • 洞察到生物医学文本具有与通用语料截然相反的“不确定性-语义密度”对应关系,即连续低置信度区间高度对应认知知识盲区而非偶发噪声。
  • 提出的组归一化加权无需训练任何辅助价值网络或引入额外超参数,仅凭滑动窗口均值这一统计量即优雅解决了领域罕见词与因果推理的梯度保留难题。
  • 搭建了生成式语言模型对齐与判别式生物学特征表示之间的桥梁,证明高质量对齐的大模型能够生成信息更为特异的基因/细胞描述文本,从而全面赋能下游组学表征与图谱集成。

局限与展望

  • 基座自引用依赖:诊断分区与训练权重的计算依然依赖模型自身的后验概率,若基座模型在某些生化小众领域存在系统性过度自信(Overconfidence),BFT 将无法感知该区域的认知盲区。
  • 知识库对齐的覆盖范围有限:在因果三元组的事实性检验中,现有结构化知识库(如 SemMedDB、DrugBank)仅能覆盖约三分之一的开放抽取三元组,其余复合状态短语与定量指标尚未获得完全自动化的事实性闭环验证。
  • 未来方向:探索引入校准良好的独立外源不确定性估计器,并将 BFT 推广至自然真实专家多轮交互数据及前瞻性临床辅助决策场景中。

相关工作与启发

  • vs Dynamic Fine-Tuning (DFT): DFT 采用绝对预测概率 \(c_{b,t}\) 作为权重以稳定训练,但其对低置信度的无差别截断直接抹杀了生物医学的专有长尾实体;BFT 引入局部均值相对比值,不仅保留了对孤立噪声的平滑抑制,且精准恢复了知识密集区的反向传播信号。
  • vs GenePT / scGPT: 传统生物表征模型要么依赖大规模单细胞测序数据预训练(如 scGPT),要么直接编码 NCBI 静态人工摘要(如 GenePT);本文证明 BFT 对齐后的通用开源大语言模型可以作为动态、精准的生物文本知识引擎,所生成的描述嵌入后在多项基因与单细胞下游任务上超越专用基准。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示生物文本中密集低置信度对应认知不确定性的结构特性,提出了极简且高效的双尺度对齐方案。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖临床问答、基因推理、细胞虚拟仿真、强化学习对齐及跨尺度生物表征,实验链条完备坚实。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰自然,消融与对照控制实验设计严密,叙述层次分明。
  • 价值: ⭐⭐⭐⭐⭐ 为生命科学与生物医学领域大模型的专业化微调提供了极具实用性与启发性的标准范式。