基于动态嵌入训练推荐模型

摘要：本文介绍了基于动态嵌入技术的推荐系统实现方案。通过TensorFlow推荐系统扩展库(TFRA)的动态嵌入功能，可以解决大规模推荐场景中静态嵌入表内存占用过高的问题。文章详细演示了如何结合TFRS与TFRA构建电影推荐模型，包括数据处理、双塔模型构建、动态嵌入层实现等关键步骤。实验表明，动态嵌入模型在保持与基准模型相当性能的同时，能根据需求自动扩容/缩容嵌入表。文中还探讨了嵌入表缩容策略及验证方案设计要点，为大规模推荐系统开发提供了实用技术参考。

目录

基于动态嵌入训练推荐模型

导入库

数据处理

构建双塔模型

训练模型

构建基于动态嵌入的模型

概述

嵌入的类型

定义嵌入层

创建并编译最终模型

训练模型

嵌入表的缩容

性能分析

总结

基于动态嵌入训练推荐模型

现代推荐系统高度依赖嵌入技术，为每位用户和候选物品生成向量表征。这些嵌入向量可用于计算用户与物品之间的相似度，从而为用户推荐更贴合其兴趣、更具相关性的候选物品。但在处理大规模数据时，尤其是在线机器学习场景下，嵌入表的规模会急剧膨胀，累积数百万甚至数十亿个条目。在这种规模下，将嵌入表完整存储在内存中成为不可能的事。此外，有很大一部分条目出现的频率极低，为这类稀有条目单独维护嵌入向量毫无意义。一种更优的解决方案是用一个通用嵌入向量来表征这些条目，这能以极小的性能损耗为代价，大幅缩减嵌入表的体积，这也是动态嵌入表诞生的核心动因。

TensorFlow 内置的tf.keras.layers.Embedding层在创建时尺寸固定，因此我们需要另一种实现方案。幸运的是，TensorFlow 特别兴趣小组（SIG）有一个专门的项目可实现该需求 ——TensorFlow 推荐系统扩展库（TFRA）。你可以从其代码仓库了解更多细节，简而言之，TFRA 借助动态嵌入技术实现了嵌入表尺寸的动态调整，相比静态嵌入能取得更优的推荐效果。TFRA 完全兼容 TensorFlow 2.0，可与大家熟悉的 Keras API 接口无缝协作，因此能轻松与 TensorFlow 推荐系统（TFRS）等其他 TensorFlow 产品集成。

在本教程中，我们将结合 TFRS 与 TFRA 构建一个电影推荐模型，所使用的 MovieLens 数据集包含用户对电影的匿名评分数据。教程的核心重点是演示如何利用 TensorFlow 推荐系统扩展库提供的动态嵌入功能，在推荐场景中实现嵌入表尺寸的动态扩容与缩容。你可以在此处找到完整的实现代码，在此处查看分步讲解。

导入库

首先导入所需的各类库：

import tensorflow as tf
import tensorflow_datasets as tfds

# TFRA会对TensorFlow进行部分补丁修改，因此必须在导入TensorFlow后再导入
import tensorflow_recommenders as tfrs
import tensorflow_recommenders_addons as tfra
import tensorflow_recommenders_addons.dynamic_embedding as de

注意，我们在导入 TensorFlow 之后才导入 TFRA 库。建议遵循这一导入顺序，因为 TFRA 需要对 TensorFlow 执行部分补丁修改操作。

数据处理

我们先基于 TensorFlow 推荐系统构建一个基准模型，遵循 TFRS 检索任务教程的范式，搭建一个双塔检索模型。其中用户塔以用户 ID 为输入，物品塔则以经过分词处理的电影标题为输入。

为处理电影标题，我们定义一个辅助函数：将电影标题转为小写，移除其中的所有标点符号，按空格切分生成词元列表；最终从标题开头截取最多max_token_length个词元，若电影标题的词元数量少于该值，则保留全部词元。该数值是经数据分析确定的，对应数据集中标题长度的 90 百分位数。

max_token_length = 6
pad_token = "[PAD]"
punctuation_regex = "[\\!\\"#\\$%&\\*\\+,-\\.\\/\\:;\\<\\=\\>\\?@\\[\\]\\\\\\^_`\\{\\|\\}~\\\\t\\\\n]"

# 定义处理电影标题的辅助函数
def process_text(x: tf.Tensor, max_token_length: int, punctuation_regex: str) -> tf.Tensor:
return tf.strings.split(
tf.strings.regex_replace(
tf.strings.lower(x["movie_title"]), punctuation_regex, ""
)
)[:max_token_length]

我们还会对分词后的电影标题做补零处理，使其长度固定；并使用相同的随机种子划分数据集，确保各训练轮次的验证结果具有一致性。你可以在笔记本的「数据集处理」章节找到详细代码。

构建双塔模型

我们的用户塔与 TFRS 检索教程中的实现基本一致（仅网络结构更深），而物品塔在嵌入查找后新增了一个GlobalAveragePooling1D层，将电影标题词元的嵌入向量做平均处理，生成一个单一的嵌入向量。

def get_movie_title_lookup_layer(dataset: tf.data.Dataset) -> tf.keras.layers.Layer:
movie_title_lookup_layer = tf.keras.layers.StringLookup(mask_token=pad_token)
movie_title_lookup_layer.adapt(dataset.map(lambda x: x["movie_title"]))
return movie_title_lookup_layer

def build_item_model(movie_title_lookup_layer: tf.keras.layers.StringLookup):
vocab_size = movie_title_lookup_layer.vocabulary_size()
return tf.keras.models.Sequential([
tf.keras.layers.InputLayer(input_shape=(max_token_length), dtype=tf.string),
movie_title_lookup_layer,
tf.keras.layers.Embedding(vocab_size, 64),
tf.keras.layers.GlobalAveragePooling1D(),
tf.keras.layers.Dense(64, activation="gelu"),
tf.keras.layers.Dense(32),
tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1))
])

训练模型

模型训练的操作十分简单，只需向模型的fit()方法传入所需参数即可，我们将使用验证数据集validation_ds评估模型性能。

history = model.fit(datasets.training_datasets.train_ds, epochs=3, validation_data=datasets.training_datasets.validation_ds)

训练结束后，输出结果如下所示：

plaintext

Epoch 3/3
220/220 [==============================] – 146s 633ms/step
……
val_factorized_top_k/top_10_categorical_accuracy: 0.0179 – val_factorized_top_k/top_50_categorical_accuracy: 0.0766 – val_factorized_top_k/top_100_categorical_accuracy: 0.1338 – val_loss: 12359.0557 – val_regularization_loss: 0.0000e+00 – val_total_loss: 12359.0557

我们在验证数据集上取得了 13.38% 的 Top-100 分类准确率。

构建基于动态嵌入的模型

概述

接下来，我们将学习如何使用 TensorFlow 推荐系统扩展库中的动态嵌入，替代静态的嵌入表。顾名思义，动态嵌入不会预先为词汇表中的所有条目创建嵌入向量，而是根据需求动态扩容嵌入表。当企业需要处理数百万、数十亿的物品和用户数据时，这一特性的优势会体现得淋漓尽致 —— 对这些企业而言，静态嵌入表体积过大、无法载入内存是常见问题，其大小甚至可达数百 GB 乃至数 TB，即便是云环境中内存配置最高的实例也无法承载。

当嵌入表的基数极大时，权重的访问会呈现高度稀疏的特点。因此，我们会采用基于哈希表的数据结构存储权重，并在每一轮迭代中从底层表结构中调取所需的权重。本教程为聚焦库的核心功能，将采用非分布式部署场景，此时 TFRA 会默认选用杜鹃哈希表，同时也支持 Redis、nvhash 等其他解决方案。

附：TFRA 库中分布式与非分布式场景下的各类嵌入解决方案对比图

使用动态嵌入时，我们会为嵌入表设置一个初始容量，模型训练过程中遇到新的 ID 时，嵌入表会根据需求自动扩容。如需了解该功能的设计动因与内部实现机制，可参考相关的请求评论文档（RFC）。

嵌入的类型

目前，TFRA 的dynamic_embedding模块提供三种嵌入类型：

定义嵌入层

我们将使用Embedding表征用户 ID，使用SquashedEmbedding表征词元 ID。需要注意的是，每个电影标题包含多个词元，因此需要将生成的词元嵌入向量归约为一个具有代表性的单一嵌入向量。

注意：Embedding 的功能在 TFRA 0.5 至 0.6 版本中发生了变更，本教程请务必使用 0.6 版本。

基于此，我们可以像构建标准模型一样定义双塔结构，不同的是，本次将使用动态嵌入层替代静态嵌入层：

def build_de_user_model(user_id_lookup_layer: tf.keras.layers.StringLookup) -> tf.keras.layers.Layer:
vocab_size = user_id_lookup_layer.vocabulary_size()
return tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(), dtype=tf.string),
user_id_lookup_layer,
de.keras.layers.Embedding(
embedding_size=64,
initializer=tf.random_uniform_initializer(),
init_capacity=int(vocab_size*0.8),
restrict_policy=de.FrequencyRestrictPolicy,
name="UserDynamicEmbeddingLayer"
),
tf.keras.layers.Dense(64, activation="gelu"),
tf.keras.layers.Dense(32),
tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1))
], name='user_model')

def build_de_item_model(movie_title_lookup_layer: tf.keras.layers.StringLookup) -> tf.keras.layers.Layer:
vocab_size = movie_title_lookup_layer.vocabulary_size()
return tf.keras.models.Sequential([
tf.keras.layers.InputLayer(input_shape=(max_token_length), dtype=tf.string),
movie_title_lookup_layer,
de.keras.layers.SquashedEmbedding(
embedding_size=64,
initializer=tf.random_uniform_initializer(),
init_capacity=int(vocab_size*0.8),
restrict_policy=de.FrequencyRestrictPolicy,
combiner="mean",
name="ItemDynamicEmbeddingLayer"
),
tf.keras.layers.Dense(64, activation="gelu"),
tf.keras.layers.Dense(32),
tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1))
])

完成用户塔和物品塔的定义后，我们即可按常规方式定义检索模型。

创建并编译最终模型

模型构建的最后一步，是创建模型并完成编译：

def create_de_two_tower_model(dataset: tf.data.Dataset, candidate_dataset: tf.data.Dataset) -> tf.keras.Model:
user_id_lookup_layer = get_user_id_lookup_layer(dataset)
movie_title_lookup_layer = get_movie_title_lookup_layer(dataset)
user_model = build_de_user_model(user_id_lookup_layer)
item_model = build_de_item_model(movie_title_lookup_layer)
task = tfrs.tasks.Retrieval(
metrics=tfrs.metrics.FactorizedTopK(
candidate_dataset.map(item_model)
),
)

model = DynamicEmbeddingTwoTowerModel(user_model, item_model, task)
optimizer = de.DynamicEmbeddingOptimizer(tf.keras.optimizers.Adam())
model.compile(optimizer=optimizer)

return model

datasets = create_datasets()
de_model = create_de_two_tower_model(datasets.training_datasets.train_ds, datasets.candidate_dataset)

注意，我们在标准 TensorFlow 优化器外层封装了DynamicEmbeddingOptimizer。这一步是必须的，因为该封装器能为训练存储在哈希表中的权重提供专属的功能支持。至此，我们即可开始训练模型。

训练模型

基于动态嵌入的模型训练流程相当直观，只是为了记录更多额外信息，需要稍作额外处理 —— 我们将通过tf.keras.callbacks.Callback回调函数实现日志记录，并将其命名为DynamicEmbeddingCallback。

epochs = 3
history_de = {}
history_de_size = {}
de_callback = DynamicEmbeddingCallback(de_model, steps_per_logging=20)

for epoch in range(epochs):
datasets = create_datasets()
train_steps = len(datasets.training_datasets.train_ds)

hist = de_model.fit(
datasets.training_datasets.train_ds,
epochs=1,
validation_data=datasets.training_datasets.validation_ds,
callbacks=[de_callback]
)

for k,v in de_model.dynamic_embedding_history.items():
if k=="step":
v = [vv+(epoch*train_steps) for vv in v]
history_de_size.setdefault(k, []).extend(v)

for k,v in hist.history.items():
history_de.setdefault(k, []).extend(v)

我们将训练轮次的循环逻辑从fit()方法中抽离出来，在每一轮训练中重新创建数据集，以实现训练数据的不同随机打乱方式，并在循环内完成单轮次的模型训练。最后，将自定义回调函数记录的嵌入表尺寸信息汇总至history_de_size，将模型性能指标汇总至history_de。

该回调函数的实现代码如下：

class DynamicEmbeddingCallback(tf.keras.callbacks.Callback):
def __init__(self, model, steps_per_logging, steps_per_restrict=None, restrict=False):
self.model = model
self.steps_per_logging = steps_per_logging
self.steps_per_restrict = steps_per_restrict
self.restrict = restrict

def on_train_begin(self, logs=None):
self.model.dynamic_embedding_history = {}

def on_train_batch_end(self, batch, logs=None):
# 触发嵌入表缩容
if self.restrict and self.steps_per_restrict and (batch+1) % self.steps_per_restrict == 0:
[
self.model.embedding_layers[k].params.restrict(
num_reserved=int(self.model.lookup_vocab_sizes[k]*0.8),
trigger=self.model.lookup_vocab_sizes[k]-2 # 排除未知词元和补零词元
) for k in self.model.embedding_layers.keys()
]

# 记录嵌入表尺寸
if (batch+1) % self.steps_per_logging == 0:
embedding_size_dict = {
k:self.model.embedding_layers[k].params.size().numpy()
for k in self.model.embedding_layers.keys()
}

for k, v in embedding_size_dict.items():
self.model.dynamic_embedding_history.setdefault(f"embedding_size_{k}", []).append(v)
self.model.dynamic_embedding_history.setdefault(f"step", []).append(batch+1)

该回调函数主要实现两大功能：

接下来，我们来理解嵌入表缩容的含义及其重要性。

嵌入表的缩容

我们尚未探讨一个重要问题：当嵌入表的规模超过预设阈值时，该如何对其进行缩容。这是一项实用性极强的功能，允许我们为嵌入表设置最大规模阈值，从而在处理大词汇表的同时，将内存占用控制在硬件的内存限制范围内。实现方式是通过DynamicEmbeddingCallback中调用嵌入层底层变量的restrict()方法，该方法接收两个参数：num_reserved（缩容后的目标尺寸）和trigger（触发缩容的临界尺寸）。嵌入表的缩容策略由层构造函数中的restrict_policy参数定义，本教程中使用的是FrequencyRestrictPolicy（频率限制策略），即从嵌入表中移除出现频率最低的条目。用户可通过设置DynamicEmbeddingCallback中的steps_per_restrict和restrict参数，定义嵌入表缩容的触发频率。

嵌入表缩容在流数据场景下的实用性会更高，以在线学习场景为例：模型需要每天（甚至每小时）基于新接入的数据进行训练，此时外层的训练轮次循环可代表训练的天数，每天接收一份新的数据集（例如前一天的用户行为交互数据），并基于上一次的模型检查点继续训练。这种情况下，当嵌入表的规模超过trigger参数定义的阈值时，即可通过DynamicEmbeddingCallback触发缩容操作。

性能分析

我们在此分析三种模型变体的性能表现：

附：动态嵌入开启 / 关闭状态下的模型准确率对比图

从图中可看出，基于动态嵌入的模型（绿色实线）在验证集上的性能与基准模型（红色实线）相当，训练集的准确率也呈现相似的趋势。在实际的大规模在线学习场景中，动态嵌入往往能带来准确率的提升。

最后我们能看到，执行缩容操作后，模型的验证准确率出现了一定程度的下降，这一现象符合预期。因为本教程使用的数据集规模相对较小、条目数量有限，缩容操作可能会移除那些本应保留在嵌入表中的嵌入向量。例如，可增大restrict函数中的num_reserved参数（如设置为int(self.model.lookup_vocab_sizes[k]*0.95)），此时模型性能会向未执行缩容操作的版本靠拢。

接下来，我们观察嵌入表的尺寸随训练过程的动态变化情况。

附：嵌入表尺寸随时间的变化趋势图

从图中可见，未执行缩容操作时，嵌入表会扩容至词汇表的完整尺寸（虚线）并保持稳定；而触发缩容操作后（点线），嵌入表的尺寸会下降，当遇到新的 ID 时又会再次扩容。

需要特别说明的是，构建合理的验证方案并非易事，需要谨慎考虑样本外验证、时间外验证、分层抽样等诸多因素。本教程为简化演示，未考虑这些因素，仅通过对现有数据集随机抽样构建了验证集。

总结

当处理包含数百万、数十亿实体的大规模物品集时，使用动态嵌入表是开展表征学习的高效方式。在本教程中，我们学习了如何利用 TensorFlow 推荐系统扩展库提供的dynamic_embedding模块实现这一技术：首先对数据进行探索，提取模型训练和评估所需的特征并构建tf.data.Dataset数据集对象；随后定义了一个使用静态嵌入表的模型，作为性能评估的基准；接着创建了基于动态嵌入的模型并在数据上完成训练，发现动态嵌入表能根据需求动态扩容，且模型性能与基准模型相当；同时还探讨了如何利用restrict功能，在嵌入表规模超过预设阈值时对其进行缩容。

我们希望本教程能为你提供关于 TFRA 和动态嵌入的基础概念讲解，助力你思考如何利用该技术优化自身的推荐系统。如需进行更深入的交流探讨，欢迎访问 TFRA 的代码仓库。

————————TensorFlow 核心框架 | TensorFlow 推荐系统 | TensorFlow 推荐系统扩展库