AI 如何学习读取 DNA 和声音

快速阅读: 据《黑客 Noon》最新报道，本文研究选择性状态空间模型（SSMs），提出一种新的选择机制以提升模型效率，并在多种任务上验证其性能。结果显示该机制在某些场景下优于传统方法，但在长格式音频任务中效果较差。文章还对比了模型的计算复杂度及内存占用情况。

作者：（1）顾阿尔伯特，卡内基梅隆大学机器学习系，同等贡献（[email protected]）；（2）特里·道，普林斯顿大学计算机科学系，同等贡献（[email protected]）。

目录
摘要与
1. 引言
2. 状态空间模型
3. 选择性状态空间模型和
3.1 动机：选择作为压缩手段
3.2 使用选择改进SSMs
3.3 选择性SSMs的有效实现
3.4 简化SSM架构
3.5 选择机制的特性
3.6 模型的其他细节
4. 实验评估
4.1 合成任务
4.2 语言建模
4.3 DNA建模
4.4 音频建模与生成
4.5 速度和内存基准测试
4.6 模型消融
5. 讨论
6. 结论、致谢与参考文献

附录A 关于选择机制的讨论
附录B 相关工作
B.1 S4变体与衍生品
B.2 SSM架构
B.3 与RNN的关系
B.4 线性注意力
B.5 长上下文模型
附录C 选择性SSMs的机制
附录D 为选择性SSMs设计的硬件感知算法
附录E 实验细节与附加结果
E.1 合成任务
E.2 语言建模
E.3 DNA建模
E.4 音频细节
E.5 效率基准

附录E.3 DNA建模
E.3.1 预训练细节
我们更详细地描述了HG38预训练任务的数据集和训练程序。
E.3.2 扩展：模型大小细节
模型。我们考虑的模型包括：
• Transformer++：具有改进架构的Transformer，特别是使用RoPE位置编码（苏等，2021）。非正式地说，我们发现这些比Vaswani等（2017）提出的普通位置编码明显更好。
• HyenaDNA：Nguyen、Poli等人（2023）以及Poli等人（2023）提出的Hyena模型，大致是一个Transformer，其中MHA块被H3块替换，该块使用由MLP参数化的全局卷积。
• Mamba：标准Mamba架构。
模型大小。我们使用以下模型大小。
E.3.3 扩展：序列长度细节
E.3.4 物种（大猿）分类
训练包括10个周期，每个周期包含1024个梯度步数。每个梯度步数使用批量大小为64，所有样本都是通过独立随机抽样一个物种、均匀随机抽样一条染色体，然后均匀随机抽样一段连续的DNA来获取的。物种分类任务的结果见表13。

E.4 音频细节
E.4.1 YouTubeMix音频预训练
模型。我们使用一个每阶段有3个块（3 × 5 = 15个总Mamba块）、池化因子p = 16、外维度D = 64的模型，大约有350万个参数。数据集。数据以8位μ律编码，因此模型在词汇量为256的情况下对离散标记进行建模。数据集由最长1分钟或长度为960000的片段组成，这些片段被子采样并分割成任何所需序列长度的段。由于架构涉及两阶段的16倍池化且我们希望最终的序列长度是8的倍数以提高硬件效率，最长可能的序列是468 × 2048 = 958464。其余的序列长度由依次将其减半并向上取最近的2048的倍数组成。表14列出了图7中使用的规格。除了变化的批量大小外，训练集中有效片段的数量在不同序列长度之间有所不同（例如，图中的每个周期的训练步数不恒定），这可能导致缩放曲线出现拐点。训练。模型在最大学习率为0.002、20K（10%）热身步数和权重衰减0.1（类似于我们在跨领域的通用预训练配方）下训练了200K训练步。图10显示，从S4 → S6（即选择机制）的变化并不总是有益的。在长格式音频波形上，它实际上显著降低了性能，这从音频是均匀采样的且非常平滑的角度来看可能是直观的，并因此受益于连续线性时不变（LTI）方法。在消融掉选择机制后，请注意，得到的模型是Mamba块内的S4层。为了区分，我们将此称为Mamba-S4而不是默认的Mamba架构Mamba-S6。然而，在右侧，我们保留了U-Net Mamba-S4的外部层，仅消融内部层。性能差异大幅缩小；这强化了这样的假设：接近原始音频信号的层应该是LTI的，但一旦它们被外部层“标记化”并压缩，内部层就不再需要是LTI的。在这种情况下，实值SSM仍然不如复值SSM表现好。图10显示，从S4 → S6（即选择机制）的变化并不总是有益的。在长格式音频波形上，它实际上显著降低了性能，这从音频是均匀采样的且非常平滑的角度来看可能是直观的，并因此受益于连续线性时不变（LTI）方法。在消融掉选择机制后，请注意，得到的模型是Mamba块内的S4层。为了区分，我们将此称为Mamba-S4而不是默认的Mamba架构Mamba-S6。

E.4.2 SC09语音生成
自回归训练主要遵循自回归语言建模协议，例如我们使用学习率为0.002，在批量大小为16的情况下训练200000步。表4中的大型Mamba模型在每个阶段有15层，外维度D = 96，池化因子为4。请注意，这个数据集较小（训练进行了100个周期），对于这个大型模型，BPB或NLL的过拟合现象显著。然而，生成样本的自动化指标在整个训练过程中持续改善。

E.5 效率基准
扫描操作。我们比较选择性SSMs的核心操作——并行扫描（第3.3节）——与卷积和注意力，使用A100 80GB PCIe GPU进行测量。请注意，这些不包括核心操作之外的其他操作的成本，如全局卷积模型中计算卷积核或注意力中计算QKV投影的成本。我们的扫描实现融合了离散化步骤和平行扫描，避免了在HBM中产生所有大参数的成本。对于卷积，我们使用PyTorch的标准实现，分别对输入和滤波器执行FFT，频率域相乘，然后进行逆FFT以获得结果。对于序列长度L，理论复杂度为O(L log(L))。对于注意力，我们与我们所知最快的实现（带因果掩码的FlashAttention-2 (Dao 2023)）进行比较。请注意，带因果掩码的FlashAttention-2比不带因果掩码的快约1.7倍，因为大约只有一半的注意力条目被计算。

端到端推理。我们测量了一个Mamba 1.4B模型和一个未训练的Mamba 6.9B模型的推理吞吐量，与标准Transformer（GPT3架构）在1.3B和6.7B大小下的吞吐量进行对比。我们使用Huggingface transformers库中的标准Transformer实现。我们将提示长度设置为2048，生成长度设置为128。我们从批量大小1、2、4、8、16、32、64到128变化，并测量生成128个标记所需的时间。然后我们计算吞吐量（标记/秒）为批量大小×128/所用时间。我们重复测量三次并取平均值。测量是在A100 80GB PCIe GPU上进行的。

内存基准。与其他大多数深度序列模型一样，内存使用量与激活张量的大小成正比。我们报告了在1个A100 80GB GPU上的125M模型的训练内存需求测量结果。每个批次由长度为2048的序列组成。我们与我们所知最高效的Transformer实现（带有来自torch.compile的内核融合和FlashAttention-2）进行比较。表15显示，Mamba的内存需求与具有高度优化实现的类似大小的Transformer相当，我们预计Mamba未来的内存占用将进一步减少。

本文可在arXiv网站上以CC BY 4.0 DEED许可发布。

(以上内容均由Ai生成)