使用LSTM+注意力机制的模型无法学习指定类别，始终输出同一类，如何优化？

def __init__(self):
    super().__init__()
    self.lstm = nn.LSTM(input_dim,
                        hidden_dim, 
                       num_layers=num_layers, 
                       bidirectional=bidirectional, 
                       dropout=dropout,
                       batch_first=True)
    self.fc = nn.Linear(hidden_dim * 2, num_classes)    
def attention_net(self, lstm_output, final_state):
    hidden = final_state.unsqueeze(2)
    attn_weights = torch.bmm(lstm_output, hidden).squeeze(2)
    soft_attn_weights = F.softmax(attn_weights, 1)
    context = torch.bmm(lstm_output.transpose(1, 2), 
     soft_attn_weights.unsqueeze(2)).squeeze(2)
    return context, soft_attn_weights.cpu().data.numpy()
def forward(self, text):       
    output, (hn, cn) = self.lstm(text)    
    hn = torch.cat((hn[-2,:,:], hn[-1,:,:]), dim = 1)
    attn_output, attention = self.attention_net(output, hn)
    return self.fc(attn_output), attention

数据层面优化 #

• 优先检查类别分布：如果第3类样本量远低于其他类别，属于典型的类别不平衡问题。可通过对第3类过采样、对多数类欠采样调整样本分布，也可直接使用加权交叉熵损失，给第3类分配更高的损失权重，强制模型关注该类特征。
• 校验标签正确性：排查第3类的标注是否存在错误，比如大量样本被误标为其他类别，导致模型没有正确的监督信号学习该类特征。

损失与训练策略优化 #

• 更换损失函数：将普通交叉熵损失替换为Focal Loss，通过降低易分类的多数类样本的损失权重，提升少数难分类的第3类的损失贡献，避免模型被多数类主导。
• 调整训练超参数：如果学习率过高会导致模型无法收敛到细粒度特征，学习率过低会导致收敛慢欠拟合，可尝试降低初始学习率，搭配余弦退火、阶梯式衰减等学习率调整策略；如果dropout参数设置过高，可适当降低dropout比例，减少模型的拟合限制。
• 增加正则约束：加入L2正则或者标签平滑策略，避免模型过拟合到多数类的特征，提升对不同类别特征的区分度。

模型结构修正与优化 #

• 修复潜在梯度问题：attention_net方法中对soft_attn_weights执行的.cpu().data.numpy()操作不会影响context的梯度传递，但如果你训练过程中需要用到返回的注意力权重计算损失，该操作会打断梯度回流，训练阶段可以直接返回张量格式的注意力权重，推理阶段再做格式转换。
• 优化注意力计算逻辑：当前注意力的Query直接采用LSTM最后时刻的隐态，泛化性有限，可以新增一层可学习的参数矩阵生成Query，提升注意力对不同类别关键特征的捕捉能力。
• 补充归一化层：在LSTM输出和全连接层之间加入层归一化（LayerNorm），稳定训练过程的数值分布，提升模型对不同类别特征的区分能力。
• 校验全连接层配置：确认nn.Linear的输出维度num_classes和实际类别数完全匹配，避免参数配置错误导致第3类没有对应的输出通道。

推理层面调整 #

如果训练阶段第3类已经有一定的召回率，仅推理阶段被压制，可以不使用默认的argmax做类别判定，针对第3类单独设置更低的分类阈值，提升该类的召回率。

内容的提问来源于stack exchange，提问作者Nampakka