2022年3月2日，Yoshua Bengio团队在arXiv预印本提交论文，介绍了一个将主动学习框架和GFlowNets生成器结合的生物序列生成模型，用于生物序列（多肽、DNA、蛋白质等）的设计。
摘要
想要设计具有所需性质的生物序列（比如蛋白质或DNA序列），通常需要经历几次包括虚拟设计和湿实验的循环迭代过程。这些实验可以分为多个阶段，随着精度和评估成本不断提高，候选序列将被筛选出来。在虚拟设计阶段，候选序列的多样性是一个关键的考虑因素。在这项工作中，我们提出了一种主动学习算法，利用认知不确定性估计和最近提出的GFlowNets作为生物序列生成器，在每一轮之后获得有用信息（由一些预测函数定义，例如，预测的抗菌肽活性）和多样的生物序列。我们还提出了一个框架：除了奖励函数之外，再纳入带标签的生物序列数据集，从而能够加速GFlowNets的学习。之后，我们展示了此模型在几个序列设计任务中的结果。我们发现，相比于现有方法，我们的方法能够生成更多样、更新颖、更高分的生物序列。
1 简介
诸如蛋白质、DNA等生物序列有着很广泛的应用，比如治疗疾病、材料设计。例如，设计新的抗菌肽很关键，其能够缓解日益增长的耐药菌对公共健康的威胁。这一点尤其令人担忧，据世卫组织的报告显示，除非找到有效控制快速增长的耐药菌的方法，否则每年将有数百万人丧生。
考虑到蛋白靶标、生物结构及进化的多样性，多样性成为设计生物序列的关键考虑因素。廉价的筛查方法（如虚拟筛选或体外实验）可能无法很好地反映动物和人类的结果（如图1所示），所以需要更多样的候选序列。为了最大限度地增加“至少有一个候选序列满足所有需求”的概率，所选取的候选序列应该尽可能地覆盖“高分数”所涉及的所有模式。
针对特定应用，此类生物序列的设计过程包括几个阶段（可能从随机库开始），如图1所示。以抗生素药物的传统评估流程为例：1）使用近似模型去虚拟评估药物的抗菌活性，大约有106个候选药物；2）用单细胞体外实验去测评抗菌药物对特定细菌的抑制性，大约有103个候选药物；3）用小鼠做实验，大约有10个候选药物；4）人体实验，大约有个位数的候选药物。这些评价通常不是完美的，不能评价候选药物所有所需的性质。
综上所述，多样的候选物能够提升“发现满足多种性质的药物”的可能性。因此也诞生了许多能够生成多样性且性质好的方法。这种多次生成批量候选物并从评判器提供的反馈中学习的情况符合主动学习的框架。贝叶斯优化、遗传算法、强化学习也是这种问题的通用解法。值得一提的是，强化学习能够最大化奖励函数，但是通常其搜索到的候选物多样性较差。此外，由于最终目标是找到与已知药物不同的新颖设计，生成模型必须能够采样到数据分布中的边缘点。
由于候选物的多样性很重要，生成候选物的另一个有趣方法是使用生成策略，该策略可以根据奖励函数（例如，代理模型上的预测函数）按比例对候选物进行采样，并可以对独立同分布的样本进行采样，以获得一组充分覆盖奖励高的模式的候选物。其中，GFlowNets不同于MCMC，前者能够提供一种复杂的策略并采样到分数高且多样的候选物，且训练时间远小于后者。因此，本文以GFlowNets为出发点，实现生物序列的设计。
在本文中，我们提出了一个将主动学习框架和GFlowNets生成器结合的生物序列生成模型。除此之外，我们还提出了改进GFlowNets的训练方法，以提高主动学习环境中的表现。我们将此方法应用于各种的生物序列设计任务。这项工作的主要贡献总结如下：
1）一个将主动学习框架和GFlowNets生成器结合的生物序列生成模型；
2）利用线下强化学习去加速GFlowNets的训练；
3）将认知不确定性纳入到奖励函数，以此提升GFlowNets的探索能力；
4）验证所提出的方法在生物序列设计上的效果。
2 问题设定及指标选取
生物序列设计问题是在一个离散的状态空间中去搜索使得判别分数最高的目标。我们规定每一批训练，评判器（Oracle）只能评判N次，这构成了主动学习算法的N轮评估。在算法之初，数据集仅有n个样本，使用评判其对这n个样本进行评分。（假定评判器Oracle足够的准确）在每一个训练轮次，使用刚训练的生成器生成一批b个样本，并用评判器进行评判。最后将每一轮生成的样本及评判分数加到当前的数据集上，以此不断丰富数据集，这就是主动学习。
正如简介中所说，在评判器不是很完美的情况下，想要搜索到一个满足所有性质要求的候选物是很困难的。所以，我们现在感兴趣的是寻找一个多样的并且分数高的候选物的集合。我们在下面描述了理想候选物的关键特征。
2.1 表现分数
我们任务的基本标准是让候选物集合包含分数高的候选物（比如在抗菌肽设计中，分数指的是肽的抗菌能力），其可以量化为：
2.2 多样性
不仅要让候选物分数高，我们还想让候选物尽可能多地覆盖分数高的模式。其中一种测量方法为：
其中，d表示的是两个样本之间的距离。
2.3 新颖性
因为我们是从初始数据集开始，后面所生成的候选物应该尽可能和已有候选物不同。对于生成的候选物集合，我们定义它们的新颖性为：
对于上述三个指标，如果只考虑其中某一个都是没有意义的。比如某个方法可以生成多样性的候选物，但是这些候选物的分数都很低且与已知的已知的候选物很接近，这显然是没有意义的。因此，对于一种生成方法，应该充分考虑上述三个指标。
3 用作序列设计的GFlowNets
GFlowNets能够学习一个复杂的策略pi，其能够使得采样到样本x的概率正相关于此样本的奖励R(x)。这一特性使GFlowNets能够在主动学习中生成具有多样性的候选物集合。在这一部分，我们仅大致介绍了基于GFlowNets的主动学习算法及主要的结果，有关GFlowNets的具体介绍，推荐读者去阅读相关论文(见参考文献)。图2展示了本方法的概览，算法2描述了本方法的细节。
3.1 背景
我们假设X中的每一个样本x都能由一系列动作构建而成。在每一个动作执行之后，我们会得到一个部分构建的样本，被称作“状态s”。比如Bengio使用GFLowNet去按顺序地构建一个分子，每次往当前分子图添加一个原子或分子碎片。在自回归的序列生成中，动作可以被看作是将标记添加到已有的序列中。例如一个RNA序列由一个包含四种碱基AUCG的序列构成，每添加一个碱基就可以看作是进行一次动作。对于一个样本x，从空的起始状态s_0开始，逐个施加动作到达不同的状态，最后以结束状态s_f结束。此种生成过程可以看作是一个状态的轨迹。
流形
Bengio在GFlowNet中定义了轨迹的流形，其可以被解释为是一个轨迹的概率密度。每一个状态或每两个状态之间，都有流形的存在。流形与最后状态的奖励有关，GFlowNet就是通过流形去选取适当的动作从而使得生成的样本具有高分、多样的特性。为了正确评估每个状态及每个状态之间的流形，GFlowNet使用了流匹配的目标函数，即使得每个状态的流入流出尽可能相等：
在实际操作中，训练GFLowNets时采样的轨迹是由GFlowNets采样器P_Fθ和随机采样器采样而成：
这种随机采样策略提升了模型的探索能力，避免模型的训练陷入在一种或少量几种模式中。
轨迹平衡目标函数
Malkin提出了一种用于训练GFlowNets的目标函数，其能够提升GFlowNets的训练速度及鲁棒性：
算法1详细阐释了GFlowNets的训练细节：
3.2 在训练时充分利用数据
在我们的主动学习框架中，GFloNets的奖励函数（评判器）由有实验数据作为标签的数据集训练得到。以流匹配方程（公式9）和轨迹平衡方程（公式11）作为目标函数的强化学习都属于线下学习。这就允许我们在训练中使用其他策略采样来的轨迹去训练模型。这些轨迹可以从给定数据集中的x来构造，方法是从终端状态x开始为每个x采样先前状态，并根据后向转移概率对父节点进行采样。在此处的自回归学习，每一个状态只有一个父节点（因为一个生物序列是不会有分叉的），因此我们可以马上恢复出末状态x到初始状态s_0的轨迹。这就提供了线下学习的轨迹。
受到线上策略和线下策略结合的强化学习的启发，在训练GFlowNets时我们提出利用数据集作为轨迹学习的方法。在每一步训练中，我们可以用由数据集的样本构造的轨迹来丰富当前策略采样的轨迹。在每一批次训练中，通过调整γ来调整线上学习和线下学习的比例。如算法1所示。在后续实验表明，在主动学习中混合线下学习的轨迹能够提升学习的速度、增加探索的能力。
3.3 认知不确定性的应用
当模型探索的样本超出了数据集的范围时，会给从数据集中学习到的奖励函数增加不确定性。在主动学习的框架下，这种不确定性可以看作是指引模型去探索新空间的信号。Bengio认为使用奖励函数的不确定性信息能够使得GFlowNets进行更有效地探索。因此我们将模型的不确定性评估融入到GFlowNet的奖励函数中。
上述操作需要两个要点：（a）奖励函数除了能够打分之外，应该能提供对输出的不确定性评估；（b）应该能将不确定分数和奖励分数结合，形成一个标量。于是在我们的实验中，我们分别采用MC Dropout和集成模型去进行不确定性评估。对于将不确定分数和奖励分数结合的函数，我们选用Upper Confidence Bound和Expected Improvement。主动学习的框架如算法2所示：
4 相关工作
已经有多种方法被应用于生物序列设计：强化学习、贝叶斯优化、深度生成模型、基于模型的深度优化、适应性的进化方法、不借助似然的推断以及使用替代模型的黑箱优化。如第3部分所述，GFlowNets有可能通过在学习过程中摊销搜索成本（例如，与MCMC的混合时间相比），从而改善这种方法，为探索空间做出决策，从而促进模型的探索能力；允许使用不完美的数据（打分较低的数据），提升评判器的准确性及采样的多样性。
5 实验
5.1 任务及评估标准
我们通过不同的生物序列设计任务来验证所提出的GFlowNet-AL算法，所涉及的任务有：
抗菌肽设计
设计一系列具有抗菌作用的短肽（长度≤50 a.a.），选取DBAASP作为数据集，其中有3219个抗菌肽和4611个无抗菌作用的短肽；此时GFlowNets-AL的动作空间包含了20种氨基酸。
TF Bind 8 DNA设计
设计长度为8的对人类转录因子有强结合能力的DNA序列，初始的数据集包含32898个样本；此时的动作空间包含了4种脱氧核糖核酸。
GFP 蛋白质设计
设计长度为237且具有高荧光能力的蛋白质序列，初始的数据集包含5000个样本；此时的动作空间包含了20种氨基酸。
为了评估模型在这些任务上的表现，我们使用本文第2部分所提出的指标：表现分数、多样性和新颖性。对于所生成的生物序列，我们选取分数最高若干个作为代表。
5.2 基线模型和实施细节
本文用作对比的基线模型有：DynaPPO、AmortizedBO、COMs。在所有的实验中，生物数据使用独热编码的序列来表示。为了公平比较，我们限制所有基线使用相同的架构（MLP）。
5.3 实验结果
对于抗菌肽设计任务，各个基线模型的表现如表1所示（不包括AmortizedBO，因为其只能生成定长序列）。
从表1可以看出，GFlowNet-AL在保证高的多样性和新颖性的条件下，所生成抗菌肽序列的表现分数能够与DynaPPO相当。
对于TF Bind 8 DNA设计任务，各模型的表现如表2所示。
从表2可以看出，GFlowNet-AL除了在序列多样性上不如MINs（其表现分数太低）之外，表现分数和新颖性都比其他模型高，足以见GFlowNet-AL的优越性。
GFP蛋白质设计的实验结果如表3所示。
本文还做了不确定性评估的消融实验，比对了两种不确定性评估方法Ensemble和MC Dropout的效果，如表4所示。
6 结语
本文将GFlowNet模型放入主动学习的框架并把不确定性评估纳入到奖励函数，实现了分数高、多样且新颖的生物序列生成。但该方法和其他涉及代理模型（评判器）和生成策略的方法的一个局限性是，我们现在有两个独立的模型，每个模型都有其超参数。另一方面，在高维搜索空间中需要使用有效的优化或生成。
考虑到这是一个持续的学习环境，未来的工作应该探索如何有效地训练代理模型。更好的信息增益估计器、利用数据底层结构的非自回归生成模型，以及以不同保真度处理多个评判器的外循环策略，也都是未来要研究的内容。
参考资料
[1] Jain M, Bengio E, Garcia A H, et al. Biological Sequence Design with GFlowNets[J]. arXiv preprint arXiv:2203.04115, 2022.
[2] Bengio E, Jain M, Korablyov M, et al. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation[J]. Advances in Neural Information Processing Systems, 2021, 34.
--------- End ---------
