大规模基础模型的缩放与量化

摘要

网络生物学中的基础模型通过对大规模生物数据进行预训练，能够通过迁移学习在各种下游任务中实现上下文感知的预测。然而，随着预训练数据的扩展，模型规模的增长也增加了下游任务中微调和推理所需的计算资源。本文首先构建了一个包含约1.04亿个人类单细胞转录组的数据集，并预训练了规模递增的模型，定义了转录组掩码学习的缩放定律。实验证明，模型量化能够保留全精度模型的上下文基因和细胞嵌入空间，在零样本和微调应用中性能相当，同时，在相同批次大小下，微调时间仅为全精度模型的15%，内存占用仅为34%。模型量化是在保留生物学知识的同时实现资源高效微调和推理的有效方法。

主要内容

基因网络图谱的绘制能够发现关键的网络调控因子和网络校正疗法。传统方法需要大量转录组数据来学习基因间的连接关系，限制了在数据有限的场景（如罕见病）中的应用。迁移学习利用大规模通用数据集预训练模型，获得基础知识，再迁移到大量下游任务中。

此前开发的网络生物学迁移学习策略，利用约3000万个单细胞转录组预训练了基础深度学习模型Geneformer。本文首先构建了更大的预训练语料库Genecorpus-104M，包含来自多种组织和疾病状态、约1.04亿个人类单细胞转录组。为了应对当前技术检测到每个细胞中基因数量增多的情况，模型输入大小扩展至4096个基因。

随后，预训练了参数规模递增的Geneformer模型，以定义转录组掩码学习的缩放定律。结果发现，参数越多的模型，对每个数据令牌的学习速度越快（类似于其他领域的基础模型），其性能随参数增加呈幂律提升。

零样本预测：在基因层面，参数最多的模型（3.16亿参数）在所有任务中均超越了全微调的替代方法，这对于缺乏任务特定数据的场景（如罕见病）尤其有价值。

模型量化：为了解决大模型微调和推理需要更多计算资源的问题，文章实施了4位精度的模型量化（使用QLoRA技术）。

• 效率提升：在相同批次大小下，4位量化模型的微调时间仅为全精度模型的15%，内存占用仅为34%。由于内存占用更低，实际可以使用更大的批次，从而获得更显著的时间节省。
• 精度保持：尽管计算需求降低，4位量化模型在基因和细胞层面任务的少样本和零样本预测精度上与全精度模型几乎相当，无显著性能变化。
• 嵌入空间保留：量化对基因和细胞嵌入向量的影响极小，与不同细胞类型或不同基因带来的生物学差异相比可以忽略不计。量化模型也能有效衡量基因敲除模拟实验中的嵌入空间偏移。

细胞层面任务：最大模型的零样本嵌入空间在区分组织、细胞类型和疾病状态方面表现最佳。在微调后，量化模型同样能匹配全精度模型在细胞分类任务上的预测精度。

推理效率：量化模型在推理时，时间仅为全精度模型的33%，内存占用仅为53%。

总结：模型量化是实现资源高效微调和推理的有效策略，能够在保留生物学知识的同时显著节省计算成本和时间，扩大生物学基础模型的可访问性。FINISHED