第五名 - GNN + 自我中心 Squeezeformer

副标题：MABe 挑战赛的神经网络解决方案

这是一场有趣的比赛，非常感谢组织者，也祝贺所有参与者。从一开始，我的想法就是构建一个统一的神经网络模型，能够支持不同的实验室和身体部位。我的重点是鲁棒性和多样性，希望它能实现跨实验室和未见视频的迁移学习。

概述/TL;DR

1. 架构: GNN (TransformerConv) + 时间 Transformer (Squeezeformer) + 成对/关系网络分类头
2. 关键洞察: 社交方面和自我中心处理对我的解决方案非常重要。高 Dropout 和模型多样性对于不同实验室/未见视频的鲁棒性和迁移学习至关重要。
3. 特征:
   • 标准特征：每只老鼠的速度、加速度、加加速度 (jerks)、角速度、运动朝向……
   • 老鼠特定特征：身体曲率、身体长度、身体 elongation 比率、耳朵张开度、身体长度变化、头部角度、尾巴角度……
   • 从每只老鼠视角的鼠间特征：质心距离、距离变化、各种身体部位距离（例如鼻子对尾巴用于嗅探等）、相对速度、接近角度等。
   • 每只老鼠 58 个特征，导致 232 维输入。（不同的训练运行/模型系列之间略有差异）
4. 训练: 偏向动作丰富窗口的采样，带类别权重的 Focal Loss，AdamW + OneCycle，FP16 训练 + 推理

数据管道与特征工程

所有的跟踪和标注 parquet 文件都经过预处理并保存为 .npy 格式以加快训练速度。
这次比赛的一个问题是不同的实验室跟踪不同的身体部位。为了解决这个问题，我采用了一种主骨架方法，即选取数据集中最具代表性的一部分身体部位。这产生了 6 个主骨架身体部位：

[nose (鼻子), ear_left (左耳), ear_right (右耳), head_center (头部中心), body_center (身体中心), tail_base (尾巴根部)]。

当某些关键点缺失时，有一个回退系统。在头部场景中，我们首先寻找头部，然后是颈部，然后是左右耳的平均值。所有关键点都通过 pix_per_cm 比例进行归一化。

特征工程包括提取每只老鼠和鼠间特征，以帮助模型处理自我动作和社交动作。模型被提供了标准特征，如速度、加速度、加加速度等，以及老鼠特定特征，如身体配置特征、朝向、肢体特征等。除了这些特征外，一个重要部分是鼠间不变特征。在模型中未使用原始坐标。所有特征都是不变的，因为它们是相对于自身身体或两只老鼠之间的。我的直觉是，攻击无论是在笼子中间还是侧面都是一样的——但老鼠之间的相对角度/距离才是将攻击描述为攻击的关键。模型的最终特征输入向量为 [Batch(32), Time(512), Mice(4), Features(58)]。在某些模型上使用了非常轻的增强，如增强特征比例、添加噪声和关键点 Dropout。

架构

模型的架构是 GNN + Squeezeformer + 分类头。对于 GNN，我使用了 TransformerConv 来实现老鼠之间的社交互动。在此之前，我尝试过将老鼠数量作为通道维度用卷积来建模老鼠之间的互动，但这只带来了很小的改进，并且 CV/LB 差距较大，所以我放弃了这种方法，因为 GNN 更有效且更好地建模了笼子内的互动。在将特征通过 GNN 后，模型现在具有社交感知能力，每只老鼠的特征都被笼子中其他老鼠的特征所丰富。然后，这个特征向量被传入 Squeezeformer，需要注意的是，我们没有折叠 [B, T, M * F] 维度，而是将 B * M 维度折叠为 [B * M, T, F]。因此，我们在时间维度上“单独”处理每只老鼠的轨迹。为了适应我的 GPU，我不得不将 Batch size 减少 4 倍，因为现在 Batch 维度中的数据点多了 4 倍。在将数据通过 4 个 squeezeformer 块后，我们将其重塑回 [B, T, M, encoder_dim]，然后通过分类器。分类器中的数据通过广播和最后维度的拼接变为成对的：[B, T, M, M, 2D]，然后通过 MLP 给出每个成对互动的最终 logits。分类头的 logits 进一步通过 behavior_mask 进行掩码，该掩码仅允许当前视频中标记的行为对损失做出贡献。最终形状为 [B, T, M*M, Actions]，其中每一帧被分类为可能的动作 + 无动作之一。

模型多样性

• 除了主要架构外，还有其他调整以获得最终集成所需的多样性。一些模型没有先在开始时将数据通过 GNN 然后 进入 Squeezeformer，而是交错层 GNN->Squeezeformer->GNN->Squeezeformer。这并没有比已经有效的方法提供实质性的收益，但它为最终集成提供了多样性。其他模型在 GNN 中省略了 edge_features 并即时学习它们，还有一些模型在图中被提供了明确的鼠间特征作为边属性。
• 增强的 RelationNetwork 分类头。一个模型系列使用了受关系网络启发的分类头，而不是成对分类头。我从过去读过的一篇论文中获得了想法/灵感，作者在其中使用一个简单的插件网络进行视觉模型中的关系推理。与成对分类头相比，增加的部分是互动向量的总和，这给出了一种“笼子氛围”。在成对分类头中，我们只看老鼠 A 对老鼠 B，而不知道笼子整体发生了什么。然而在 RelationNetwork 头中，我们也看笼子的全局上下文——如果我们看到有老鼠在睡觉或处于低能量状态，那么老鼠 A 和老鼠 B 之间发生攻击的可能性就较小。同样，由于模型已经相当不错，这仅提供了微妙的改进，但更重要的是它增加了多样性。

训练策略

模型被输入来自实验室视频的窗口块。每个窗口长 512 帧。在训练期间，动作丰富帧的采样很重要，因为数据集不平衡且标记帧 << 总帧数。我们检测所有有动作的帧并将索引存储为“活跃帧”。在采样期间，我们偏置数据加载器从这些活跃帧中采样一个窗口，并进一步随机偏移窗口的开始，以便我们可以在动作的开始/停止中引入变化。采样动作与非动作窗口的概率通过偏置超参数进行调整，范围从 0.5 到 0.8。然而，这还不足以对抗不平衡。即使我们以 0.5 的概率采样动作丰富窗口，仍然存在另一个不平衡，即每个窗口中动作与非动作帧的持续时间。大多数动作的中位持续时间低于 100 帧，而窗口是 512 帧，这留下了超过~400 个非动作帧，这意味着即使我们在 50% 的时间内偏置数据加载器采样动作丰富窗口，平均而言我们采样的非动作帧仍然比动作帧多得多。为了对抗这一点，我使用了带每个动作类别权重的 Focal Loss，这些权重是通过每个动作的长度与采样窗口长度的比率计算的。当一个动作在 512 帧窗口中仅跨越 12 帧时，模型对将这 12 个动作帧分类错误的惩罚与对 500 个非动作帧分类错误的惩罚相同。

训练使用 3e-4 和 1e-3 的学习率，AdamW 优化器 + OneCycle 调度器，0.15% 预热。Batch size 为 32。Focal Loss gamma 为 2.0。训练和推理期间使用 FP16。不同运行/模型系列的 Dropout 范围从 0.1 到 0.35。

后处理与验证

模型预测后几乎没有后处理。在获得每一帧的概率后，我们只需找到连续片段并将其分类为动作片段。有一个最小持续时间过滤器可去除闪烁，还有可选的片段合并逻辑，可找到相距 N 帧的相同动作片段并将其合并（这在最终提交中未使用，因为我没有时间调整它，也不想冒过拟合的风险）。验证是通过在重叠窗口中扫描整个视频完成的。步长是窗口大小的一半（窗口 = 512 -> 我们移动 256 帧并运行模型），最后在帧重叠的地方平均最终概率。大部分收益来自于将不同的模型系列集成到最终解决方案中。一个模型在一个 T4 GPU 上大约需要 15 分钟完成预测。最终提交中使用了 23 个模型。在这 23 个模型中，大多数只是同一模型系列的 4 折 CV 平均。

有效的方法

1. 用 GNN 处理老鼠并在 Transformer 中以自我中心轨迹单独查看每只老鼠是最重要的改进之一。此前我尝试过在老鼠维度上使用 CNN，以及在将老鼠特征传入 squeezeformer 之前混合所有老鼠特征，尽管模型学习并得分不错（没有 GNN 且将混合老鼠特征为单个特征向量的 6 个单折模型集成大约在 16 名左右），但这并不是最优的。
2. 高 Dropout 和多样性。一些模型使用 0.35 的 Dropout 进行训练。模型仍然收敛得很好并保持鲁棒性，我认为这对于跨实验室和未见视频的迁移学习非常重要。另一个主要点是多样性。一个模型系列先在开始时训练 GNN 然后是 Squeezeformer。另一个交错 GNN 和 Squeezeformer 块，以便它们相互更新。下一个系列没有明确给出 edge_features，必须自己学习，另一个版本在 TransformerConv 层中将计算出的鼠间特征作为明确的 edge_features 提供。一个模型系列具有所有特征的自我中心视图，其中我们计算旋转矩阵并旋转世界/其他老鼠，以便老鼠总是从其 POV 查看世界。这种多样性和高 Dropout 带来的鲁棒性实现了第 5 名的私有榜成绩。

无效的方法

• 更深或更宽的模型。这对分数影响很小。我尝试过不同的深度，最多 12-16 层，宽度最多 512 维，但没有提供实质性的改进。我认为数据的“路径”和架构比规模更重要。
• 反向时间集成。众所周知的想法是向前 + 向后处理视频并平均概率。但由于某种原因我无法使其工作。可能是我实现中的一些 bug，反向视频总是表现不佳并拖累了集成。
• 使用普通 1D CNN。CV/LB 差距比 squeezeformer 大得多，所以我放弃了这种方法。
• 添加滚动窗口/傅里叶特征。想法是检测周期性动作，如梳理毛发。然而，在我的第一个实现中，计算傅里叶和滚动窗口特征的开销太大，所以我从特征生成中放弃了它们。
• 在缺失帧中插补关键点 - 这降低了本地 CV
• 更大/更小的窗口大小。我尝试过 128、256 和 768 帧长的窗口，但差别不大。

记录但未尝试的想法

由于我加入得有点晚，我写下了一些我想尝试但没有时间的想法。其中一些是：

1. 动作类别 + 开始/停止边界回归。使用类似于 ActionFormer 论文的方法，除了预测当前帧的动作类别外，我们还尝试预测当前帧到真实片段开始/停止的距离。据我与 LLM 交谈时的理解，我认为这比仅仅预测当前帧要好得多。因为在 ActionFormer 方法中（据我理解），动作片段内的许多帧可以“投票”决定它们认为相对于它们的开始/停止帧在哪里。因此，即使动作片段开始时的一些帧损坏，中间的帧也可以投票决定它认为动作开始的地方，我们可以利用这一点从所有帧中获得概率质量，投票决定它们认为动作开始的地方。我仍然不完全理解它，需要花更多时间，但似乎是一个合理的想法。
2. 不同的时间步长/帧率（每第 2 帧等）
3. 在 90% 未标记数据集上预训练的关键点插补器，用于插补缺失的身体部位。用它从 6 点升级到更详细的 14 点主骨架。
4. 由于实验室不变，我想制作针对特定实验室微调的每实验室适配器。

感谢这次机会和学习经历。