摘要

本文介绍了一种名为M5的新型线性注意力机制，该机制能够扩展编码器仅变换器的上下文长度，达到百万级单核苷酸分辨率的基础模型，预训练于细菌全基因组。M5模型通过近似全二次注意力机制，实现了在低关键查询嵌入维度下的简单轻量级实现。该模型在单个A100 GPU上训练和测试，训练时支持高达196K核苷酸，测试时支持高达2M核苷酸。文章还展示了随着全基因组细菌序列长度的增加，M5-small模型的性能显著提升，并证明了所使用的全多头注意力近似的稳定性。

原理

M5变换器编码器使用了一种线性注意力机制，通过给定的关键查询嵌入维度d = dk，近似值exp(q ◦ k + m)，其中◦表示两个向量q, k ∈ Rdk的点积，m是一个可能依赖于q和k向量域的数。这种近似通过非线性变换θm, ϕ实现，使得exp(q ◦ k + m) ≈ θm(q) ◦ ϕ(k)。这种近似可以在给定有界区间内尽可能精确，通过使用多项式在给定有界区间内均匀近似指数函数的事实。文章还详细介绍了如何通过多项式近似和二项式公式重写这种近似，以及如何通过高数量的注意力头和小关键查询维度来实现这种近似。

流程

M5模型的训练和测试流程包括使用掩蔽的WGS（Whole Genome Sequencing）片段，最大长度为196,608核苷酸，从细菌基因组中采样。模型在单个A100-40gb GPU上进行训练和测试，输入片段长度可达200万核苷酸。训练过程中，12%的已知核苷酸被掩蔽，3%保持不变并包含在预测中。此外，训练期间还会隐藏一个长度均匀随机分布的连续片段，最大长度为4096核苷酸，占上下文长度的15%。上下文长度从1024核苷酸开始逐渐增加，每次训练会话之间翻倍。

应用

M5模型在细菌基因组学领域具有广泛的应用前景，特别是在处理超长序列和单核苷酸分辨率方面。该模型能够有效地处理和分析细菌基因组数据，有助于发现基因组特征如转录测试位点和调节元素，以及测量区域的适应性和保守性。此外，M5模型还可以扩展到多GPU/TPU设置，进一步提高模型输入上下文的处理能力。