Google - 快或慢？预测AI模型运行时竞赛第14名解决方案

上下文

• 业务背景：竞赛概述
• 数据背景：数据集说明

方法概述

数据预处理

• 开发了2种图压缩技术降低问题复杂度（从配置节点的N跳归约和配置元图）
• 归一化数值特征，删除标准差为0的特征
• 对操作码进行独热编码

特征工程

• 创建多个节点特定特征和配置特定特征
• 创建应用于整个图的全局特征

模型设计

所有模型基本遵循以下结构：

1. 图/配置/操作码连接
2. 使用图表示进行图卷积（不同模型略有差异）
3. 全局平均池化与全局特征连接
4. MLP输出层

Tile数据集结果采用单模型设计，包含3个GraphSAGE层和3个线性层，使用ListMLE损失训练。Layout数据集结果采用模型集成，所有模型使用GeLu激活函数，但在其他方面有所不同。输出损失使用ListMLE和Pairwise Hinge。

验证方法

我们保持竞赛数据集中提供的相同训练/验证集划分。

方法细节

图归约

每个布局图被转换为两种不同的图：

1. 3跳图（从可配置节点出发）：可配置节点是不同图之间可能不同的节点，因此图归约应尽量保留这些节点。N跳图转换将保留可配置节点，以及距离任何可配置节点N跳以内的节点（和边）。通过验证分数，我确定了3跳图，但如果有时间调优，4或5跳图可能会获得更好的结果。
2. "配置定位"图：该图移除所有非可配置节点，但在任意两个可配置节点之间绘制边，只要它们之间存在不经过另一个可配置节点的路径。我的直觉是，配置不良的节点相对于下游可配置节点的相对位置可能对其对运行时的整体贡献产生有意义的影响。

特征工程

首先需要强调的是，归一化在这个问题上绝对至关重要。如果不归一化数据，就无法获得好的分数。

节点特征

• 形状稀疏度（形状和/形状积）
• 维度（活动形状的数量）
• 步幅交互
• 填充比例
• 反转比率
• 是否可配置

配置特征

为输出、输入和内核部分计算了额外的配置特征。每个特征都为这些部分重复：

• is_default（全为负值）
• active_dims（非负值的数量）
• max order（最大值）
• contiguity rank（最长连续顺序计数/活动维度）
• section variance（部分方差）

此外，我们计算了以下相似性指标：

• 输出-输入
• 输出-内核
• 输入-内核

操作码

操作码仅进行独热编码。

全局特征

• 图中最长路径长度
• 连接组件间的平均最短路径长度
• 节点数量
• is_default

引入is_default标志是因为随机分布与默认分布似乎有足够差异，具有预测价值，因为测试集中也包含此信息。它似乎为验证分数提供了小而可靠的提升。

模型架构

布局模型

对于布局问题，我们定义了以下GraphBlock，使用GAT（输出通道/2）处理3跳图的节点特征，使用GAT或GraphSAGE（输出通道/2）处理配置定位图的特征。

然后我们通过残差连接分层GraphBlock，并添加一些密集前馈层（同样带有残差），将全局特征连接到这些层。最终结果是一个略有不同版本的模型集成（变化隐藏维度、线性层等）。

以下是布局集中使用的多边缘块的代码（Tile集模型基本相同，但结构更简单。您可以在此处查看。其余代码可在此处获取）：

class MultiEdgeGATBlock(nn.Module):
    def __init__(
        self,
        *,
        input_dim: int,
        output_dim: int,
        heads: int = 4,
        with_residual: bool = True,
        dropout: float = 0.5,
        main_block: Literal["gat", "sage"] = "gat",
        alt_block: Literal["gat", "sage"] = "sage",
    ):
        """应用两种不同边缘卷积到图并连接结果的块。使用边缘掩码确定将主块应用于哪些边，将交替块应用于哪些边。"""
        super().__init__()
        output_dim_per_block = output_dim // 2

        if main_block == "gat":
            self.main_edge_block = GATBlock(...)
        else:
            self.main_edge_block = SAGEBlock(...)

        if alt_block == "gat":
            self.alternate_edge_block = GATBlock(...)
        else:
            self.alternate_edge_block = SAGEBlock(...)

        self.with_residual = with_residual
        self.output_dim = output_dim

    def forward(self, data: Data):
        # 前向传播实现...
        return data.update(new_data)

Tile模型

Tile数据集的模型基本相同，只是使用了3个图层、3个线性层，并且没有图归约（因为没有可配置节点）。

训练过程

NLP和XLA分别训练。我很想尝试统一模型，但时间和计算资源有限，而分开训练效果良好。

所有模型使用：

• AdamW优化器
• 16的批量大小
• GeLu激活函数
• 图和MLP中的LayerNorm
• 图块后的全局平均池化
• 无学习率调度器

输出从XLA-1在3个epoch后收集，以及XLA-2在训练中的2个快照（epoch 2和3结束时）基于其验证分数收集。

NLP模型甚至从未完成一个epoch，损失似乎达到平台期，而我的计算资源有限。这表明我可能可以更好地调整学习率或正则化。

集成方法

对于给定的文件ID（例如"xla:default:abc..."），每个模型输出有N（1000或1001）个预测。我们对预测进行最小-最大归一化，使其全部在0到1范围内，然后将每个模型输出的分数逐元素相加。我们使用求和后的配置分数来推导排序。

归一化在这里很重要，因为如果在使用排序损失，模型输出的数值不一定在同一尺度上。

我还尝试了简单排序平均和Borda计数，两者都有效，但不如最小-最大平均效果好。这很可能是因为这些方法无法考虑"第2名比第3名好多少"这类信息，而最小-最大归一化集成可以。