[论文解读] Toward Understanding BERT-Like Pre-Training for DNA Foundation Models

Chaoqi Liang, Lifeng Qiao+10 more|arXiv (Cornell University)|Oct 11, 2023

Genomics and Phylogenetic StudiesBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

本文提出 RandomMask，一种针对 DNA 序列的 BERT 类模型新型预训练策略，通过逐步扩大掩码边界来增加掩码难度，克服了重叠 K-mer 分词方法导致的快速收敛和次优预训练问题。RandomMask 在 28 个数据集的 7 项下游任务中均达到最先进性能，包括表观遗传标记预测中 65.83% 的 Matthews 相关系数——较基线高出 14.02%，较最先进方法高出 4.82%。

ABSTRACT

With the success of large-scale pre-training in language tasks, there is an increasing trend of applying it to the domain of life sciences. In particular, pre-training methods based on DNA sequences have received increasing attention because of their potential to capture general information about genes. However, existing pre-training methods for DNA sequences largely rely on direct adoptions of BERT pre-training from NLP, lacking a comprehensive understanding and a specifically tailored approach. To address this research gap, we provide the first empirical study with three insightful observations. Based on the empirical study, we notice that overlapping tokenizer can benefit the fine-tuning of downstream tasks but leads to inadequate pre-training with fast convergence. To unleash the pre-training potential, we introduce a novel approach called RandomMask, which gradually increases the task difficulty of BERT-like pre-training by continuously expanding its mask boundary, forcing the model to learn more knowledge. RandomMask is simple but effective, achieving state-of-the-art performance across 6 downstream tasks. RandomMask achieves a staggering 68.16\% in Matthew's correlation coefficient for Epigenetic Mark Prediction, a groundbreaking increase of 19.85\% over the baseline and a remarkable 3.69\% improvement over the previous state-of-the-art result.

研究动机与目标

• 探究现有 BERT 类预训练方法在 DNA 序列建模中的局限性，特别是直接套用自然语言处理方法而未考虑 DNA 特异性特征的问题。
• 理解重叠与非重叠 K-mer 分词方式对 DNA 基础模型预训练动态及下游微调性能的影响。
• 通过设计更有效的预训练策略，解决重叠分词导致的快速收敛与优化不足问题。
• 开发一种简单但高效的方法，增强模型学习 DNA 序列中局部核苷酸级与区域序列级模式的能力。
• 在多样化的下游 DNA 任务中实现最先进性能，包括表观遗传标记预测、转录因子结合位点识别和启动子检测。

提出的方法

• 提出 RandomMask，一种动态掩码策略，预训练过程中逐步扩大掩码跨度，从单个标记开始，逐步扩展至更长跨度。
• 应用重叠 K-mer 分词（窗口大小为 K，步长为 1）生成标记序列，以实现更丰富的局部上下文表征。
• 采用掩码语言建模目标，使模型预测逐步扩大的掩码跨度内的原始核苷酸，随时间推移逐步增加任务难度。
• 采用课程学习训练方式，随机且分步地扩展掩码边界，以防止模型过度拟合于短程依赖关系。
• 使用标准超参数和评估指标，在 28 个数据集的 7 项下游任务上对生成的模型进行微调。
• 在相同训练与微调条件下，将 RandomMask 的性能与 DNABERT、Nucleotide Transformer 和 HyenaDNA 等基线方法进行对比。

实验结果

研究问题

• RQ1 重叠 K-mer 分词对 DNA 基础模型的预训练动态及下游微调性能有何影响？
• RQ2 为何重叠分词会导致 BERT 类模型在 DNA 序列预训练中出现快速收敛但优化不足的问题？
• RQ3 一种随时间逐步增加任务难度的动态掩码策略，能否提升 DNA 基础模型的表征能力？
• RQ4 RandomMask 在多样化的下游 DNA 任务中，包括表观遗传标记预测和转录因子结合预测，相较于现有方法的性能提升程度如何？
• RQ5 RandomMask 如何增强模型捕捉 DNA 序列中局部核苷酸级变化与更广泛区域序列模式的能力？

主要发现

• 在表观遗传标记预测任务中，RandomMask 达到 65.83% 的 Matthews 相关系数（MCC），较基线 DNABERT 提升 14.02%，较此前最先进方法高出 4.82%。
• 在增强子活性预测中，RandomMask 达到 69.56% 的皮尔逊相关系数（PCC），超过基线（68.43%）及最先进方法。
• 在剪接位点预测中，RandomMask 达到 87.20% 的 MCC，较基线 DNABERT（85.44%）提升 1.76%。
• 在转录因子预测中，与基线相比，RandomMask 在小鼠数据集上提升 5.97%，在人类数据集上提升 5.63%。
• 在 7 项下游任务的 28 个数据集中，RandomMask 在 26 个数据集中达到最先进性能，展现出一致且广泛的性能增益。
• 该方法有效缓解了由重叠分词引发的对短程依赖关系的过拟合问题，从而实现更深层次、更鲁棒的表征学习。