581. Google - Fast or Slow? Predict AI Model Runtime | predict-ai-model-runtime

祝贺所有完成Fast or Slow竞赛的参与者,并衷心感谢我的优秀队友@drhabib。在竞赛公告发布后,我们对其中涉及的有趣问题感到非常兴奋。然而,经过更详细的分析发现,评估仅在少量图上进行,这意味着评估可能非常不稳定,出现两位数大幅波动的风险很大(我很好奇组织者是否采用了基于配置的划分而非基于图的划分来避免这种情况?)。因此,我们将这场比赛视为一次抽奖,只投入了约1.5周的时间,同时还在参加另一场竞赛。
文献中经常考虑的方法通常基于早期融合或晚期融合,即配置输入的注入要么在最初进行,要么在图表示被压缩为单个嵌入向量之后进行。然而,第一种方法由于必须为每个配置独立执行计算,计算效率较低,当需要考虑数千个配置且图包含数万节点时会出现问题。与此同时,晚期融合方法在将所有节点压缩为单个嵌入向量的融合时刻会丢失节点特定信息,即模型可能难以将特定配置分配给相应的节点。
我们方法的关键思想是使用中间融合,它在计算效率和融合时刻图表示的表达能力之间提供了良好的平衡。因此,我们可以同时训练数千个配置和大型图(最多10^4个节点),而不会产生显著开销。在批次中使用大量配置对于像listMLE这样的序列排序损失至关重要。同时,在融合时刻,节点不会被压缩成单个嵌入向量,特定配置可以直接与节点关联。
一些节点特征描述了数据和内核的维度,维度的乘积可能高达10^6。这对模型来说是不可接受的,因此我们简单地对输入特征取对数并加3,以获得可比较的尺度。在我们的实验中,我们还尝试通过定制组织者的代码提取额外的29个图特征,但结果与使用基本的140个特征相当。
在布局管道中,我们没有将所有内容存储在内存中(这可能比tile数据大得多),而只加载图特征(仅几GB)。同时,配置以1000个为一块进行采样,并在训练期间以实时方式加载,而不是加载所有配置却只使用其中1000个。这消除了数据加载瓶颈。
我们的模型基于针对消费级硬件有效工作的图数据Transformer。
(Tile模型)具体来说,tile模型使用一系列残差图块,在全局混合之前创建局部邻域表示,然后通过自注意力块进行全局混合,最后通过交叉注意力Transformer块检索特定配置的节点信息。使用此设置,训练达到0.97+减速指标的模型仅需不到1分钟。我们的纯tile模型在私有和公共排行榜上的分数分别为0.196和0.197(满分0.2)。不同折的组合在私有和公共排行榜上分别获得0.196和0.198的分数。
(Layout模型)在layout模型中,我们必须将标准点积注意力替换为多头线性注意力,其复杂度与token数量呈线性关系。这种注意力机制能够在整个10^4节点图中实现全局感受野,同时计算成本相对较低,使得在4090或A6000等GPU上运行12层Transformer成为可能。我们将注意力中的ReLU非线性函数替换为更稳健的ELU+1函数。由于掩码填充token的困难,我们未在tile模型中使用SLA。同时,Transformer块中的MLP部分被图网络替代,使相邻节点间的特征能够局部混合。我们考虑了APPNP、SAGE、GAT、GPRGNN图模块,其中APPNP获得了稍好的结果。我们同时考虑了单向和双向图模型。在一组实验中,我们还考虑了4层和6层的DiFFormer模型。图Transformer之后是可配置节点的池化,将其与相应配置连接并通过MLP混合,最后在所有可配置节点间进行全局池化并预测整体排名。模型架构如帖子开头的图所示。
通过中间融合方法,即使是相对宽(n=256)且深(最多12层)的模型也非常快速,在tile设置上每轮仅需几秒钟,XLA布局上每轮几分钟,NLP布局上每轮30分钟到1小时,同时使用多达1000个配置。tile模型收敛需要约5轮,layout模型需要约16轮。
我们使用了适用于Transformer的标准AdamW优化器,采用余弦退火,学习率lr=5e-4,权重衰减wd=1e-2。tile模型训练5轮,layout模型训练16轮。在某些训练中,我们使用了权重EMA平均和梯度累积。我们使用listMLE作为损失函数,这是因为我们在每个批次中考虑了给定图的约1000个配置。
由于XLA和NLP图在图大小上表现出巨大差异,我们决定将其拆分为两组独立训练的模型。具体来说,XLA图非常庞大(消耗大量显存)且配置数量较少。而NLP图相对较小但配置数量巨大(即训练时间更长)。因此,我们可以在24GB显存GPU上训练每图256个配置的XLA模型,在48GB显存GPU上训练每图256-1000个配置的模型。而NLP模型可以轻松训练每图1000个配置,显存占用低于24GB。
在运行实验时,我们很快意识到评估集上的评估并不稳定,几乎所有考虑的设置与用不同种子重新运行相同设置相比,都获得了统计上不显著的差异。这使我们无法评估不同设置下模型的真实改进,并报告消融研究。
我们设置了11折交叉验证方案(验证集+训练集的10折划分),以提供足够的评估精度来区分所考虑设置中的微小性能差异。然而,鉴于我们在比赛中仅剩约1周时间,且NLP布局每折总训练时间为8-16小时,我们没有在此方面进一步深入,仅进行了几折实验以获得多样性。
由于我们无法可靠评估性能,我们的主要策略是运行尽可能多的配置并将它们混合在一起,以减轻可能波动的影响。我们为XLA和NLP模型构建了10-20种不同的配置。此设置在公共和私有排行榜上分别获得了0.712和0.696分。
此外,我们还选择了公共排行榜上得分最高的提交,该提交包含2个APPNP XLA模型和2个DIFFormer NLP模型。它在公共和私有排行榜上分别达到了0.721和0.703分。这对我们来说是一个非常令人惊讶的结果。
我们最好的私有排行榜提交在公共和私有排行榜上分别为0.706和0.715分。在此提交中,我们使用了带有门控单元而非残差和的APPNP模型。