1 引言
上文 精通算法>推荐算法26:行为序列建模之BST— Transformer建模用户行为序列-CSDN博客
讲解了BST的背景和模型结构,本文给出其代码实现,供大家参考。
2 BST核心代码
Transformer已经成为了算法工程师的必备技能,因此这一节给出BST的Transformer层代码。代码基于Keras深度学习库实现,其中Keras版本为2.10。
先看Multi-Head Self Attention的实现。
class MultiHeadSelfAttention(keras.layers.Layer):def __init__(self, emb_dim, num_heads=4, **kwargs):"""构建多头自注意力@param emb_dim: 输入embedding维度,必须为多头数目的整数倍@param num_heads: 多头的数目,一般取8、4、2等"""self.emb_dim = emb_dimself.num_heads = num_heads# 定义q、k、v三个权重矩阵self.w_query = keras.layers.Dense(256)self.w_key = keras.layers.Dense(256)self.w_value = keras.layers.Dense(256)# 每个头内的向量维度,等于总维度除以多头数self.dim_one_head = 256 // num_heads# 全连接单元,将多个头连接起来。输出向量与原始输入向量维度相同,从而可以进行残差连接。self.w_combine_heads = keras.layers.Dense(emb_dim)super(MultiHeadSelfAttention, self).__init__(**kwargs)def attention(self, query, key, value):"""attention计算,softmax((q * k) / sqrt(dk, 0.5)) * v@param query: q向量@param key: k向量@param value: v向量@return:"""# 1 内积,计算q和k的相关性权重,得到一个标量score = tf.matmul(query, key, transpose_b=True)# 2 计算向量长度,dkdim_key = tf.cast(tf.shape(key).shape[-1], tf.float32)# 3 缩放,向量长度上做归一化score = score / tf.math.sqrt(dim_key)# 4 Softmax归一化,使得权重处于0到1之间att_scores = tf.nn.softmax(score, axis=-1)# 权重乘以每个v向量,再累加起来,最终得到一个向量,维度与q、k、v相同att_output = tf.matmul(att_scores, value)return att_output, att_scoresdef build_multi_head(self, x_input, batch_size):"""分割隐向量到多个head中,所以多头并不会带来维度倍增@param x_input: 输入向量,可以为q、k、v等向量@param batch_size:@return: 多头矩阵"""x_input = tf.reshape(x_input, shape=(batch_size, -1, self.num_heads, self.dim_one_head))return tf.transpose(x_input, perm=[0, 2, 1, 3])def call(self, inputs, **kwargs):"""多头自注意力计算部分@param inputs:@param kwargs:@return:"""x_query = inputs[0]x_key = inputs[1]batch_size = tf.shape(x_query)[0]# 得到q向量,原始输入经过线性变换,然后再进行多头切割query = self.w_query(x_query)query = self.build_multi_head(query, batch_size)# 得到k向量key = self.w_key(x_key)key = self.build_multi_head(key, batch_size)# 得到v向量, v和k用同一个输入value = self.w_value(x_key)value = self.build_multi_head(value, batch_size)# attention计算att_output, att_scores = self.attention(query, key, value)att_output = tf.transpose(att_output, perm=[0, 2, 1, 3]) # [batch_size,seq_len,num_heads,dim_one_head]att_output = tf.reshape(att_output, shape=(batch_size, -1, self.dim_one_head * self.num_heads))# 多头合并,并进行线性连接输出output = self.w_combine_heads(att_output) # [batch_size,seq_len,emb_dim]return output
代码采用Keras标准的自定义Layer实现,重点看call函数和attention函数。call函数先将原始输入向量,通过线性连接,得到query、key、value三个向量。然后通过attention函数计算得到每个头的输出。最后通过一层线性连接来融合多头,得到一个与输入向量维度相同的输出向量。Attention函数的实现步骤为:
- 计算query和key的内积,得到二者相关性权重
- 计算query和key向量长度,二者长度相同,后续归一化会用到
- 对第一步得到的权重进行缩放,相当于在向量长度上做归一化
- Softmax归一化,使得权重处于0到1之间
- 权重乘以每个v向量,再累加起来,最终得到一个向量,维度与q、k、v相同
再看Transformer的整体实现。
class Transformer(keras.layers.Layer):def __init__(self, seq_len, emb_dim, num_heads=4, ff_dim=128, **kwargs):# position emb,位置编码向量self.seq_len = seq_lenself.emb_dim = emb_dimself.positions_embedding = Embedding(seq_len, self.emb_dim, input_length=seq_len)# multi-head self attention层self.att = MultiHeadSelfAttention(self.emb_dim, num_heads)# 两层feed-forward全连接self.ffn = keras.Sequential([keras.layers.Dense(ff_dim, activation="relu"),keras.layers.Dense(self.emb_dim)])# 两次layerNorm,一次为input和attention输出残差连接,另一次为attention输出和全连接输出self.ln1 = keras.layers.LayerNormalization()self.ln2 = keras.layers.LayerNormalization()self.dropout1 = keras.layers.Dropout(0.3)self.dropout2 = keras.layers.Dropout(0.3)super(Transformer, self).__init__(**kwargs)def call(self, inputs, **kwargs):"""结构:一层multi-head self attention, 叠加两层feed-forward,中间加入layeNorm和残差连接@param inputs:原始输入向量,可以包括物品id、类目id、品牌id等特征的Embedding@param kwargs:@return:单层Transformer结构的输出"""# 1 构建位置特征,并进行Embeddingpositions = tf.range(start=0, limit=self.seq_len, delta=1)positions_embedding = self.positions_embedding(positions)# 2 原始输入和位置向量加起来,也可以采用BST原文中的Concat方法x_key = inputs + positions_embedding# 3 multi-head self attention,利用layerNorm和输入进行残差连接att_output = self.att(inputs)att_output = self.dropout1(att_output)output1 = self.ln1(x_key + att_output)# 4 两层feed-forward全连接,利用layerNorm和输入进行残差连接ffn_output = self.ffn(output1)ffn_output = self.dropout2(ffn_output)output2 = self.ln2(output1 + ffn_output) # [B, seq_len, emb_dim]# 5 平均池化,对序列进行压缩result = tf.reduce_mean(output2, axis=1) # [B, emb_dim]return result
重点看call函数,其实现步骤为:
- 构建位置特征,并进行Embedding。此处为了方便,用位置的index来表示时序关系,与BST有一点细微区别。
- 将原始输入向量和位置向量相加,得到Transformer层的输入。也可以采用BST原文中的Concat方法。
- 将Transformer层的输入送到Multi-Head Self Attention模块中,经过LayerNorm后,再和输入进行残差。得到Attention模块的输出。
- 将Attention的输出送入到Feed-Forward模块中,它是两层全连接结构,经过LayerNorm后,再和Attention的输出进行残差连接。得到Feed-Forward模块的输出。
- 将Feed-Forward模块输出的多个向量,进行平均池化,从而将不定长的序列压缩为一个固定长度的向量。
3 BST总结和思考
BST利用Transformer来建模用户行为序列,可以有效解决传统RNN模型容易出现的梯度弥散、串行耗时长等一系列问题,并取得了不错的业务效果。目前Transformer已经广泛应用在搜索推荐广告中,特别是用户行为序列建模中。BST作为一次完整的Transformer在大规模工业场景中的落地,其各种细节的处理和参数的配置,都非常值得学习和借鉴。
4 作者新书推荐
历经两年多,花费不少心血,终于撰写完成了这部新书。本文在5.4节中重点阐述了。
源代码:扫描图书封底二维码,进入读者群,群公告中有代码下载方式
微信群:图书封底有读者微信群,作者也在群里,任何技术、offer选择和职业规划的问题,都可以咨询。