QuData团队第11名解决方案：注意力机制、图神经网络与集成学习

首先，我们非常高兴能够参与这场意义深远且极具趣味性的竞赛，衷心感谢主办方提供的支持，以及所有无私分享知识与见解的参赛者们：@rasmusrse、@rsmits、@solverworld、@edguy99、@iafoss、@pellerphys……如果遗漏了任何人，请原谅我们，帮助过我们的人实在太多了！

非常感谢我们优秀的QuData团队 @synset、@alexz0 和 @semenb，你们太棒了！我们都来自乌克兰，在此衷心感谢所有正在支持我们国家的人们！

本次竞赛为我们提供了将相关领域知识（团队中有两人曾从事理论物理研究）与数据科学技能相结合的绝佳机会。参与这项大型天体物理实验既新颖又令人兴奋，我们希望为探索宇宙的运行规律贡献了自己的一份力量。

源代码与详细技术报告

由于我们的报告内容较为丰富，难以在一篇帖子中完整呈现，因此我们提供了以下链接：

整体架构

我们的解决方案结构如下：

"GNN" 是对Neutrino Telescope Event Reconstruction图神经网络架构的改进（基于 GraphNet），公开分数为 0.992。

"Transformer" 结合了全连接模块、注意力机制和循环层的混合架构，公开分数为 0.995。

"Ensemble" 是一个神经网络，用于聚合"GNN"和"Transformer"模型的输出并预测方向，公开分数为 0.976。

GNN模型详解

模型架构

在比赛期间，我们对图神经网络架构进行了多次改进（详细描述见 GNN架构页面），最终形态如下：

该模型的输入是一个图，其中每个节点对应一个脉冲，特征由脉冲特征与事件聚合特征组成。为构建图的拓扑结构，我们采用 knn 算法。图依次通过多个层，每层都会修改特征空间并更新拓扑结构。所有层的输出被合并后送入 Layers MLP 层，将特征数降至256。随后执行池化操作，使用 min、max、mean 函数聚合图中所有节点的特征。最后在 Classification MLP 节点，将得到的嵌入向量转换为类别ID，每个类别对应一个特定的方向。

模型训练与增强过程

训练基础模型 (1.018 → 1.012)

由于计算资源有限，我们决定基于最佳公开GNN模型进行再训练。从 GraphNeT Baseline Submission（特别感谢 @rasmusrse）获得，其LB分数为 1.018。

我们使用除最后一组外的所有批次进行训练，最后一组用于验证。训练过程中采用阶梯式学习率衰减策略，当验证集指标约100个epoch未改善时降低学习率。经过956个epoch后，指标降至 1.0127。

增加网络层 (1.012 → 1.007)

在已有网络基础上增加一个 EdgeConv 层。为避免从头训练，新架构中除新增层外全部冻结。

冻结层加载前一阶段的权重并继续训练。该模式下模型学习迅速，指标很快达到源模型水平。随后解冻网络，降低学习率，对整个模型进行联合训练。经过1077个epoch后，我们达到 1.007 的指标。

增加邻居数量 (1.007 → 1.003)

在原始Graphnet库中，构建图时邻居数量设为 8。我们尝试将其增加到 16。在 EdgeConv 模块中，相同的MLP模型应用于所有邻居，然后求和。因此邻居数增加不会改变网络参数数量，但会使计算量翻倍。

基于前一阶段的模型，我们采用新的邻居数量重新训练，经过1444个epoch后，指标降至 1.003。

扩展Layers MLP (1.003 → 0.9964)

由于网络层数增加，我们认为接收各层输出的 Layers MLP 应相应增加参数量。我们将该模块第一层从336扩大至2048。同样采用冻结策略，加载前一阶段权重并继续训练。经过1150个epoch后，指标降至 0.9964。

回归转分类 (0.9964 → 0.9919)

在研究最佳方案时，我们注意到 Tensorflow LSTM Model Training TPU（感谢 @rsmits），借鉴了其将回归问题转化为分类问题的思路。

将方位角均匀划分为24个区间，天顶角同样分为24个区间（使用天顶角的余弦值，因其分布更均匀），共得到 24×24=576 个类别。将MLP最后一层从[128,3]扩展为[512,576]，损失函数改为 CrossEntropyLoss。

冻结除最后一层外的所有模块，加载前一阶段权重并继续训练。经过967个epoch后，指标达到 0.9919。这是我们独立GNN模型的最佳结果，随后用于集成学习。

效果不明显的尝试

在竞赛过程中，我们尝试了许多其他方法，但均未带来显著改进：

• 分别预测天顶角和方位角
• 将卷积类型改为SAGE、GAT
• 在Layers MLP后插入Transformer
• 用RNN(GRU)替代池化层
• 插入TopKPooling层
• 使用4个或更多特征构建事件图

未充分探索的方向

由于时间和硬件资源限制，部分思路未能充分实现：

• 从头训练包含散射、吸收和DOM嵌入特征的GNN
• 使用更细粒度的分类区间训练分类器
• Transformer与GNN模型的并行训练

Transformer模型

我们命名的 "Transformer" 模型是包含全连接模块、注意力机制和循环层的混合架构，典型指标为0.995。

模型架构

Transformer模型接收两个张量作为输入：(B,T,PF) 和 (B,EF)，其中 B 是批次中的样本索引，T 是脉冲序列中的位置索引。每个脉冲由 PF=9 个特征描述。第二个张量 (B,EF) 描述整个事件的属性，包含24个所有脉冲的聚合特征和另外24个aux=0脉冲的特征。

在某些架构变体中，aux=0的聚合特征通过单层MLP压缩到更低维度，然后与脉冲特征张量 (B,T,PF) 结合。压缩的动机是事件特征数(24)显著多于脉冲特征数(9)。

拼接后的张量 (B,T,F) 输入第二个MLP块（特征生成器），将其特征数扩展至 E=128。随后进入10-12个Transformer块的序列，通过注意力机制，每个脉冲与序列中所有脉冲"交互"。

Integrator块将Transformer输出张量 (B,T,E) 转换为 (B,E') 维张量。训练初期采用平均池化和最大池化拼接作为Integrator；最终阶段替换为RNN层来整合事件脉冲，使用单向GRU单元，脉冲按时间倒序排列，使最后一个GRU单元接收最重要的首个脉冲。

所有脉冲和aux=0脉冲的事件特征张量被附加到Integrator输出的 (B,E') 张量，形成 (B,E'+2*EF) 张量，输入MLP输出预测方向向量的三个笛卡尔分量。

Transformer采用标准架构，注意力块包含16个头，跳连权重可训练。未使用位置编码（早期实验过），脉冲的时间特征似乎已足够。

数据集处理

Transformer训练要求批次内所有样本具有相同数量的脉冲。标准解决方案是使用掩码token对齐序列长度，但计算资源有限时这种方法非常浪费。

我们的替代方案：每5个训练批次打包成100万个样本，按序列长度排序并分组，组内划分小批次，最后混合所有组的小批次。一个训练epoch包含一个数据包。

另一个节省内存的思路是动态批次大小：内存需求随序列长度T二次方增长。因此长序列使用较小批次，短序列使用较大批次：

batch_size = min(int(batch_size * (T_max/T)**2), batch_max)

这使得长短序列的批次内存消耗大致相同。

训练实现

以DOM聚合脉冲的典型训练开始为例：从零开始学习，在6000-7000万个样本（300个批次）后达到约1.020的误差。在标准T4显卡上约需12小时，使用Adam优化器，恒定学习率lr=1e-3。

达到约1.002的误差需要更长时间，学习率逐步降至1e-5。在1.004水平时切换到更长的序列长度，立即使误差降低1-1.5个点。进一步训练需要切换到非DOM聚合脉冲，在A100显卡上耗时数天。

训练过程中，我们分析了模型各块权重值和梯度传播：

集成方案

我们训练的 Transformer 和 GNN 架构结果相关性不强，使得集成学习能够显著提升最终指标。此外，主办方提供的大规模数据集使我们能够使用训练好的集成进一步提升性能。

集成模型选择

每个架构选择多个不同超参数或架构选项的模型，使用80个批次通过模型生成集成数据集：

id   ang_err  az_err  ze_err   model
 0   0.9935   1.035   0.606  | gnn_1
 1   0.9922   1.033   0.595  | gnn_2
 2   0.9963   1.039   0.598  | gnn_3
 3   0.9954   1.018   0.612  | att_1
 4   0.9993   1.022   0.636  | att_2
---------------------------------------------------------------
     0.9846   1.021   0.540  | simple mean
    

模型选择标准是角度误差和模型间角度误差的相关系数：

     0        1        2        3        4
 0            0.959    0.950    0.807    0.781
 1   0.959             0.934    0.820    0.795
 2   0.950    0.940             0.808    0.783
 3   0.807    0.820    0.808             0.940
 4   0.782    0.795    0.783    0.940    
    

简单平均各模型预测向量得到 0.9846 的误差。

进一步，我们构建了带可训练权重的集成：

    n = w₁ n₁ + ... + w₅ n₅
    

这将误差降至 0.9827，各模型权重为：

gnn_1: 0.934,  gnn_2: 1.346,  gnn_3: 0.753, att_1: 1.466, att_2: 0.477
    

训练集成网络

通过神经网络训练集成进一步提升指标。考虑过多种架构（包括Transformer和添加事件聚合特征），但最终简单的单隐藏层MLP效果最佳。

输入张量维度为 (B,N,3)（B为批次大小，N为模型数量，每个模型输出三维方向向量），展平为 (B,3*N)，经MLP后输出 (B,3)。

隐藏层神经元数从128到2048，但效果基本相同。最佳损失函数是预测向量与目标向量的余弦相似度，学习率较高（lr = 1e-3 - 1e-4），使用Adam优化器。

在首批5个批次上的验证最佳指标为 0.9796。

源代码

我们的解决方案源代码位于 https://github.com/QuDataAI/IceCube 和Kaggle公开笔记本：

• IceCube-Dataset-kaggle - 创建Transformer数据集
• IceCube-Transformer-kaggle - Transformer训练
• IceCube-GNN-train-colab - GNN训练
• IceCube-Ensemble-kaggle - 集成训练
• IceCube-Ensemble-submit - 集成提交
• IceCube_QuData_Line-fit_Method - Line-fit方法改进
• IceCube_QuData_Anisotropy - 各向异性分析

未完成的想法

因时间和资源限制未能充分探索的方向：

• 从头训练带散射、吸收和DOM嵌入特征的GNN
• 使用更细分类区间的GNN分类器
• Transformer与GNN的并行训练
• 在Transformer中使用双向循环层作为Integrator
• 使用预训练Transformer回归模型检测噪声脉冲并过滤
• 在Transformer中将回归器替换为大规模方向分类器（球面上的点）
• 针对特定事件模式（如单串、两串事件）训练专用模型
• 在Transformer中添加位置编码

总结

感谢主办方提供的大规模数据集，使我们能够训练复杂架构并将其组合成集成。没有这个数据集，我们无法取得如此成果。