最近在生成式人工智能领域的进展见证了模型大小和数据维度的急剧膨胀。然而,这些大型模型随之而来的是增加的计算需求,这禁止了许多传统概率推理算法的使用。迫切需要新的推理算法,这些算法足够高效,能够在大型模型和现代架构上运行,同时也足够强大,能够处理高维度和大数据集。 在这篇论文中,我们通过设计与模型规模兼容的算法成分来应对这一挑战,如并行化、摊销推理和神经函数逼近。我们提出了多种技术来改进采样和推理,从而实现更快的样本速度、样本查询的更好灵活性,以及对推理目标的更准确估计。我们的方法适用于各种架构的大型模型,如扩散模型、自回归模型和掩蔽自编码器模型。我们在图像和文本生成、游戏玩耍和机器人领域的应用中展示了这些发现。这些洞见导致了大规模生成模型高效部署的实际改进和新视角。
近年来,大规模生成模型的成功极大地改变了机器学习领域。得益于训练数据量、计算力和架构规模的显著增加,大型生成模型的改进拓宽了跨模态可能性的范围,如文本[BMR+20b]、图像[RBL+22]、视频[PX23]、游戏玩耍[FBB+22],并开启了如推理[WWS+22]和上下文中学习[WTB+22]等新能力。 然而,规模的增长带来了计算效率和算法设计的新挑战。较大的模型和更高的数据维度自然更加昂贵,需要更长的训练时间和更慢的推理时间。因此,优化大型模型的效率需求更为迫切。这些挑战导致了围绕高效训练和推理的发展,例如在多个节点上的大规模训练[SPP+19]和基于注意力架构的快速推理[DFE+22]。 在这篇论文中,我们主要关注流水线的推理阶段,包括模型的采样和推理。由于当前模型的大规模设定了计算可行性的更严格边界,因此迫切需要在算法设计上取得进展。例如,需要多个连续前向传递的采样算法变得不便地慢。随维度增长而表现不佳的传统概率推理方法变得日益禁止。在采样过程中需要访问训练数据的框架变得过时。这种模型规模增长的趋势要求概率推理方法的进步,以跟上大型模型的资源限制。 因此,我们的目标是开发足够高效以在大型模型和现代架构上运行,且强大到能处理高维数据和大数据集的推理算法。我们如何设计算法,以推进大型模型的采样和推理的前沿呢?
为新一代大型模型设计采样和推理算法时,我们可以回顾传统机器学习系统中的推理算法。相对而言,在小型模型规模和低数据维度的情况下,采样和推理算法有广泛的设计选择是可行的。采样算法可以大量依赖马尔科夫链蒙特卡洛(MCMC)技术,甚至是简单方法如拒绝采样。许多随维度增长而表现不佳的算法仍然变得流行,如维度增长表现不佳的信念传播,或数据点数量增加表现不佳的高斯过程。那些前向传递线性依赖数据维度的方法,如祖先采样,极为高效。
然而,在大型模型和大数据维度的新格局下,大多数传统的采样和推理算法已不再实用。维数灾难使得MCMC技术过于不可靠,图形模型算法的运行时间复杂度增长得太差。大量的数据点使得如高斯过程这样的非参数模型不切实际。即使是大多数算法看似基本要求的线性数量的前向传递,在处理具有数十亿参数的模型时也变得不便地慢!
显然,当处理大型模型时,实用的推理算法范围受限,因为微小的低效率被放大成巨大的障碍。鉴于这种格局,当只有预算进行几次模型的前向传递时,哪些方法仍然可行?为了取得进展,我们专注于一些与模型规模兼容的关键成分:并行化、摊销推理和神经函数逼近。使用这些可扩展的成分,我们提出了多种技术来改进大型模型的采样和推理,包括不同架构如扩散模型、自回归模型和掩蔽自编码器模型。在我们的论文过程中,我们提出了包括并行化看似顺序的计算图、使用摊销推理驯服难以处理的问题以及利用神经函数逼近进行精确近似推理的洞见。
我们将论文组织为两部分,共享同一个总体目标:1) 大型模型的采样,以及2) 大型模型的可行推理。在第一部分,我们研究了预训练模型的采样,特别是改进了从扩散和自回归模型的底层概率分布中抽取样本的技术。在第二部分,我们将注意力转向可行的推理,这涉及估计生成模型的不同量(如期望值、条件似然)的问题。我们展示了改进的架构和算法,以提高任意顺序自回归模型(AO-ARMs)和各种图形模型上的可行推理效果。