摘要

• 基于Transformer的解决方案
• 基于分块的数据加载与缓存及长度匹配
• 傅里叶编码器
• 相对时空区间偏置
• von Mises-Fisher损失函数 + 竞赛评估指标

我们的代码已发布于 此代码库。

引言

尽管中微子是宇宙中最丰富的粒子之一，可能携带剧烈天体物理事件的信息，但其基本特性使得探测极为困难，需要巨大的物质体积才能发生相互作用。IceCube中微子观测站是首个此类探测器，由一立方公里的冰构成，旨在搜寻近乎无质量的中微子。中微子与物质罕见相互作用产生的带电粒子会发射切伦科夫辐射，被浸没在冰中的光电倍增管记录。本次挑战的目标是根据光子探测数据重建中微子径迹方向。
首先，我要衷心感谢组织者和Kaggle团队使这场比赛成为可能。阅读竞赛说明时，我立即意识到Transformer可能是非常合适的解决方案。由于我想深入研究Transformer，我决定参赛。同时，我的优秀队友@drhabib参赛则是希望访问南极洲。

详细方案

我们的解决方案基于Transformer，它似乎是处理该任务的理想选择。以下是方法的关键组成部分。

为什么选择Transformer

图神经网络（GNN）在某些方面存在局限：
(1) GNN主要在局部层面运作，基于相邻节点更新。然而这些模型需要预测全局量（如径迹方向）。组织者尝试通过动态建边解决，但这种方法相比Transformer的多头自注意力机制相形见绌。而且动态建边无法为邻居选择提供梯度反馈。
(2) 初看GNN可能比Transformer更快，因其只考虑固定数量的邻居。然而对于合理的序列长度，GNN中大量稀疏操作和数据重排实际上比Transformer高度优化的稠密张量操作慢得多。实验中，组织者基线模型与我们的参考T模型计算成本相同，但参数仅1.4M vs T模型的7.4M。
Transformer是增强版GNN。可视为在全连接图上通过注意力动态估计边权的GNN。但这种能力需要海量训练数据。

数据处理

组织者提供了庞大的训练数据集，管理起来颇具挑战。由于Transformer需要大量数据训练，建立合适的数据管道至关重要。我们的方法包含三个核心组件：(1) 缓存选定数据块以最小化加载和预处理成本，(2) 采用基于数据块的随机采样有效利用缓存，(3) 实施长度匹配的数据采样（采样近似等长的批次）以减少填充标记的计算开销。组件(1)(2)提供了计算高效且内存友好的数据处理方式，已在竞赛期间通过"基于分块的数据加载与缓存"笔记本分享。组件(3)显著加速训练，使我们能在训练时使用最大序列长度192，推理时使用768（推理时间无明显增加）。768长度推理比192长度提升25个基点(bps)。最初我们也尝试使用Hugging Face数据集，虽支持分块数据但需额外文件转换，最终选择了数据块缓存管道。
实验中，我们发现仅使用基础特征（时间、传感器位置、电荷、辅助标志和总探测数）与加入冰透明度和吸收等额外特征之间无显著差异。这合乎情理，因为对于足够大的训练集，额外特征仅提供了z坐标的替代表示，相关透明度和吸收分布会被显式学习。为确保模型多样性，我们在实验中同时使用了基础和扩展特征集。特征归一化方式与组织者基线类似。必须同时考虑辅助标志为真和假的探测数据。对于超过最大序列长度的事件，采用以下选择流程：首先从辅助标志为假的子集中随机选择探测数据，若该子集不足则随机采样含噪声的辅助标志为真的探测。有趣的是，这种随机选择方法优于更复杂的方法（如基于电荷选择或只考虑最大电荷探测附近特定时间边界内的数据）。

模型架构

模型结构如下图所示：

傅里叶编码器。我们的方案基于Transformer模型，将每个事件视为探测序列。关键在于将连续输入（如时间或电荷）处理成适合Transformer的形式。我们采用傅里叶编码表示，该方法常用于语言模型中描述序列位置。这可视为将连续输入信号软数字化为傅里叶频率决定的编码集。该过程应用于所有连续输入变量，离散辅助标志则使用可学习嵌入。我们将归一化输入变量乘以1024-4096以获得足够分辨率。例如，归一化时间乘以4096后时间分辨率达7.3纳秒，模型因傅里叶编码的连续性可理解更小时间变化。该乘法至关重要，系数从128增至4096可带来20 bps提升。
所有变量编码被连接后通过一个GELU层，再投影到Transformer维度。傅里叶编码的替代方案可能是学习连续变量表示的多层网络，但需要更多层数。
GraphNet编码器。除傅里叶编码外，部分模型中还加入了GraphNet特征提取器，不仅学习输入变量表示，还构建考虑邻居关系的特征。
相对时空区间偏置。狭义相对论中存在时空区间量：ds²=c²dt²-dx²-dy²-dz²。对于近光速运动的粒子，该值接近零。更有价值的是，由单个中微子反应产生的所有粒子和光子也应具有接近零的ds²（因冰折射率约1.3，光速低于c，故不完全精确）。因此，通过计算给定事件内所有探测对之间的ds²，相对容易区分源自同一中微子的探测和噪声，如下图所示。该准则可自然地作为相对偏置引入Transformer。因此在构建注意力矩阵时，Transformer会自动按源事件分组探测，有效过滤噪声。我们使用ds²的傅里叶编码表示构建偏置项。使用相对偏置带来40 bps的性能提升。

模型。我们使用标准BEiT模块（带可学习捷径的Transformer）。为融入相对偏置，前几个Transformer块根据该论文修改。模型大小采用ViT标准命名：T - 极小（维度192），S - 小（维度384），B - 基础（维度768）。我们共使用16个Transformer块：4个带相对偏置的块 + 12个带cls token的块。输入序列中加入额外的cls token，可在整个模型中逐步收集径迹方向信息。注意力头大小在32-64间变化：较小头尺寸性能更好，尤其对T模型（10 bps）。各模型详情总结如下表。

训练配置。模型使用AdamW优化器（权重衰减0.05）训练4-5个epoch（使用全部数据）。学习率采用带预热的余弦退火调度。首个epoch最大学习率：T和S模型为5e-4，B模型为1e-4。后续epoch最大学习率为0.5e-5 - 2e-5。有效批次大小均为4096（使用梯度累积）。从第二epoch开始使用随机权重平均。
前2-3个epoch使用von Mises-Fisher损失（模型预测3D向量，kappa基于输出向量长度计算），剩余epoch使用竞赛指标作为目标函数 + 0.05 von Mises-Fisher损失（为模型提供向量长度与置信度关系的反馈，简化集成步骤）。使用竞赛指标作为损失带来55 bps提升。

结果

下表总结了最终模型性能。CV在L=512下评估。(*)表示在2×A6000上训练的模型（相比rtx4090支持更大批次）。L=256的进一步微调可能带来约5 bps提升。

模型尺寸从T→S→B每增加一级约提升30 bps，但参数量增加4倍。然而，提供的数据可能不足以训练超过B尺寸的模型。同时，GNN与傅里叶提取器结合使用的效果相互矛盾，可能受硬件配置差异影响。
因此，即使最小的T模型（仅7.53M参数，rtx4090上每epoch训练时间不足12小时）也能单独进入竞赛前5名。同时，通过L=256的额外微调，我们最佳模型B+gnn d32作为单模型达到0.9628 CV512、0.9608公开榜和0.9618私有榜。

模型集成

使用von Mises-Fisher损失时，向量长度对应模型置信度。因此，预测向量的加权平均自动融入置信度，使预测偏向最自信的方向。
我们使用6个模型的线性集成获得最佳提交：公开榜0.9594，私有榜0.9602（各模型贡献通过在5个随机5折分割中拟合权重的平均值评估）。集成相比最佳单模型提交（公开榜0.9608，私有榜0.9618）仅有微小提升。我们也尝试使用更多特征（如序列长度、首次探测时间、总电荷等）进行更复杂集成，但未超越简单线性集成。

训练与推理时间

在rtx4090上，T模型推理时间可短至每100万样本3分钟（无推理优化），B模型可能增至8-10分钟。在落后几代的GPU（如Kaggle的P100）上，推理时间可能增加10倍。训练时间从T模型的每epoch 12小时到B+gnn模型的56小时。在2×A6000机器上速度比单rtx4090快约10%。

鲁棒性

关于输入参数扰动下模型预测的鲁棒性，我们考虑了以下情况：(1) 时间噪声：为每个事件探测时间添加正态分布随机变量；(2) 电荷噪声：将电荷乘以10^正态分布随机数（方差见横轴）；(3) 辅助标志噪声：以给定概率随机翻转辅助标志；(4) 长度缩减：移除给定百分比的探测。下图显示了参考T模型的结果（无噪声CV512为0.9704）。

尽管模型未显式训练处理上述噪声类型，大多数情况下预测在最大噪声幅度下退化不超过15%（例如添加200 ns随机时间误差或将探测数减半）。训练时使用特定数据增强可能提升鲁棒性。然而可能存在误差急剧增加的阈值（如10 ns时间噪声或探测电荷变化40%）。
该模型足够通用，可应用于任何与IceCube探测概念相似的任意中微子望远镜，无论探测器数量与布局。然而，由于训练中探测器位置固定，若探测器移动模型可能无法完全理解几何结构而性能下降。例如，T模型在沿X或Y轴翻转所有探测器位置时性能从0.9704降至1.004。为避免性能下降并提升模型对几何结构的理解，可在每次探测事件时随机生成探测器配置进行训练。加入几何增强（如平移、翻转、旋转）也可提升模型对探测器几何变化的鲁棒性。次优方案是为新望远镜生成新训练数据。考虑3D旋转平移等变注意力模型，通过注意力而非直接输入融入探测器位置，可能解决几何问题。但该改进可能增加计算成本和显存需求，更实际的方案可能是使用空间增强和随机化探测器位置训练简单模型。
下图展示了自信与不自信预测的误差分布。

有趣发现

似乎大多数参赛者无意识地利用了以下数据泄露。提供的数据是模拟结果，0时刻有特定含义，揭示了生成过程的某些细节（首次探测相对于0时刻的时间可能提供中微子能量和基于从边界到探测器传播时间的近似方向限制等信息）。模型学会了利用这一泄露提升性能。例如，我们使用T模型的实验显示：(1) 保留0时刻参考的数据训练，(2) 从所有探测时间中减去首次探测时间，首个epoch结束时的性能差距达145 bps。
在实际探测器数据中，0时刻参考不可用，时间可能仅相对于首次探测事件。对使用0时刻参考训练的模型，实验数据应平移平均首次探测时间。但这样性能可能下降150+ bps，与无0时刻泄露训练的模型结果相当。

未奏效的方法

本次竞赛中，我们尝试了多种未获改进的方法：
图神经网络。构建首个Transformer管道后，GNN明显落后（首次模型0.98-0.99 vs 1.0榜）。我们进行了一个有趣的比较：按kappa将预测分为30%自信和70%不自信（噪声输入和极少数探测）。两种模型在不自信预测上性能相当，自信预测的改进是Transformer与GNN性能差距的关键。这完全符合我们最初预期：Transformer全局考虑输入数据，更适合预测径迹方向这类全局量。因此后续工作集中在增强Transformer方案。最终我们仅将GNN作为傅里叶特征提取器的可选补充以保证多样性。
冰性质。未发现使用冰性质等额外输入能提升模型性能。
局部注意力。我们花费数周实验局部注意力。基本思想简单：在将探测序列输入考虑全连接交互的Transformer前，先在局部邻居内处理特征。理论上，局部注意力应快速且内存高效，只需检查k=8个邻居而非L×L注意力。但实践中，局部注意力因大量数据重排（类似GNN）而相当慢，某些实现甚至比优化Transformer消耗更多显存。最快的局部注意力竟是使用掩码屏蔽k个邻居外注意力的常规Transformer注意力。这种掩码在训练初期有帮助，但后期全注意力更优或相当。我们也尝试了全局-局部注意力交替（夹层结构）或并行块配置，但简单的4层全注意力加相对偏置+12层常规Transformer配置超越了所有局部注意力方案。
分类损失。我们尝试将角度空间细分并使用分类损失对抗角度不确定性。真实标签表示为提供标签周围的2D高斯，如下图所示。我们引入8个cls token预测64×128角度bin图（每token 32×32）。不幸的是，该模型最终CV约1.0，未进一步验证。但该方法值得提及，可能提供良好的预测可视化。

L和H模型（4亿和8亿参数）。由于Nvidia虚拟地址空间驱动bug，2×rtx4090直到竞赛结束前1-2周才能协同工作。即使修复后，禁用的p2p使这些卡的多GPU训练效率低下。因此我们无法扩展实验，停留在相对较小的B模型。每次尺寸提升预期约30 bps，L和H单模型可能分别达到0.958和0.955-0.956榜，但需巨大的训练和推理算力。