一种用于分子生成的混合 Transformer 架构及其量化自注意力机制

自注意力机制在经典机器学习模型中的成功，启发了开发量子类比以降低计算开销的趋势。自注意力机制将可学习的查询矩阵和键矩阵整合以计算序列中所有 token 对之间的注意力得分。这些得分随后乘以可学习的值矩阵，以获得输出自注意力矩阵，从而使模型能够有效捕获输入序列中远程依赖关系。本文提出一种混合量子-经典自注意力机制，作为大型语言模型（LLM）底层的转换器解码器的一部分。为演示其在化学领域的实用性，我们在 QM9 数据集上训练该模型进行条件生成，使用 SMILES 字符串作为输入，每个输入都标注有一组物理化学性质作为推断过程中的条件。我们的理论分析表明，查询-键点积的时间复杂性在经典模型中为 $\mathcal{O}(n^2 d)$，而在我们的量子模型中降低至 $\mathcal{O}(n^2\log d)$，其中 $n$ 和 $d$ 分别表示序列长度和嵌入维度。我们使用 NVIDIA 的 CUDA-Q 平台进行仿真，该平台旨在实现高效的 GPU 可扩展性。本工作为量子增强自然语言处理（NLP）提供了一条有前景的道路。