BERT的MLM是我见过的比较复杂的优化任务了,我以此为例分享一下对于这种比较复杂的损失函数如何计算。先说下这个MLM是什么,如下图所示(摘自BERT的原始论文):
所以这个任务就是将sentence中一些token进行掩盖,模型会输出这些掩盖的token的隐藏状态,将这些隐藏状态输入softmax可以得到候选单词的概率分布,这样根据ground truth就可以计算cross entropy了。原理比较简单,要想使用tf高效实现,博主着实费了一些脑细胞。我的计算方法如下:
import tensorflow as tf
final_output=tf.random.normal([3,4,8],
name='final_output') #[batch_size,seq_len,category_size]
mask=tf.constant([[0,1,0,0],
[1,0,0,0],
[0,0,1,0]],dtype=tf.int32)
mask_value=tf.constant([[0,3,0,0],
[6,0,0,0],
[0,0,5,0]],dtype=tf.int32)
final_output就是模型最终的输入了,即softmax的输出结果,咱们假设类别共有8类。mask标记哪些token被掩盖了,掩盖的标1。mask_value是被掩盖的token的真实类别。我的想法就是先不管是否掩盖,计算所有token的cross entropy,然后只保留被掩盖的token的cross entropy,其余的cross entropy记为0,最后相加就平均就好啦。
mask=tf.cast(mask,dtype=tf.float32)
loss_all=tf.keras.losses.sparse_categorical_crossentropy(mask_value,final_output) #[batch_size,seq_len]
loss=tf.reduce_sum(loss_all*mask)/tf.reduce_sum(mask)
大家要是有更高效的计算方法,欢迎留言交流哈~~