Trie-Aware Transformers for Generative Recommendation¶
会议: NeurIPS 2026
arXiv: 2602.21677
OpenReview: https://openreview.net/forum?id=oxoJ8MqY3p
论文: https://arxiv.org/abs/2602.21677
领域: 推荐系统
关键词: 生成式推荐 / 前缀树 / 位置编码 / 自回归模型 / 层次量化
一句话总结¶
针对现有生成式推荐模型将层次语义 ID 展平为一维序列而忽略前缀树拓扑的缺陷,本文提出模型无关且零超参的 TrieRec,通过结构感知绝对位置编码(TAE)聚合前缀路径消除语义歧义,并借助拓扑感知相对位置编码(TRE)注入基于最近公共祖先的注意力偏置,在三大主流骨干模型上取得平均 8.83% 的相对增益。
研究背景与动机¶
生成式推荐(Generative Recommendation, GR)顺应生成式人工智能的发展浪潮,将传统的候选排序评分任务重构为端到端的自回归 token 生成任务。典型的生成式推荐包含两阶段流程:首先是物品离散化(Item Tokenization),通过残差量化变分自编码器(如 RQ-VAE)或层次聚类将物品嵌入由粗到细编码为离散语义标识符序列;其次是自回归生成,利用 Transformer 架构以用户历史交互物品的 token 流为条件,逐 token 生成目标物品的标识符序列。由于层次量化具有从粗粒度到细粒度的先验,全部物品标识符在逻辑上自然构成一棵层次前缀树(Trie)。在此结构下,预测下一物品等价于在 Trie 上逐层选择子节点,不仅大幅缩小了每一步的搜索候选空间,更将传统级联推荐系统的漏斗筛选机制优雅地内化于单一生成模型中。
然而,现有的生成式推荐方法在自回归建模阶段存在一个根本性的设计缺陷:它们普遍将多层级 token 序列直接打平为扁平的线性一维流,仅仅依赖标准的自注意力机制和一维线性序列位置编码进行建模。这种扁平化处理完全丢弃了前缀树中蕴含的两类极其关键的结构归纳偏置:一是绝对结构位置信息,由于 RQ-VAE 的 codebook 索引本质上是量化残差编号而非自然语言中具备固定语义的词汇,同一个 token 在不同物品间频繁复用,仅看单点 token 极易产生严重的语义歧义,必须结合其在 Trie 中的深度与祖先-后代路径上下文才能精准锚定其语义;二是 token 对之间的拓扑关联信息,属于同一物品路径的 token 具有强依赖,而共享更近最近公共祖先(LCA)的 token 往往属于同类或相似物品,其相关程度显著高于处于遥远分支的 token 对。
尽管前缀树是语义编码的核心组织形态,但在自回归生成阶段显式利用 Trie 拓扑的研究仍极度匮乏。少数尝试引入树结构先验的工作(如 SEATER)主要通过在目标函数中添加对比损失来实现,这不仅需要为树中海量节点额外维护和更新嵌入表,带来高昂的显存与计算开销,而且仅约束局部的父子关系,未能捕捉高阶相对拓扑。核心 idea:显式将物品前缀树(Trie)的层次拓扑先验注入 Transformer 自回归模型中,设计结构感知的绝对位置编码(TAE)聚合前缀路径的层级与上下文信息以消除 token 语义歧义,并构建拓扑感知的相对位置编码(TRE)基于最近公共祖先(LCA)及树距离学习注意力偏置,以极低计算开销实现模型无关的结构化生成式推荐。
方法详解¶
整体框架¶
TrieRec 的核心理念是在不改动底层物品分词机制与自回归主干架构的前提下,通过对 Transformer 注入两套互补的结构先验来赋能生成式推荐:其一是输入端的树结构感知绝对位置编码(TAE),它显式融合物品路径上的纵向祖先与后代节点表征,使每个 token 感知其所处的绝对层级与局部上下文;其二是自注意力机制内部的拓扑感知相对位置编码(TRE),它针对序列中任意 token 对提取最近公共祖先及其树上距离三元组,转化为可学习的注意力偏置,精准引导注意力权重按拓扑亲疏度进行分配。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["用户交互物品序列<br/>X = [c_1, ..., c_T]"] --> B["Token 语义嵌入提取<br/>e_{i,l} in R^d"]
B --> C["结构感知绝对位置编码 (TAE)<br/>同路径上下文拼接 + MLP聚合 + 层级专属投影"]
C --> D["注入绝对位置先验<br/>e_hat_{i,l} = e_{i,l} + p_{i,l}"]
D --> E["拓扑感知相对位置编码 (TRE)<br/>计算成对 (d^LCA, delta_u, delta_v) 三元组"]
E --> F["拓扑偏置注意力计算<br/>Softmax(QK^T / sqrt(d) + B(d^LCA, delta_u, delta_v))"]
F --> G["Transformer 多层特征变换与自回归解码"]
G --> H["前缀树受限束搜索<br/>逐层生成目标物品标识符"]
关键设计¶
1. 结构感知绝对位置编码(TAE):多层级纵向上下文拼接与层级专属投影 针对离散量化 token 语义歧义以及一维线性位置编码无法反映树深与隶属路径的痛点,TAE 显式收集该 token 所在物品从根到叶的完整纵向路径邻居(包含所有祖先与后代节点),将同一物品路径上的 \(L\) 个 token 嵌入向量拼接为一个联合上下文表征 \(\mathbf{x}_i^{con} = [\mathbf{e}_{i,1}^\top; \mathbf{e}_{i,2}^\top; \dots; \mathbf{e}_{i,L}^\top]^\top \in \mathbb{R}^{L \cdot d}\)。随后将其送入两层多层感知机(MLP,配备 GELU 激活函数与层归一化)进行信息融合,输出聚合后的多层级交互表征 \(\hat{\mathbf{x}}_i \in \mathbb{R}^{L \cdot d}\)。与具有置换不变性的自注意力机制不同,MLP 的权重与维度位置严格绑定,能够针对不同层级深度显式分配独立的交互参数,从而实现跨 token 与跨维度的细粒度结构建模。最后,为了适配不同层级 token 独特的语义粒度,TAE 使用层级专属的线性投影参数 \(\mathbf{W}_l \in \mathbb{R}^{d \times (L \cdot d)}\) 和偏置 \(\mathbf{b}_l\) 生成该深度的绝对位置编码: $\(\mathbf{p}_{i,l} = \mathbf{W}_l \hat{\mathbf{x}}_i + \mathbf{b}_l, \quad l=1,\dots,L\)$ 该编码被直接加和到第一层 Transformer 输入的 token 嵌入上:\(\hat{\mathbf{e}}_{i,l} = \mathbf{e}_{i,l} + \mathbf{p}_{i,l}\)。为平衡效果与计算效率,TAE 仅在第一层注入且仅作用于已完整给出的历史交互物品,生成阶段尚不具备完整后代路径的预测 token 则不加 TAE,确保了自回归推理的因果严格性与极低耗时。
2. 拓扑感知相对位置编码(TRE):基于最近公共祖先三元组的结构注意力偏置 标准自注意力机制仅依赖 Query-Key 内积计算相关性,缺乏对两两 token 在前缀树中所处拓扑关系的结构感知。TRE 通过对序列内任意 token 对 \((u, v)\) 提取三个关键拓扑特征,全面量化其层级结构亲疏度:① 最近公共祖先的树深度 \(d^{LCA}_{u,v}\),该值直接反映了两节点所属物品在 Trie 中共享前缀的长度,值越大表明两物品在粗粒度类目或属性上越接近;② 节点 \(u\) 到最近公共祖先的垂直距离 \(\delta_u\);③ 节点 \(v\) 到最近公共祖先的垂直距离 \(\delta_v\)。三元组 \((d^{LCA}_{u,v}, \delta_u, \delta_v)\) 不仅给出了两节点间的树上最短路径长度 \(\delta_u + \delta_v\),而且通过比对 \(\delta_u\) 与 \(\delta_v\) 能够明确二者的相对层级高低(例如判断谁更偏向粗粒度或细粒度属性)。基于该三元组,TRE 维护一个紧凑的可学习注意力偏置张量 \(B(d^{LCA}_{u,v}, \delta_u, \delta_v)\),并将其直接叠加在原始缩放点积注意力得分上: $\(\tilde{A}_{u,v} = \text{Softmax}_v \left( \frac{(\mathbf{e}_u \mathbf{W}^Q)(\mathbf{e}_v \mathbf{W}^K)^\top}{\sqrt{d}} + B(d^{LCA}_{u,v}, \delta_u, \delta_v) \right)\)$ 偏置表 \(B\) 在 Transformer 的所有注意力层间完全共享,因此全序列成对拓扑特征只需预先计算一次并查表复用,无需随网络加深重复计算,避免了深层网络的额外开销。
3. 轻量高效与模型无关的工程架构:结构缓存与零超参即插即用 不同于 SEATER 等在损失函数层面引入全节点嵌入对比训练的方法,TrieRec 具备三大突出特性:① 模型无关(Model-Agnostic),其设计完全解耦于具体的离散分词方式与推荐骨干架构,可无缝嵌入 TIGER(标准 RQ-VAE)、CoST(对比量化)以及 LETTER(协同正则化量化)等主流框架;② 极低复杂度与缓存加速,对长度为 \(T\) 个物品、每物品 \(L\) 层的序列,TAE 的序列计算复杂度为 \(O(T L^2 d^2)\),TRE 的成对 LCA 计算仅有 \(O(h T^2 L^2 \log L)\) 且可借助哈希缓存去重加速,偏置查找为 \(O(1)\),与 Transformer 自身 \(O(k h T^2 L^2 d)\) 的自注意力复杂度相比量级极小,实测额外时间开销不足 10%;③ 零额外超参数(Hyperparameter-Free),整个框架无需调整任何损失平衡权重或结构衰减系数,避免了繁重的超参网格搜索。
损失函数 / 训练策略¶
TrieRec 继承标准生成式推荐的端到端自回归极大似然训练范式,无需引入任何辅助正则化损失。给定目标物品 token 序列 \(\mathbf{y} = [y_1, y_2, \dots, y_L]\),模型直接最小化交叉熵自回归生成损失: $\(\mathcal{L} = -\sum_{t=1}^{L} \log P_{\mathcal{M}}(y_t \mid y_{<t}, \mathbf{X})\)$ 在推理阶段,模型采用带前缀树约束的束搜索(Constrained Beam Search),仅在 Trie 允许的合法前缀子节点空间中扩展,从而保证生成的 token 序列严格对应商品库中的真实有效物品,束宽统一设为 \(K=10\)。
实验关键数据¶
主实验¶
评估基于四个公开基准数据集:Amazon-Beauty、Amazon-Toys、Amazon-Food 和 MovieLens-1M,采用留一法评估,评价指标为 NDCG@K 和 Recall@K(\(K \in \{5, 10\}\))。下表展示了以经典及当前最先进骨干模型(TIGER、CoST、LETTER)为基础,集成 TrieRec 后的主要对比结果(基线模型摘录自原论文表 2):
| 骨干与模型 | Beauty N@5 | Beauty R@10 | Toys N@5 | Toys R@10 | Food N@5 | Food R@10 | ML-1M N@5 | ML-1M R@10 |
|---|---|---|---|---|---|---|---|---|
| SASRec | 0.0177 | 0.0612 | 0.0229 | 0.0739 | 0.0210 | 0.0716 | 0.0440 | 0.1391 |
| BERT4Rec | 0.0196 | 0.0545 | 0.0207 | 0.0507 | 0.0245 | 0.0678 | 0.0420 | 0.1232 |
| TIGER (基线) | 0.0318 | 0.0719 | 0.0332 | 0.0765 | 0.0337 | 0.0792 | 0.0713 | 0.1713 |
| TIGER + TrieRec | 0.0335 | 0.0723 | 0.0368 | 0.0766 | 0.0368 | 0.0843 | 0.0841 | 0.1929 |
| (相对提升) | (+5.35%) | (+0.56%) | (+10.84%) | (+0.13%) | (+9.20%) | (+6.44%) | (+17.91%) | (+12.61%) |
| CoST (基线) | 0.0322 | 0.0703 | 0.0331 | 0.0756 | 0.0335 | 0.0811 | 0.0719 | 0.1730 |
| CoST + TrieRec | 0.0337 | 0.0720 | 0.0391 | 0.0771 | 0.0367 | 0.0851 | 0.0847 | 0.1909 |
| (相对提升) | (+4.66%) | (+2.42%) | (+18.13%) | (+1.98%) | (+9.55%) | (+4.93%) | (+17.80%) | (+10.35%) |
| LETTER (基线) | 0.0345 | 0.0765 | 0.0347 | 0.0781 | 0.0341 | 0.0827 | 0.0779 | 0.1858 |
| LETTER + TrieRec | 0.0385 | 0.0794 | 0.0403 | 0.0801 | 0.0373 | 0.0861 | 0.0914 | 0.2031 |
| (相对提升) | (+11.59%) | (+3.79%) | (+16.14%) | (+2.56%) | (+9.38%) | (+4.11%) | (+17.33%) | (+9.31%) |
在横向对比其他专有位置编码及树模型方面,TrieRec 显著超越了代码领域树编码迁移方法(HiT、HiRoPE)以及粗粒度推荐编码(LayerPE、DiscRec)。与同为树感知方法的 SEATER(基于树对比学习损失)对比,在相同骨干(LETTER)下,TrieRec 在 ML-1M 上的 NDCG@5 达 0.0914(对比 SEATER 的 0.0726,提升 25.9%),充分证明直接注入网络拓扑表征远优于外挂辅助损失。
消融实验¶
在 LETTER 骨干上对各核心组件进行消融分析(数据摘自原论文表 3):
| 变体配置 | Beauty N@5 | Beauty R@10 | Toys N@5 | Toys R@10 | ML-1M N@5 | ML-1M R@10 | 说明 |
|---|---|---|---|---|---|---|---|
| TrieRec (Full) | 0.0385 | 0.0794 | 0.0403 | 0.0801 | 0.0914 | 0.2031 | 完整模型 |
| (1) w/o TAE | 0.0360 | 0.0781 | 0.0383 | 0.0800 | 0.0740 | 0.1779 | 移除绝对位置编码,ML-1M N@5暴跌19.0% |
| (2) w/o TRE | 0.0363 | 0.0779 | 0.0381 | 0.0791 | 0.0847 | 0.1951 | 移除相对位置偏置,全线出现稳定退化 |
| (3) TAE-Blocked | 0.0362 | 0.0790 | 0.0375 | 0.0799 | 0.0762 | 0.1816 | 限制MLP仅在同维度跨深度交互,缺乏跨维融合 |
| (4) TAE-Prefix | 0.0343 | 0.0766 | 0.0339 | 0.0760 | 0.0734 | 0.1814 | 用简单前缀和累加替代MLP,性能劣于原始基线 |
| (5) TRE-w/o LCA | 0.0374 | 0.0787 | 0.0385 | 0.0799 | 0.0901 | 0.2007 | 去除最近公共祖先深度,仅保留垂直相对距离 |
| (6) TRE-w/o \(\delta\) | 0.0371 | 0.0786 | 0.0397 | 0.0795 | 0.0900 | 0.2017 | 去除相对垂直距离,仅保留LCA共享深度 |
关键发现¶
- TAE 的纵向路径上下文对性能起决定性支撑:完全移除 TAE(变体 1)在密集交互的 MovieLens-1M 上导致 NDCG@5 从 0.0914 骤降至 0.0740(相对跌幅达 19.0%),表明若无法结合整条路径的上下文,RQ-VAE 的残差 token 将承受严重的语义混淆。
- 简单累加不可行,非线性结构变换至关重要:若采用朴素的前缀和累加(TAE-Prefix),性能甚至低于未经任何修改的 LETTER 原型(0.0343 vs 0.0345),说明多层级量化特征之间的关联是非线性的,必须通过具备跨维度交互能力的 MLP 与深度感知专属投影才能有效解构层级语义。
- LCA 深度与树距离具有双向协同增益:TRE 中单独剔除 LCA 深度(变体 5)或剔除垂直距离(变体 6)均会导致指标受损,证明共享前缀长度所代表的宏观语义共性与垂直树距离所代表的微观相对位置缺一不可。
- 低 K 指标提升尤为显著:在各数据集上,NDCG@5 的提升幅度普遍高于 NDCG@10(如 Toys 上 TIGER+TrieRec 在 N@5 提升 10.84%,在 N@10 仅提升 6.22%),这表明树拓扑归纳偏置能极大增强模型在顶层最相关候选上的辨识度与置信度。
亮点与洞察¶
- 将前缀树由“外部解码约束”升格为“内部归纳偏置”:过去的研究普遍仅在测试阶段把 Trie 当作受限解码字典,或在损失端施加弱监督;本文首次将 Trie 的拓扑距离显式映射为 Transformer 的内部绝对表征与注意力偏置,让模型在注意力计算的第一步就内生性地理解树结构。
- 三元组 \((d^{LCA}, \delta_u, \delta_v)\) 的精妙拓扑参数化:无需像图神经网络或语法树编码那样显式遍历整条路径的嵌入,仅用标量三元组查表即可完整描述任意两节点的树亲缘关系与方向性,在极度轻量化的前提下实现了充分的结构表达。
- 高兼容性的即插即用范式:不调整分词器训练过程,不增加任何辅助目标损失,不引入额外的敏感超参数,可以直接赋能既有的各类先进生成式推荐框架。
局限与展望¶
- 依赖预构建静态 Trie 的拓扑质量:模型假设分词器构建的前缀树真实且完美地反映了物品间的层次语义关联;若 RQ-VAE 发生严重的码本坍塌或层次聚类失真,结构位置编码可能引入误导性的拓扑偏置。
- 长序列全量成对 LCA 检索的瓶颈:尽管当前依靠路径哈希缓存将开销压制在 10% 以内,但若用户历史交互物品序列极长(如数百上千项),成对注意力偏置的构建与显存占用仍将面临扩展性挑战。
- 未来方向:探索结合动态树结构的端到端自适应前缀树更新机制,并将树感知位置编码推广至多模态物品特征混合生成的更大规模跨域推荐场景中。
相关工作与启发¶
- vs TIGER / LETTER / CoST:上述经典方法关注如何通过重构、对比学习或协同信号优化 RQ-VAE 的离散码本,但在自回归生成阶段均将 token 视作扁平序列;本文正交于其分词创新,专注于自回归阶段的树拓扑注入,与其完全兼容互补。
- vs SEATER:SEATER 通过在损失函数中增加前缀共享与父子节点对比损失来引导树结构学习,模型需额外维护庞大的树节点全局表征且无法建模高阶亲缘关系;TrieRec 直接通过轻量位置编码修改网络内部注意力流,精度显著更高且显存占用更小。
- vs HiT / HiRoPE:HiT 与 HiRoPE 针对编程语言抽象语法树(AST)设计,强调函数块嵌套与变量作用域,无法直接适配物品推荐中基于残差量化建立的前缀树语义特性;TrieRec 专为生成式推荐的层次量化 Trie 定制,实测效果大幅占优。
评分¶
- 新颖性: ⭐⭐⭐⭐ [针对生成式推荐扁平化建模的固有痛点,首创性地提出基于 Trie 的绝对与相对位置双编码机制]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 4 个标准数据集、3 大主流骨干模型,提供详尽的消融、效率开销与基准对比]
- 写作质量: ⭐⭐⭐⭐⭐ [动机链条严密,方法设计清晰干练,数学公式与图表表达严谨规范]
- 价值: ⭐⭐⭐⭐⭐ [模型无关、零超参、轻量高效,极易在工业级生成式推荐管线中无缝落地推广]