三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

TokenFormer:统一多域与序列推荐的Transformer框架解析

TokenFormer:统一多域与序列推荐的Transformer框架解析

1. 项目概述:当多域推荐遇上序列推荐

最近在梳理推荐系统领域的前沿论文,看到一篇挺有意思的工作,标题叫《TokenFormer: Unify the Multi-Field and Sequential Recommendation Worlds》。光看这个标题,就感觉它想干一件大事——把两个看似不同的推荐世界给统一起来。一个是多域推荐,简单说就是一个平台上有多种类型的商品或内容(比如电商里同时有图书、服装、家电),我们需要同时考虑用户在这些不同领域的偏好。另一个是序列推荐,这是我们更熟悉的,根据用户历史点击或购买的顺序,预测下一个可能感兴趣的东西。这两个问题通常被分开研究,各有各的模型套路,但这篇论文提出用一个统一的框架TokenFormer来搞定,这个思路本身就很有启发性。

为什么说统一这两个“世界”有价值呢?从实际业务角度看,用户的行为从来不是割裂的。一个用户可能先在视频平台看了几个美食纪录片(序列行为),然后去电商平台搜索了相关的厨具(跨域兴趣)。传统的序列模型可能只盯着看视频的序列,忽略了“美食”这个领域属性;而传统的多域模型可能只统计了用户对“家电”、“图书”等大类的偏好,却丢失了“先看食谱视频,再买烘焙模具”这种精细的时序因果关系。TokenFormer的核心思想,就是用Transformer那套强大的建模能力,把用户行为序列中的每个物品,以及这个物品所属的“域”信息,都当成一种“令牌”来处理,通过自注意力机制让它们充分交互,从而同时捕捉序列模式和跨域关联。

这篇笔记,我就结合自己的理解,拆解一下TokenFormer是怎么设计的,它背后的动机是什么,以及我们在实际尝试复现或借鉴其思想时需要注意哪些坑。无论你是刚入门推荐算法的新手,还是想寻找模型创新点的老手,希望这篇深度解读都能带来一些收获。

2. 核心思路拆解:用“令牌”统一一切

2.1 多域推荐与序列推荐的固有矛盾

在深入TokenFormer之前,我们得先看清楚它要解决的核心矛盾。多域推荐和序列推荐,长期以来在数据表征和模型设计上走的几乎是两条路。

多域推荐的典型做法,是把“域”作为一个重要的特征字段。例如,我们有一个“领域ID”特征,取值可能是“图书”、“电子”、“服装”。在模型里,通常会为每个域学习一个独立的嵌入向量,或者通过一些参数共享与分离的机制来建模域间的共性与特性。它的目标往往是提升跨域的整体CTR/CVR,或者解决冷启动问题(用一个域的数据帮助另一个域)。但这类模型大多基于用户的历史统计特征(如过去30天在各域的点击分布),对行为发生的顺序不那么敏感。

序列推荐则恰恰相反,它的灵魂就是顺序。从早期的马尔可夫链到后来的GRU4Rec、SASRec、BERT4Rec,模型的核心任务是捕捉序列中的转移模式。比如,用户买了手机,下一个很可能会买手机壳。在这里,物品的ID是核心,至于这个手机属于“电子产品”域,手机壳属于“配件”域,这个信息往往只是作为一个辅助特征,甚至有时会被忽略。序列模型擅长捕捉“接下来是什么”,但对于“这个推荐结果是否平衡了用户在不同领域的长期兴趣”这类问题,就显得力不从心。

TokenFormer的作者认为,这种割裂限制了模型的表达能力。用户选择看一部电影(娱乐域),紧接着购买相关的电影原著小说(图书域),这个行为链条里既包含了强时序性,也包含了跨域的兴趣迁移。一个理想的模型应该能无缝地理解这一点。

2.2 TokenFormer的统一建模哲学

TokenFormer的解决方案非常“Transformer”——既然Transformer能用统一的注意力机制处理各种NLP任务中的令牌,那为什么不能把推荐系统中的一切元素都“令牌化”呢?

在这个框架里,一个用户的一次交互记录(比如点击了一件商品),被表示成两个核心令牌:

  1. 物品令牌:代表被交互物品本身的ID。
  2. 域令牌:代表该物品所属的领域。

此外,为了标记序列的开始和捕获全局信息,还会引入一个特殊的[CLS]令牌。这样一来,一个用户的行为序列就被转化成了一个令牌序列,例如:[CLS], 物品1, 域1, 物品2, 域2, ..., 物品T, 域T

这个简单的转换是统一建模的关键。接下来,所有这些令牌被送入一个多层Transformer编码器。自注意力机制在这里大显神威:

  • 物品令牌之间可以计算注意力,从而捕捉序列中物品到物品的转移规律(这是序列推荐的核心)。
  • 域令牌之间可以计算注意力,从而学习域到域的宏观关联(例如,买过家电的用户可能也对家装建材感兴趣)。
  • 物品令牌与域令牌之间可以交叉计算注意力,这正是建模“特定物品与其所属领域关系”以及“跨域物品关联”的桥梁。例如,模型可以学到“购买某品牌手机(物品)”与“数码配件(域)”之间存在强关联。

最终,那个[CLS]令牌的最终隐藏状态,聚合了整个序列的信息,可以用来做下一个物品的预测。通过这种方式,Transformer的注意力矩阵自然地成为了融合序列信号和领域信号的熔炉。

注意:这里的一个关键设计细节是位置编码。由于我们引入了额外的域令牌,原始序列的长度翻倍了。因此,需要精心设计位置编码,确保物品令牌和其对应的域令牌在位置关系上是紧耦合的,同时又能区分序列的先后顺序。论文中通常会对奇偶位置进行不同的处理或使用成对的位置编码。

2.3 与经典模型的对比思考

理解了TokenFormer的统一思想后,我们可以把它和几个经典模型做个对比,这样能更清楚它的创新点在哪。

  • 对比SASRec/BERT4Rec:这类纯序列模型只有物品令牌。它们能很好地学习物品A -> 物品B的转移概率,但无法显式地建模“物品A属于域X”这个信息对预测的影响。TokenFormer通过引入域令牌,让模型在关注物品相似性的同时,也能意识到领域的约束或促进关系。
  • 对比MMoE/PLE等多任务多域模型:这类模型结构上通常有共享层和多个专家网络(对应不同域),它们的目标是优化多任务损失。但它们的输入往往是特征向量,对序列的建模能力较弱。TokenFormer则用标准的Transformer块统一处理,结构更简洁,且对序列动态的捕捉能力天生更强。
  • 对比FiBiNet等特征交互模型:这些模型擅长处理并列的稀疏特征域,并通过注意力、交叉网络等方式进行特征交互。但它们通常不处理序列数据。TokenFormer可以看作是将其思想扩展到了序列数据上,并且利用了Transformer更强的全局依赖建模能力。

简单来说,TokenFormer找到了一条用最经典的Transformer架构,同时解决两个重要推荐问题的优雅路径。它没有增加特别复杂的模块,而是通过输入表示的创新,释放了现有架构的潜力。

3. 模型核心细节与实现要点

3.1 令牌嵌入层:一切的开端

TokenFormer的输入层是整个模型的基石,需要将原始的、离散的特征转化为稠密的向量表示。这里主要涉及三种嵌入:

  1. 物品嵌入:一个可学习的查找表,将每个物品ID映射为一个d维向量。这是推荐系统里最常规的操作。
  2. 域嵌入:同样是一个可学习的查找表,将每个域ID映射为d维向量。这里需要注意的是,域的划分粒度是一个超参数。可以按产品大类(如3C、服饰),也可以按更细的频道划分。粒度越细,域令牌携带的信息越具体,但也会增加参数和过拟合风险。
  3. 位置嵌入:这是实现中的第一个难点。由于序列中插入了域令牌,序列长度变为2L+1(L为原始交互序列长度,加1是[CLS])。如果使用标准的正弦余弦位置编码,需要确保其长度能覆盖这个扩展后的序列。更常见的做法是使用可学习的位置嵌入,直接为扩展后序列的每一个位置分配一个d维向量。这样模型可以自由学习最适合这种“物品-域”交错结构的位置关系。

初始的令牌序列表示,就是对应嵌入的加和:H⁰ = [e_cls + p_0, e_item1 + p_1, e_field1 + p_2, e_item2 + p_3, e_field2 + p_4, ...]

3.2 Transformer编码器:注意力如何工作

得到初始嵌入H⁰后,将其送入一个由N层Transformer编码器块堆叠而成的模块。每一层都包含多头自注意力机制和前馈神经网络,并伴有残差连接和层归一化。

自注意力机制是融合信息的关键。我们以单头注意力为例,看看注意力分数是如何计算的。对于第l层的隐藏状态H^l,我们计算查询Q、键K、值VQ = H^l W^Q, K = H^l W^K, V = H^l W^V注意力权重矩阵A = softmax(QK^T / sqrt(d_k)),这个A矩阵的每一个元素A_{ij},就代表了第i个令牌对第j个令牌的关注程度。

在TokenFormer中,这个A矩阵会呈现出非常有趣的模式:

  • 物品->物品:这捕捉了序列推荐的核心,比如物品t(位置1, 3, 5...)会高度关注物品t-1
  • 域->域:这反映了领域间的宏观转移,比如域“图书”(位置2, 4, 6...)可能关注域“影视”,因为存在IP衍生关系。
  • 物品<->域:这是TokenFormer的特色。物品“《三体》”会高度关注域“图书”(这是归属关系),同时也可能关注域“科幻电影”(这是跨域关联)。反过来,域“家电”也会关注那些热销的物品,以了解该领域的当前趋势。

通过多层这样的交互,信息在物品和域之间充分流动,最终使得[CLS]令牌的表示包含了这种混合的、丰富的用户兴趣表征。

3.3 预测层与训练目标

经过N层Transformer后,我们取最后一层[CLS]令牌对应的输出向量h_cls^N作为整个用户序列的表示。

对于下一个物品的预测,这是一个标准的分类问题。我们计算h_cls^N与所有候选物品嵌入向量e_item的内积(或通过一个投影层后再内积),然后通过softmax得到每个物品成为下一个交互物品的概率:P(item_next = i | sequence) = softmax(h_cls^N · e_i^T)

训练目标就是最大化用户真实下一个物品的似然概率,即使用交叉熵损失:L = -log P(item_next = i* | sequence)其中i*是真实的下一个物品。

这里有一个重要的实操细节:在训练时,由于我们预测的是下一个“物品”,因此计算损失时只考虑物品令牌对应的位置。这也符合我们的直觉:我们预测用户会买什么具体商品,而不是预测下一个“域”。当然,域的信息已经通过注意力机制编码进了h_cls^N中,间接影响了预测。

3.4 实现中的关键技巧与超参数选择

  1. 序列长度与掩码:和大多数序列模型一样,需要设定一个最大序列长度L。对于短序列进行填充,对于长序列进行截断(通常保留最近的L次交互)。在训练时,为了防止信息泄露,必须使用因果掩码(或叫序列掩码),确保在预测位置t时,只能看到t之前(以及当前物品对应的域)的令牌,而不能看到t之后的。
  2. 嵌入共享:一个有趣的技巧是,论文中探索了是否让物品嵌入层和预测层中的物品嵌入矩阵共享权重。共享权重可以减少参数量,并可能起到正则化的效果,但有时也会限制表达能力。这是一个需要根据数据集大小进行实验的选项。
  3. 层数与注意力头数:典型的配置可能是N=24层,注意力头数heads=48。对于推荐数据,过深的网络容易过拟合,因为用户行为序列本身噪声较大,且长度有限。
  4. Dropout与正则化:在嵌入层后、注意力权重计算后、以及前馈网络中都应使用Dropout,这是防止Transformer在推荐任务上过拟合的标配。此外,层归一化对训练稳定性至关重要。

4. 实验复现与效果分析视角

虽然我们无法在此完全复现论文实验,但可以梳理出其验证逻辑和我们可以借鉴的评估视角,这对于我们自己设计实验至关重要。

4.1 数据集构造与评估指标

论文通常在多个公开数据集上进行验证,例如Amazon、MovieLens等。关键是如何为这些数据集定义“域”。

  • 粗粒度:在Amazon数据集中,可以直接使用产品类别(如“Books”, “Electronics”)作为域。
  • 细粒度:在MovieLens中,可能需要根据电影的标签(genres)进行聚合或划分来定义域。

评估指标主要来自序列推荐:

  • Hit Rate @ K (HR@K):真实的下一个物品是否出现在预测的Top-K列表中。这是最直观的指标。
  • Normalized Discounted Cumulative Gain @ K (NDCG@K):不仅考虑是否命中,还考虑命中的排名位置,排名越靠前得分越高。
  • MRR (Mean Reciprocal Rank):计算真实物品排名的倒数,再取平均。

为了证明其统一建模的优势,论文除了在整体指标上对比,还应进行更细致的分析:

  1. 跨域推荐效果:可以统计当用户历史序列和待预测物品属于不同域时,模型的预测准确性。对比纯序列模型,看TokenFormer是否有提升。
  2. 序列推荐效果:在历史序列和待预测物品属于同一域时,对比性能。验证加入域令牌是否干扰了原有的序列建模能力。
  3. 冷启动场景:对于历史交互很少的用户,域令牌提供的先验信息可能更有帮助。可以按用户活跃度分组评估。

4.2 与基线模型的对比分析

论文的对比实验部分应包含以下几类基线:

  • 序列模型基线:如GRU4Rec, SASRec, BERT4Rec。用于证明TokenFormer在序列建模能力上不弱,甚至更强。
  • 多域/多任务模型基线:如MMoE, STAR, PLE。用于证明TokenFormer在利用域信息上的有效性。
  • 结合序列与特征的模型:如将域特征拼接后输入序列模型。用于证明TokenFormer“令牌化统一处理”的方式优于简单的特征拼接。

一个有力的证明是:TokenFormer在保持与顶尖序列模型(如SASRec)相当的序列推荐性能的同时,在跨域推荐和冷启动场景下显著优于它们;同时,它的性能也优于那些为多域任务专门设计的模型,尤其是在序列相关性强的场景下。

4.3 消融实验与可视化理解

消融实验是理解模型每个组件贡献度的关键:

  1. 移除域令牌:将模型退化为一个标准的物品序列Transformer。性能下降,尤其是在跨域预测上,说明域令牌的必要性。
  2. 替换域令牌为特征拼接:将域ID作为普通特征,与物品嵌入拼接后输入Transformer。性能通常不如独立的域令牌,说明“令牌化”让域信息在注意力层面有了独立的、可交互的表示,比作为特征值更有效。
  3. 注意力权重的可视化:这是最有趣的部分。可以可视化最后一层[CLS]令牌对序列中其他令牌的注意力权重。理想情况下,我们应该能看到[CLS]既关注了最近几个关键的物品令牌,也关注了某些有代表性的域令牌,从而直观地看到模型是如何综合信息的。

5. 实战思考与避坑指南

将TokenFormer的思想应用到实际业务中,远不止照搬论文结构那么简单。以下是我基于经验总结的一些实战思考和可能遇到的“坑”。

5.1 如何定义“域”?

这是落地时第一个也是最关键的问题。论文里用的“域”概念清晰(如商品大类),但实际业务中,“域”的边界可能很模糊。

  • 场景一:综合电商平台。可以按一级类目(服饰、数码、家居)划分,但用户买“手机”和买“手机壳”虽然类目不同,关联性却极强。此时,可以考虑更细的划分(如“手机及配件”作为一个域),或者引入“知识图谱”,将强关联的类目在嵌入空间拉近。
  • 场景二:内容信息流。“域”可能是内容类型(短视频、文章、图集),也可能是主题(娱乐、体育、财经)。如果按类型分,那么“娱乐类短视频”和“娱乐类文章”的域就不同,但主题相似。这时,或许可以设计双通道令牌,一个表示类型,一个表示主题,让模型同时学习两种划分下的关联。
  • 场景三:域的定义动态或稀疏。有些业务中,新的领域会不断出现,或者某些用户涉及的领域很少。这就需要域嵌入层有良好的泛化能力,可以考虑使用元学习或对域嵌入进行归一化等技巧。

实操心得:没有银弹。最好的方法是进行A/B测试。可以尝试几种不同的域划分方案,在离线评估(尤其是跨域预测和冷启动评估)后,选择效果最好的一个。也可以将“域ID”作为一个可学习的聚类中心,让模型自动从数据中学习出最有用的划分。

5.2 如何处理超长序列与计算效率?

Transformer的自注意力复杂度是序列长度的平方级。TokenFormer的序列长度因为引入域令牌而翻倍,这对长序列用户是个挑战。

  • 高效注意力机制:可以考虑使用线性注意力局部窗口注意力Reformer的LSH注意力等改进方案,来降低计算复杂度。对于推荐场景,用户最近的行为往往最重要,局部窗口注意力是一个很实用的选择。
  • 层次化建模:另一种思路是,先在一个较短的“会话”粒度(例如一次登录期间的行为)内部使用TokenFormer进行精细建模,得到会话表征,再将多个会话表征输入一个更上层的序列模型(如RNN或更轻量的Transformer)进行融合。这既控制了计算量,又保留了长周期兴趣。

5.3 负采样策略与大规模物品库

在训练时,计算所有物品的softmax在物品库巨大(百万甚至千万级)时是不现实的。通用的做法是使用采样softmax负采样

  • 随机负采样:最简单,但可能不够“硬”,模型学起来太容易。
  • Batch内负采样:将同一批次内其他用户的正样本作为当前用户的负样本。效率高,能提供一定难度的负例。
  • 重要性采样:根据物品的流行度进行采样,流行物品被采为负例的概率更高,这符合实际情况(用户更可能避开不感兴趣的流行商品)。
  • 混合负采样:结合多种策略。例如,一部分随机负样本,一部分是“困难负样本”(通过一个初步模型筛选出的、得分较高的非正例物品)。

对于TokenFormer,由于引入了域信息,在负采样时也可以做一些设计。例如,在训练跨域预测时,可以有意识地多采样一些与目标物品同域但不同的负例,以及不同域的负例,让模型更好地学习域内判别和跨域关联。

5.4 在线服务与性能优化

模型离线效果好,上线还要过性能关。

  • 模型轻量化:可以通过知识蒸馏,将大的TokenFormer模型学到的知识迁移到一个更小的模型(如双塔DNN)中,用于线上召回。或者对Transformer层进行剪枝、量化。
  • 向量化检索:线上推理时,我们可以将用户最近的序列输入TokenFormer,得到[CLS]向量作为用户实时兴趣向量。然后使用高效的向量检索库(如Faiss),从海量物品向量中快速检索Top-K候选。这里的物品向量就是训练好的物品嵌入。这要求物品嵌入离线更新及时,且与用户向量在同一空间。
  • 缓存策略:对于活跃用户,其兴趣向量可以短期缓存,避免每次请求都进行完整的Transformer前向计算。当用户有新的交互行为时,再更新缓存。

5.5 一个简单的代码框架示意

虽然无法给出完整代码,但以下PyTorch风格的伪代码框架可以帮助理解核心实现:

import torch import torch.nn as nn class TokenFormer(nn.Module): def __init__(self, num_items, num_fields, max_len, d_model, nhead, num_layers): super().__init__() self.item_embed = nn.Embedding(num_items, d_model) self.field_embed = nn.Embedding(num_fields, d_model) self.cls_token = nn.Parameter(torch.randn(1, 1, d_model)) # [CLS]令牌 self.pos_embed = nn.Embedding(2*max_len + 1, d_model) # 可学习位置编码 encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, batch_first=True) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) # 输出层:通常与item_embed共享权重 self.output_layer = nn.Linear(d_model, d_model) def forward(self, item_seq, field_seq): """ item_seq: [batch_size, seq_len] field_seq: [batch_size, seq_len] """ batch_size, seq_len = item_seq.shape # 1. 获取嵌入 item_emb = self.item_embed(item_seq) # [B, L, D] field_emb = self.field_embed(field_seq) # [B, L, D] # 2. 交错拼接物品和域令牌,并在开头加上[CLS] # 重塑:[B, L, D] -> [B, L*2, D] interleaved = torch.stack([item_emb, field_emb], dim=2).flatten(1, 2) cls_tokens = self.cls_token.expand(batch_size, -1, -1) tokens = torch.cat([cls_tokens, interleaved], dim=1) # [B, 1+2L, D] # 3. 添加位置编码 positions = torch.arange(0, 1+2*seq_len).expand(batch_size, -1).to(item_seq.device) pos_emb = self.pos_embed(positions) tokens = tokens + pos_emb # 4. 创建因果注意力掩码 seq_mask = torch.triu(torch.ones(1+2*seq_len, 1+2*seq_len), diagonal=1).bool() # 5. 通过Transformer encoded = self.transformer(tokens, mask=seq_mask) # [B, 1+2L, D] # 6. 取[CLS]位置输出 cls_output = encoded[:, 0, :] # [B, D] # 7. 预测分数 (示例:与所有物品嵌入计算内积) # 实际训练中,这里会接采样softmax损失 item_emb_all = self.item_embed.weight # [num_items, D] scores = torch.matmul(self.output_layer(cls_output), item_emb_all.T) # [B, num_items] return scores

这个框架省略了Dropout、层归一化细节以及训练循环,但清晰地展示了TokenFormer从输入到输出的数据流动过程。在实际实现时,需要特别注意位置编码的设计和注意力掩码的生成,确保符合因果建模的要求。

← 返回列表