第5名解决方案 - Stanford Ribonanza RNA Folding

第5名解决方案

作者: sroger (Kaggle MASTER)

团队成员: Ayaan Jang, junseonglee11

比赛: Stanford Ribonanza RNA Folding

发布时间: 2023年12月8日

得票: 17票

我要祝贺我的队友 @ayaanjang 和 @junseonglee11 在这场艰苦的比赛中获得的（master triplet）胜利！我们对比赛各自有着不同的理解，最终将这些理解融合成了集成模型。下面我们将介绍最重要的性能贡献因素；我们集成中的每个基于transformer的模型都采用了以下方法的不同组合。

代码

Roger: https://github.com/s-rog/StanfordRibonanza2023
Ayaan: https://github.com/ehdgnsdl/2023-Stanford-Ribonanza-RNA-Folding
Junseong: https://github.com/JunSeongLee1102/2023-Kaggle-Stanford-Ribonanza-RNA-Folding
Docker镜像: 查看镜像

Ribonanza模型图模型架构图 Junseong模型图

要点概览

我们基于 @iafoss 的基准模型进行了以下改进：

利用低信噪比（SN）样本
利用Eterna BPP（碱基对概率）
将Layer Norm替换为RMS Norm
将绝对位置编码替换为相对位置偏置
将编码器层的pre-norm结构替换为ResiDual-norm结构
从QKV和FFN层中移除偏置项
将头维度从32增加到48
在每个编码器层末尾添加单层双向GRU
多阶段伪标签训练
其他：逆平方根学习率调度，0.1权重衰减

数据采样

我设置了每个epoch的固定序列数量，并使用 WeightedRandomSampler 根据信噪比（SN）加权无放回采样。在SN上添加了作为超参数的偏置项，并按调度设置，在训练过程中逐步增加数据使用量和正则化强度。调度超参数的形式如下：
{(i * 10): v / 10 for i, v in enumerate(range(-12, -7))}

数据增强

我们在一定程度上利用了翻转（flip）增强。虽然使用翻转训练通常会导致更长的训练时间和略微更差的CV分数，但通过翻转测试时增强（TTA）能获得性能提升。
在BPP上添加高斯噪声也改善了CV分数和长距离knot的可见性。

归一化

我尝试了pre-norm、post-norm（结合 Admin 和 ResiDual）。ResiDual-norm获得了最佳的CV分数和训练稳定性。从Layer Norm改为RMS Norm带来了进一步的提升。

注意力偏置：位置编码

在所有12层编码器的注意力机制中使用相对位置偏置（所有12层使用相同的位置偏置）。我们也测试了Alibi和动态位置偏置，但相对位置偏置在我们的架构上表现最佳。旋转位置编码（Rotary embedding）效果也较差。感谢 @lucidrains 的实现代码。

GRU

@junseonglee11 在每个编码器层末尾添加了GRU，当无法进一步扩展注意力层时，这提升了性能（ResiDual的残差输出不与GRU交互）。我在模型中也加入了这一改进，并尝试了多种结构，但似乎越简单越好。我也尝试用1D FusedMBConv块替代GRU，这提升了CV分数但恶化了LB分数。

多头RNN和LSTM

类似于多头注意力结构，可以通过将隐藏维度减半来用多个RNN层替代单个RNN层（增加层数并减少隐藏维度相互抵消）。@junseonglee11 测试了2、3、4、8头结构替代单RNN层。其中，结合GRU层和LSTM层的2头RNN结构表现最佳。

BPP BMM卷积块

@ayaanjang 使用 "1D conv + ResidualBPPAttention (using bmm)" 来整合BPP，代码如下：

class ResidualBPPAttention(nn.Module):
    def __init__(self, d_model:int, kernel_size:int, dropout:float):
        super().__init__()
        self.conv1 = Conv(d_model, d_model, kernel_size=kernel_size, dropout=dropout)
        self.conv2 = Conv(d_model, d_model, kernel_size=kernel_size, dropout=dropout)
        self.relu = nn.ReLU()

    def forward(self, src, attn):
        h = self.conv2(self.conv1(torch.bmm(src, attn)))
        return self.relu(src + h)

class Conv(nn.Module):
    def __init__(self, d_in:int, d_out:int, kernel_size:int, dropout=0.1):
        super().__init__()
        self.conv = nn.Conv1d(d_in, d_out, kernel_size=kernel_size, padding=kernel_size // 2)
        self.bn = nn.BatchNorm1d(d_out)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        return self.dropout(self.relu(self.bn(self.conv(src))))

这个BPP注意力块被添加到FFN和GRU之间，同样不影响ResiDual的残差输出。BPP的引入显著降低了CV分数。

注意力偏置：BPP

我使用了不同的方法来添加BPP。在编码器的前6层中，BPP在卷积后被添加到现有的位置/掩码偏置中（每一层使用相同的BPP偏置）。对于6个注意力头，单通道BPP通过1x1卷积层输出6个通道。减少BPP使用的层数改善了CV分数和长距离knot的可见性（从12层降至6层）。减少BPP头数、添加Contra BPP、使用3x3卷积均无提升效果。

伪标签：2阶段流程

为了生成伪标签，使用翻转训练一个5折集成模型（推理时使用TTA），然后按标准差分位数0.75过滤。过滤后的值将其反应活性设为nan，感谢 @ayaanjang 的这套流程。这个parquet文件随后作为训练数据训练一个预训练模型，再用于在原始比赛训练集上训练5折模型。这个方法来自 @junseonglee11，非常高效。最后阶段在单张4090上每折耗时不到2小时。

伪标签：3阶段流程

@ayaanjang 使用了不同的流程：仅训练集 -> 训练集 + 伪标签集 -> 仅训练集。各阶段学习率分别为：2e-3 -> 2e-4 -> 2e-5。伪标签也使用上述相同方法过滤，不过滤则无性能提升。第三阶段微调用于减少过拟合。

GRU之谜

现在是个谜题... 以下是我的GRU代码：

class GRU(nn.Module):
    def __init__(self, d_model: int, p_dropout: float):
        super().__init__()
        self.gru = nn.GRU(d_model, d_model // 2, batch_first=True, bidirectional=True)
        self.dropout = nn.Dropout(p_dropout)

    def forward(self, x: Tensor) -> Tensor:
        B, L, D = *x.shape[:2], x.size(-1) // 2
        x = x.view(B, L, D, 2).transpose(2, 3).flatten(2)
        x = self.gru(x)[0]
        x = x.view(B, L, 2, D).transpose(2, 3).flatten(2)
        return self.dropout(x)

起初我添加转置和展平操作，以为它们会将GRU的每个方向委托给一半（3个）注意力头，或反过来在头之间混合。进一步检查后，我认为两者都不是...但它却显著提升了性能。如果你能给出满意的解释，请在评论区留言！

致谢

再次感谢 @ayaanjang 和 @junseonglee11 的辛勤工作
感谢 @iafoss 的优秀基准模型，教会了我很多，因为这是我第一个transformer模型
感谢 @lucidrains 的各种transformer实现，我过度参考了它们
感谢 xformers 和 apex 的开发者，大大加快了训练速度
感谢众多论文的作者，他们的研究指导了我们的方法
感谢所有参与我们讨论的Kaggle选手
感谢比赛主办方的一切努力！

5th Place Solution