[第 47 名解决方案] Pytorch Lightning 框架 + 按列集成

祝贺所有参与者！在过去的几个月里我学到了很多，感谢 Kaggle 和主办方组织这次比赛。这是一次真正有益的经历，我很兴奋与大家分享我的解决方案。

背景

• 业务背景：https://www.kaggle.com/competitions/leap-atmospheric-physics-ai-climsim/overview
• 数据背景：https://www.kaggle.com/competitions/leap-atmospheric-physics-ai-climsim/data

代码

开源代码可在 Github 仓库 获取。

简要总结 (TLDR)

我们的解决方案利用 PyTorch Lightning 框架进行训练。所有模型架构均基于 Transformer 编码器。旋转位置编码 (RoPE) 被证明是有效的。数据准备涉及归一化、对数变换、最小 - 最大缩放和截断 (clamping)。后处理采用了按列集成 (column-wise ensemble) 方法。虽然 KAN 线性嵌入没有提高分数，但它作为 ensemble 中的模型之一做出了贡献。我们使用 90% 的数据进行训练，剩余 10% 用于验证。

训练 - Pytorch Lightning 框架

我们的解决方案提供了一种方便且整洁的方法，使用 PTLit (PyTorch Lightning 模块) 类来训练机器学习模型。训练和验证循环完全封装在类中，简化了流程。当使用多个模型时，该类会自动平均它们的输出以创建集成。

• 训练特定模型：

mdlit = PTLit(mask, learning_rate, step_size, gamma, [YOUR_MODEL()])

• 将多个模型作为集成进行训练：

mdlit = PTLit(mask, learning_rate, step_size, gamma, [YOUR_1ST_MODEL(), YOUR_2ND_MODEL()])

• 从检查点加载模型并将它们作为集成进行训练：

base_models = PTLit.load_from_checkpoint("PATH_TO_THE_1ST_SAVED_CHECKPOINT").models
base_models += PTLit.load_from_checkpoint("PATH_TO_THE_2ND_SAVED_CHECKPOINT").models
mdlit = PTLit(mask, learning_rate, step_size, gamma, base_models)

模型

旋转位置编码 (Rotary positional encoding)

在我们的实验中，将原始位置编码切换为 旋转位置编码 (RoPE) 后，公共排行榜分数提高了约 0.005。我们最好的单个模型结合了 RoPE 和基本 Transformer 编码器，在公共排行榜上得分为 0.751，在私有排行榜上得分为 0.744。

模型架构 - 其他尝试

我们的 JNet 模型架构处理形状为 (B, 60, 25) 的输入序列，通过一系列 Conv1d 层，每层后跟 GELU 激活函数，以提取层次特征。这些层的输出被连接成一个组合特征表示，即嵌入 (embedding)。该嵌入丰富了旋转位置编码 (RoPE) 以提供位置信息，随后由 Transformer 编码器处理。最后，一个线性层生成形状为 (B, 60, 14) 的输出序列。

我们的 KANformer 模型架构从形状为 (B, 60, 25) 的输入序列开始，并对序列的每个 token 应用一系列 60 个 KAN 线性层 (25, 256)。这些层将输入转换为形状为 (B, 60, 256) 的嵌入。然后由 Transformer 编码器处理该嵌入以捕获序列内的依赖关系。最后，一个线性层生成形状为 (B, 60, 14) 的输出序列。虽然 KAN 线性嵌入没有提高分数，但它作为集成中的模型之一做出了贡献。

数据预处理

在没有领域知识的情况下，我们逐一分析了所有特征列的分布直方图，以确定是否需要进行对数变换来解决偏斜问题。之后，我们手动决定对每个特征列应用归一化还是最小 - 最大缩放。

对于标签列，直接归一化会导致最小值和最大值出现极端值。为了缓解这种情况，我们首先对标签列进行截断 (clamp)，然后应用归一化。虽然截断使模型在初期更容易学习，但在收敛后并没有带来更好的分数。然而，这种方法仍可用于课程学习 (curriculum learning)。我们将基于截断数据训练的模型作为最终提交集成的一部分。

数据后处理 - 按列集成

我们使用 optuna 搜索为每一列集成模型的最佳权重。代码如下：

num_targets = label.size(1)
preds_em = torch.zeros_like(preds[0])
alphas = torch.zeros(num_targets, num_models)

def objective(trial, i):
    weights = []
    remaining_sum = 10
    for j in range(num_models - 1):
        w_i = trial.suggest_float(f'weight_{j}', 0, remaining_sum, step=.5)
        w = w_i / 10.
        weights.append(w)
        remaining_sum -= w_i
    col = sum(weight * preds[m][:, i] for m, weight in enumerate(weights))
    return r2_score(col, labeln[:, i])

for i in tqdm(range(num_targets)):
    if mask[i]:
        study = optuna.create_study(direction='maximize')
        study.optimize(lambda trial: objective(trial, i), n_trials=50)
        best_weights = []
        remaining_sum = 10.0
        for j in range(num_models - 1):
            w = study.best_params[f'weight_{j}']
            best_weights.append(w / 10.)
            remaining_sum -= w
        best_weights.append(remaining_sum / 10.) 
        alphas[i] = torch.tensor(best_weights)
        preds_em[:, i] = sum(weight * preds[m][:, i] for m, weight in enumerate(best_weights))

无效尝试

我们尝试将线性输出层改为卷积输出层和 KAN 线性输出层。然而，这些修改增加了过拟合的可能性。由于时间限制，我们无法实施足够的正则化技术来有效缓解这个问题。

参考文献

[1] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021, April 20). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv.org. https://arxiv.org/abs/2104.09864

[2] Liu, Z., Wang, Y., Vaidya, S., Ruehle, F., Halverson, J., Soljačić, M., Hou, T. Y., & Tegmark, M. (2024, April 30). KAN: Kolmogorov-Arnold Networks. arXiv.org. https://arxiv.org/abs/2404.19756