[论文解读] Bio-xLSTM: Generative modeling, representation and in-context learning of biological and chemical sequences

Niklas Schmidinger, Lisa Schneckenreiter+8 more|arXiv (Cornell University)|Nov 6, 2024

Machine Learning in Bioinformatics被引用 5

一句话总结

Bio-xLSTM 将 xLSTM 架构定制用于 DNA、蛋白质和 SMILES，使得长上下文生成建模、丰富表示和带线性运行时的上下文学习成为可能。

ABSTRACT

Language models for biological and chemical sequences enable crucial applications such as drug discovery, protein engineering, and precision medicine. Currently, these language models are predominantly based on Transformer architectures. While Transformers have yielded impressive results, their quadratic runtime dependency on the sequence length complicates their use for long genomic sequences and in-context learning on proteins and chemical sequences. Recently, the recurrent xLSTM architecture has been shown to perform favorably compared to Transformers and modern state-space model (SSM) architectures in the natural language domain. Similar to SSMs, xLSTMs have a linear runtime dependency on the sequence length and allow for constant-memory decoding at inference time, which makes them prime candidates for modeling long-range dependencies in biological and chemical sequences. In this work, we tailor xLSTM towards these domains and propose a suite of architectural variants called Bio-xLSTM. Extensive experiments in three large domains, genomics, proteins, and chemistry, were performed to assess xLSTM's ability to model biological and chemical sequences. The results show that models based on Bio-xLSTM a) can serve as proficient generative models for DNA, protein, and chemical sequences, b) learn rich representations for those modalities, and c) can perform in-context learning for proteins and small molecules.

研究动机与目标

• 激励并开发超越 Transformer 架构的生物和化学序列的长上下文语言模型。
• 将 xLSTM 调整为领域特定变体（DNA-xLSTM、Prot-xLSTM、Chem-xLSTM），并实现生成、修复与就地上下文学习（ICL）。
• 在基因组学、蛋白质和化学序列任务上评估 Bio-xLSTM，并与最先进基线进行比较。
• 展示对 DNA 的 RC 等变性，并评估在分类和设计任务中的下游性能。
• 展示就地上下文学习能力和不进行微调的领域条件生成。

提出的方法

• 用针对生物和化学序列调优的 sLSTM 与 mLSTM 块扩展 xLSTM。
• 开发三种领域特定变体：DNA-xLSTM、Prot-xLSTM、Chem-xLSTM，配备上下文窗口和 RoPE 以实现长距离依赖。
• 实现建模模式：因果语言建模 (CLM)、掩蔽语言建模 (MLM)、中间填充 (FIM) 以及就地上下文学习 (ICL)。
• 通过事后连接 (PH) 或参数共享 (PS) 纳入反向互补（RC）等变性。
• 在人体基因组上对 DNA-xLSTM 进行 RC 等变变体训练，并与 HyenaDNA、Mamba、DNA-Mamba 及 Transformers 进行比较。
• 用 FIM 对未对齐的同源体执行同源性感知、无比对输入训练 Prot-xLSTM，并评估其生成能力及变体适应性预测。
• 训练 Chem-xLSTM 以进行无条件 SMILES 生成和领域条件的 ICL，并在 SMILES 现实性上进行评估。
• 使用 RoPE 评估长上下文能力，蛋白质上下文规模达 262k 个标记，DNA 达 32k。

Figure 1: Overview of Bio-xLSTM. Top left: xLSTM for natural language processing tasks. Top right: Considered modeling approaches for biological sequences: masked language modeling, equivariance to reverse complementary sequence, and in-context learning. Bottom left: DNA-xLSTM models are trained on

实验结果

研究问题

• RQ1 Bio-xLSTM 变体是否能够以线性内存规模有效建模长序列的生物与化学序列？
• RQ2 DNA-xLSTM、Prot-xLSTM、Chem-xLSTM 是否在各自任务上提供相对于领域特定 Transformer 与基于 SSM 的模型的有竞争力或更优性能？
• RQ3 RC 等变性设计（PH/PS）是否对 DNA 建模及下游任务有益？
• RQ4 Prot-xLSTM 是否能够进行同源性条件的就地上下文学习，用于生成设计和残基级预测？
• RQ5 Chem-xLSTM 是否能够在不微调的情况下实现分子生成的领域条件就地上下文学习？

主要发现

• DNA-xLSTM 在人类基因组上进行 CLM 与 MLM 预训练，参数量为 2M，表现胜过 Transformer、Mamba 和 HyenaDNA。
• DNA-xLSTM-2M (PH/PS) 在下游基因组分类任务中，在 18 项任务中的 12 项达到或超过基线，且参数数少于 2M。
• Prot-xLSTM-102M 在同源性条件下的蛋白质生成中实现更优的困惑度和生成质量，特别在更长上下文下超越 ProtMamba 和 Transformer++ 基线。
• 尽管总训练 token 少，Prot-xLSTM-102M 的表现仍优于 ProtMamba-107M，表明高效的长上下文学习。
• Chem-xLSTM 在无条件 SMILES 生成中实现最低的 Fréchet ChemNet Distance (FCD) 并具备竞争力的困惑度，表明输出具有现实的化学性质。

Figure 2: Pre-training of 2M-parameter DNA models on the human reference genome (GRCh38). Models are trained at single-nucleotide resolution with a context length of 1024 bases. Left: causal language modeling . Learning curves display NTP loss ( $\downarrow$ ) on a test set, plotted against the numb