MiniOneRec 论文解读:生成式推荐框架与代码分析
概述
MiniOneRec 是一个开源的生成式推荐框架,支持从 0.5B 到 7B 的模型扩展。核心流程包括 SID 构建(RQVAE)、SFT(0.5b~7b)以及 RL(受限解码 + 混合奖励)。实验表明模型规模越大效果越好,且将 SID 与文本对齐至关重要。
代码结构
主要入口文件位于 MiniOneRec/sft.py 和 MiniOneRec/rl.py。
SFT 原理
数据集
- SidSFTDataset: 输入用户最近交互 item 列表,预测下一个交互 item。
- SidItemFeatDataset: 执行 sid2title 或 title2sid 任务。
- FusionSeqRecDataset: 带意图识别的商品推荐。
训练细节
- 传入主干 LLM 即 base_model。
- SID token 被映射到普通 tokens 中,而非 special tokens。
- 可选择冻结 LLM 自带参数,仅训练新传入的 SID token embedding。
- 处理上述三种类型数据集并进行训练。
SID 索引扩展
SID 索引示例存储在 JSON 文件中,格式为 {index: [token_list]}。函数 get_new_tokens 返回 SID 的子 token 集合。注意代码逻辑可能读取所有实体的 SID token 集合,而非码本的实际利用集合,部分未见过的实体可能被分配至未使用的 SID 索引 token。
新 Token Embedding 训练策略
通过冻结所有参数并设置 embedding 层权重梯度,配合 hook 控制保留特定 token 的梯度。
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 to )")

