使用Bert预训练模型处理序列推荐任务

hhy不许摸鱼2023-08-10 9:11

最近的工作有涉及该任务，整理一下思路以及代码细节。

流程

总体来说思路就是首先用预训练的bert模型，在训练集的序列上进行CLS任务。对序列内容（这里默认是token id的sequence）以0.3左右的概率进行随机mask，然后将相应sequence的attention mask（原来决定padding index）和label（也就是mask的ground truth）输入到bert model里面。

当然其中vocab.txt并不存在的token是需要add进去的，具体方法不再详述，网上例子很多，注意word embedding也需要初始化就行。

模型定义：
self.model = AutoModelForMaskedLM.from_pretrained('./bert')

模型的输入：
result = self.bert_model(tail_mask, attention_mask, labels)

得到模型训练的结果之后，要做一个选择：

（1）transformer的bert model可以输出要预测时间步的hidden state，可以选择取出对应的hidden state，其中需要在数据处理的时候记录下每个sequence的tail position，也就是要预测位置的idx。另外我认为既然要进行序列推荐，那么最后一个tail position的token表征一定是最重要的，所以需要对tail position的idx专门给个写死的mask，效果会好一些。然后与sequence中item的全集进行相似度的计算，再去算交叉熵loss。

py 复制代码

bert_hidden = result.hidden_states[-1]
bert_seq_hidden = torch.zeros((self.args.batch_size, 312)).to(self.device)
for i in range(self.args.batch_size):
	bert_seq_hidden[i,:] = bert_hidden[i, tail_pos[i], :]
logits = torch.matmul(bert_seq_hidden, test_item_emb.transpose(0, 1))
main_loss = self.criterion(logits, targets)

（2）同时也可以result.loss直接数据mask prediction的loss，我理解这个loss面对的任务是我要求sequence中的各个token表征都要尽可能准确，都要考虑，（1）可能更加注重最后一个位置的标准的准确性。

然后在evaluate阶段，需要注意输入到模型的不再是tail_mask，而是仅仅mask掉tail token id的sequence，因为我们需要尽可能准确的序列信息，只需要保证要预测的存在mask就够了。

由于是推荐任务，而且bert得到的hidden state表征过于隐式，所以需要一定的个性化引导它进行训练。经过个人的实验也确实如此，而且结果相差很多。

以上就是我个人的总结经验，欢迎大家指点。