RoBERTa融合BiGRU模型训练失效问题排查求助

尝试微调RoBERTa并通过BiGRU模块融入外部知识，但模型无法学习——训练Loss维持在0.8左右且无下降。数据集无问题，其他基于RoBERTa的模型在该数据集上表现正常。

模型架构 #

class CustomRoberta(PreTrainedModel):

    def __init__(self, config, num_labels, max_em_len_1, max_em_len_2, no_bert_layers=2):
        super(CustomRoberta, self).__init__(config)
        self.num_labels = num_labels
        self.bert = RobertaModel.from_pretrained("roberta-large")
        self.hidden_size = self.config.hidden_size
        self.emotion_embeddings = nn.Embedding(max_em_len_1, self.config.hidden_size)
        self.opinion_embeddings = nn.Embedding(max_em_len_2, self.config.hidden_size)
        self.dropout = nn.Dropout(0.1)
        self.dense = nn.Linear(self.config.hidden_size * 3 * 3 * 2, num_labels) # avd_pool max_pool, last hidden state
        self.apply(self.init_bert_weights)
        self.linear = nn.Linear(self.config.hidden_size, num_labels)

        self.gru = nn.GRU(self.config.hidden_size * 3, self.config.hidden_size * 3, bidirectional=True, batch_first=True)


    def init_bert_weights(self, module):
        """ Initialize the weights."""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.xavier_uniform_(module.weight)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()

    def get_att(self, hiddes, emotion_embd):
        concat = torch.cat([hiddes, emotion_embd], -1)
        g = self.att_lin(concat)
        alpha = F.softmax(g, dim=0)
        att_hidden = alpha * hiddes
        return att_hidden

    def forward(self, input_ids, opinion_ids, attention_mask, emotion_ids, return_indices=False):
        bert_encoded_layers_raw = self.bert(input_ids, attention_mask).last_hidden_state 
        
        bert_encoded_layers = self.dropout(bert_encoded_layers_raw)

        
        emotion_embeddings = self.emotion_embeddings(emotion_ids)
        opinion_embeddings = self.opinion_embeddings(opinion_ids)
        eks = torch.cat((opinion_embeddings,  emotion_embeddings), -1)
        concat = torch.cat((bert_encoded_layers, eks), -1)
        gru_all_hidden, gru_last_hidden = self.gru(concat, torch.zeros(2, concat.shape[0], self.config.hidden_size * 3).to(device))

        gru_last_hidden_dir0 = gru_last_hidden[0, :, :]
        gru_last_hidden_dir1 = gru_last_hidden[1, :, :]


        gru_last_hidden_stacked = torch.cat(
            (gru_last_hidden_dir0, gru_last_hidden_dir1), dim=1
        )

        gru_avg = torch.mean(gru_all_hidden, dim=1)
        gru_max, _ = torch.max(gru_all_hidden, dim=1)
        gru_complete_concatted = torch.cat(
            (gru_last_hidden_stacked, gru_avg, gru_max), dim=1
        )
        logits = self.dense(gru_complete_concatted)
        
        return logits

训练代码 #

roberta = CustomRoberta(config=config, num_labels=3, max_em_len_1=3, max_em_len_2=4)
roberta = roberta.to(device)
param_optimizer = list(roberta.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.1},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]


batch_size = 16
epochs = 5
optimizer = AdamW(optimizer_grouped_parameters, lr=5e-6)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0.1*total_steps, num_training_steps=total_steps)
loss_fn = torch.nn.CrossEntropyLoss()


for epoch in range(epochs):
            loss_epoch = []
            for step, batch in enumerate(train_dataloader):
                model.train()
                input_ids, opinion_ids, input_mask, emotion_ids, label_ids = batch
                logits = model(input_ids=input_ids.to(device), opinion_ids=opinion_ids.to(device), attention_mask=input_mask.to(device), emotion_ids=emotion_ids.to(device))
                loss = loss_fn(logits, label_ids.to(device))
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                loss_epoch.append(loss.item())

                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

已排查项 #

• 反向传播时权重有变化
• DataLoader工作正常且shuffle=True
• 尝试在小子集上训练多轮，无法过拟合
• 调整学习率（升高/降低）无效
• 移除GRU模块，仅训练RoBERTa+全连接层仍无效

解决方向 #

1. 修复预训练权重被覆盖问题 #

__init__中调用self.apply(self.init_bert_weights)会对所有Linear、Embedding层执行xavier初始化，包括RoBERTa的预训练权重，直接废掉预训练优势。需改为仅初始化新增的自定义层：

# 替换原self.apply(self.init_bert_weights)
self.init_bert_weights(self.emotion_embeddings)
self.init_bert_weights(self.opinion_embeddings)
self.init_bert_weights(self.gru)
self.init_bert_weights(self.dense)
self.init_bert_weights(self.linear)

2. 压缩全连接层输入维度 #

当前dense层输入维度为1024*3*3*2=18432（roberta-large的hidden_size为1024），维度爆炸会导致梯度消失、参数难以训练。可：

• 先降维：在拼接特征后添加Linear层将维度压缩至1024，再连接到num_labels
• 减少拼接特征：比如只保留gru_last_hidden_stacked+gru_avg，去掉gru_max

3. 优化GRU初始隐藏状态 #

手动传入的全0初始隐藏状态不可训练，可改为可训练参数：

# 在__init__中添加
self.gru_h0 = nn.Parameter(torch.zeros(2, self.config.hidden_size*3))
# forward中替换原h0初始化代码
gru_all_hidden, gru_last_hidden = self.gru(concat, self.gru_h0.unsqueeze(1).repeat(1, concat.shape[0], 1).to(device))

4. 调整优化器参数 #

• 权重衰减0.1过高，建议改为0.01
• 若修复预训练权重后仍无改善，可尝试将学习率调整至2e-5~5e-5区间

5. 清理冗余代码 #

get_att函数未被使用，且att_lin未定义，属于无效代码，建议删除避免干扰。

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