MiniOneRec 框架解析
MiniOneRec 是一个基于大语言模型的开源推荐框架,主要包含 SID 构建(RQVAE)、SFT(0.5b~7b)和 RL(受限解码 + 混合奖励)三个核心阶段。实验证明了 LLM 的 Scaling 特点,模型越大效果越好,且对齐 SID-text 对性能提升至关重要。
代码结构
SFT 训练流程
前置:数据集
- SidSFTDataset: 输入用户最近交互过的 item 列表,预测下一个交互 item。
- SidItemFeatDataset: sid2title 或 title2sid 任务。
- FusionSeqRecDataset: 带意图识别的商品推荐。
TokenExtender: SID
SID 索引示例存储在 JSON 文件中,例如 Industrial_Scientific.index.json。indices 是一个字典,如 indices["0"]=["<a_236>", "<b_231>", "<c_226>"]。函数 get_new_tokens 返回 SID 的子 token 集合。注意代码逻辑中读取的是所有实体的 SID token 集合,而非码本 token 集合,部分未见过实体可能分配给未使用的 SID 索引 token。
只对新 token 的 embedding 做训练
通过冻结所有参数,打开 embedding 参数的 requires_grad,并利用 register_hook 控制梯度保留。直接对张量切片设置 requires_grad 属性不可行,需使用 hook 将旧 vocab 部分的梯度置零。
print("Freezing LLM parameters, only training new token embeddings")
for param in model.parameters():
param.requires_grad = False
if sid_index_path and os.path.exists(sid_index_path) and new_tokens:
embedding_layer = model.get_input_embeddings()
if embedding_layer.weight.shape[0] > original_vocab_size:
embedding_layer.weight.requires_grad = True
def mask_grad(grad):
grad[:original_vocab_size].zero_()
return grad
embedding_layer.weight.register_hook(mask_grad)
print(f"Unfrozen {len(new_tokens)} new token embeddings (indices {original_vocab_size} to {len(tokenizer)-1})")
RL 训练流程
数据集
- SidDataset: 输入用户历史交互 item SID 序列,预测下一个 item SID。
- RLTitle2SidDataset: title2sid 和 description2sid 任务,返回 prompt-completion pairs。
- RLSeqTitle2SidDataset: 输入用户交互序列(带 title)预测下一个 item SID。
奖励模型
使用 SASRec 作为奖励模型基础。
建模细节
框架组成
- Tokenizer (RQVAE)
- LLM-text 对齐(利用 LLM 的世界知识)
- SFT next token prediction
- RL (GRPO)
Task
序列推荐任务。先分词:用户 u 有历史交互商品序列 $H_u=[i_1, i_2, ..., i_T]$。每个商品 $i_t$ 通过 RQVAE 编码为 3 层 SID ${ c_0^{i_t}, c_1^{i_t}, c_2^{i_t} \u007d$。后训练 LLM $π_θ$ 读取历史序列预测下一个商品。推理时使用 k beams search。
Item Tokenization
采用标准 RQVAE。为避免码本坍缩,使用第一个训练 batch 的 k-means 中心作为 codebook 初始化码本。代码中包含 Sinkhorn-Knopp algorithm 代码,这也是缓解码本坍缩的 trick。
Align with LLM
对齐 LLM 世界知识和 SID 信号。包括序列推荐任务和对齐 SID 与文本描述任务。GitHub 仓库中更新了新技术 GPR-inspired SFT with Value-Aware Fine-Tuning (VAFT),基于模拟物品价值实现加权损失,但暂无实验结果对比。
RL with verifiable rewards (RLVR)
- 混合动态采样: SID 空间小,容易采样到相同 SID。最终采用 beam search 方法。
- 稀疏排序信号: 使用 NDCG 作为奖励,如果是正确商品分数再加 1。
训练配置
- 商品描述: Qwen3-Embedding-4B 编码
- 分词器: RQVAE,单卡训练,batchsize=20480,lr=1e-3,epochs=10000
- SFT: AdamW,Qwen2.5-Instruct。8 卡训练,单卡 batchsize=128, 10 epochs+early stop(patience=1),lr=3e-4,cosine decay。
- RL: GRPO,2 epoch,KL 权重 β=0.1,lr=1e-5,batchsize=512
- 推理: beam search width=16
评估
使用亚马逊数据集 Office+Industrial,指标为 hitrate+NDCG。
- Scaling: 训练 + 评估损失随模型增大而减小。
- Baseline 对比: LLM 系列和非 LLM 系列对比说明世界知识重要性;Ours 和 LLM 系列对比说明 RL 的重要性。
Transferability
在 Industrial 上训练并在 Office 上评估,证明 RL 的有效性。未做 SFT 是因为 SFT 容易领域过拟合影响迁移。
消融实验
- Language-SID 重要性: 对比不做对齐、仅 RL 对齐、仅 SFT 对齐等不同配置。
- 采样策略: 直接 topk、采 1.5 倍 budget+ 筛选、beamsearch(效果最好)。
- 奖励设计: 01 奖励、SASRec 模型 logits(效果差,存在 reward hacking)、NDCG。
- 是否预训练: 预训练效果更好。
相关代码
sft_gpr 模块实现了 GPR-inspired SFT with Value-Aware Fine-Tuning (VAFT),基于模拟物品价值实现加权损失。
