在拿下2025腾讯广告算法大赛之前我们可以做些什么——官方数据集和 Baseline 代码解读
数据集
数据集概述
TencentGR_1k/是一个腾讯广告推荐系统的数据集，包含用户行为序列、物品特征、多模态嵌入等多种数据格式。该数据集主要用于推荐系统模型的训练和评估。
目录结构
TencentGR_1k/
├── seq.jsonl # 用户行为序列数据
├── item_feat_dict.json # 物品特征字典
├── indexer.pkl # ID映射索引
├── seq_offsets.pkl # 序列偏移指针
├── README.md # 原始说明文档
└── creative_emb/ # 多模态嵌入特征
├── emb_81_32/ # 嵌入特征81，维度32
├── emb_82_1024/ # 嵌入特征82，维度1024
├── emb_83_3584/ # 嵌入特征83，维度3584
├── emb_84_4096/ # 嵌入特征84，维度4096
├── emb_85_4096/ # 嵌入特征85，维度4096
└── emb_86_3584/ # 嵌入特征86，维度3584
核心数据文件详解
1. seq.jsonl - 用户行为序列数据
文件格式: JSON Lines格式，每行一个用户的行为序列
数据样例:
[
# item interaction
[480, 53961, null, {"112": 19, "117": 103, "118": 125, "119": 87, "120": 126, "100": 2, "101": 40, "102": 8559, "122": 5998, "114": 16, "116": 1, "121": 52176, "111": 5630}, 0, 1746077791],
[480, 50654, null, {"112": 15, "117": 285, "118": 737, "119": 1500, "120": 1071, "100": 6, "101": 22, "102": 10420, "122": 2269, "114": 16, "115": 43, "116": 13, "121": 12734, "111": 6737}, 0, 1746094091],
[480, 23149, null, {"112": 15, "117": 84, "118": 774, "119": 1668, "120": 348, "100": 6, "101": 32, "102": 6372, "122": 2980, "114": 16, "116": 15, "121": 30438, "111": 34195}, 0, 1746225104],
...
]
数据格式说明:
-
每一行为一个用户的行为序列，按时间排序
-
用户序列中每一个record的数据格式为：
[user_id, item_id, user_feature, item_feature, action_type, timestamp]
字段详解:
-
user_id: 用户ID（已重新编号） -
item_id: 物品ID（已重新编号） -
user_feature: 用户特征（dict格式，key为feature_id，value为feature_value） -
item_feature: 物品特征（dict格式，key为feature_id，value为feature_value） -
action_type: 行为类型-
0 = 曝光（看到但未点击）
-
1 = 点击（看到并点击）
-
-
timestamp: 时间戳
特殊说明:
-
每一条record为记录
user profile和item interaction其中的一个。 -
若当前record为
user profile，则item_id、item_feature、action_type为null -
若当前record为
item interaction，则user_feature为null -
所有ID都从1开始重新编号（re-id），便于模型embedding lookup。
-
原始值到re-id的映射参考下方的indexer.pkl文件。
2. indexer.pkl - ID映射索引
文件格式: Python pickle格式
说明: 记录了原始wuid(website-uid), creative_id和feature_value从1开始编号后的值（以下记为re-id）
用途: 记录原始ID到重新编号后的映射关系，便于模型训练时的embedding查找。
indexer['u'] # 记录了wuid到re-id, key为wuid, value为re-id
indexer['i'] # 记录了creative_id到re-id, key为creative_id, value为re-id
indexer['f'] # 记录了各feature_id中feature_value到re-id，例如：⬇️
indexer['f']['112'] # 记录了feature_id 112中的feature_value到re-id的映射
数据结构:
indexer = {
'u': {wuid: re_id}, # 用户ID映射
'i': {creative_id: re_id}, # 物品ID映射
'f': { # 特征值映射
'112': {feature_value: re_id},
'113': {feature_value: re_id},
# ... 其他特征ID
}
}
数据样例:
# data.keys()
dict_keys(['f', 'i', 'u', 'a'])
# data['u']
{'user_01060923': 1,
'user_00300038': 2,
'user_00122945': 3,
'user_00311455': 4,
'user_00427257': 5,
...}
# data['i']
{20002650278: 1,
20002695270: 2,
20005426854: 3,
20004710150: 4,
20004532262: 5,
...}
# data['f'].keys()
dict_keys(['122', '102', '121', '118', '111', '119', '103', '120', '117', '115', '107', '114', '101', '108', '112', '109', '110', '116', '106', '105', '104', '100'])
# data['f']['122']
{1220085321: 1,
1220026889: 2,
1220093353: 3,
1220088233: 4,
1220061609: 5,
...}
# data['a']
{1: 1, 0: 2}
auto. item_feat_dict.json - 物品特征字典
文件格式: JSON格式
用途: 记录训练集中所有物品的特征，方便训练时进行负采样
数据结构:
{
"item_re_id": {
"feature_id": "feature_value",
// ... 其他特征
}
}
数据样例:
{
"47086": {
"112": 14,
"117": 80,
"118": 249,
"119": 1174,
"120": 118,
"100": 6,
"101": 18,
"102": 6779,
"122": 8307,
"114": 16,
"116": 6,
"121": 28737,
"111": 42338
},
"53920": {
"112": 14,
"117": 80,
"118": 249,
"119": 305,
"120": 11,
"100": 6,
"101": 28,
"102": 10123,
"122": 5741,
"114": 4,
"116": 15,
"121": 1485
},
...
}
4. seq_offsets.pkl - 序列偏移指针
文件格式: Python pickle格式
用途: 记录seq.jsonl中每一行的起始文件指针offset，提升I/O效率，便于快速定位数据位置
数据样例:
[0, 18809, 38087, 50037, 66624, ..., 15044950, 15060592, 15079757, 15097201, 15114487]
5. creative_emb/ - 多模态嵌入特征
说明: 存储不同的多模态特征，多模态特征ID为81-86
目录结构: 包含6个不同维度的嵌入特征文件夹
嵌入特征详情:
-
emb_81_32/: 嵌入特征81，维度32 -
emb_82_1024/: 嵌入特征82，维度1024 -
emb_83_3584/: 嵌入特征83，维度3584 -
emb_84_4096/: 嵌入特征84，维度4096 -
emb_85_4096/: 嵌入特征85，维度4096 -
emb_86_3584/: 嵌入特征86，维度3584
文件格式: 每个文件夹包含多个part文件，采用Spark输出格式
数据样例:
{"anonymous_cid": 20001973219, "emb": [0.05808627977967262, -0.10434802621603012, ...]}
字段说明:
-
anonymous_cid: 匿名化的创意ID（对应item_id） -
emb: 嵌入向量，维度根据文件夹名称确定
Baseline 代码
项目概述
这是一个基于Transformer架构的序列推荐系统baseline实现，主要用于处理用户行为序列和多模态特征数据。项目结合了注意力机制、多模态特征融合和向量量化技术，为推荐系统提供了一个完整的训练和推理框架。
项目结构
baseline/
├── main.py # 主训练脚本
├── dataset.py # 数据集处理模块
├── model.py # 基础模型实现
├── model_rqxvae.py # RQ-VAE向量量化模型
├── infer.py # 推理脚本
└── run.sh # 运行脚本
Pipeline
1. 数据集处理 (dataset.py)
1.1 数据加载
# 用户序列数据格式: [(user_id, item_id, user_feat, item_feat, action_type, timestamp)]
user_sequence = self._load_user_data(uid) # 动态加载用户数据
# 扩展用户序列，插入用户token和物品token
ext_user_sequence = []
for record_tuple in user_sequence:
u, i, user_feat, item_feat, action_type, _ = record_tuple
if u and user_feat:
ext_user_sequence.insert(0, (u, user_feat, 2, action_type)) # 用户token
if i and item_feat:
ext_user_sequence.append((i, item_feat, 1, action_type)) # 物品token
# 输出序列类似 [用户A, 物品1, 物品2, 物品3]
1.2 序列填充和负采样
一些概念：
处理流程: 首先对短序列进行左填充到固定长度，然后从后往前遍历用户序列，为每个位置生成对应的ID、类型和特征信息。当预测下一个物品时，会生成正样本（真实下一个物品）和负样本（随机未交互物品）用于训练。
核心变量:
-
seq/pos/neg: 存储序列ID、正样本ID、负样本ID，每个位置对应一个用户或物品的ID -
token_type/next_token_type: 标识当前和下一个token的类型（1=物品，2=用户），用于区分序列中的用户token和物品token -
seq_feat/pos_feat/neg_feat: 存储对应的特征字典，每个位置包含该用户或物品的所有特征信息（稀疏特征、多模态特征等）
训练目标: 模型学习预测用户序列中每个位置的下一个物品，通过对比正负样本提高推荐准确性。
数据示例:
# 假设用户序列: [用户A, 物品1, 物品2, 物品3]
seq = [0, 用户A, 物品1, 物品2, 物品3] # 序列ID (0为填充)
pos = [0, 物品1, 物品2, 物品3, 0] # 下一个正样本ID
neg = [0, 随机物品, 随机物品, 随机物品, 0] # 下一个负样本ID
token_type = [0, 2, 1, 1, 1] # 1=物品, 2=用户
next_token_type = [0, 1, 1, 1, 0] # 下一个token类型
# 特征字典示例 (每个位置存储一个字典)
seq_feat[1] = {'103': 123, '104': 456, '106': [1,2,3]} # 用户A的特征
seq_feat[2] = {'100': 789, '117': 321, '81': [0.1,0.2]} # 物品1的特征
pos_feat[1] = {'100': 789, '117': 321, '81': [0.1,0.2]} # 物品1特征(正样本)
neg_feat[1] = {'100': 999, '117': 888, '81': [0.3,0.4]} # 随机物品特征(负样本)
1.3 特征处理
feat_types = {
'user_sparse': ['103', '104', '105', '109'], # 用户稀疏特征 - 用户ID、性别、年龄等离散特征
'item_sparse': ['100', '117', '111', '118', '101', '102', '119', '120', '114', '112', '121', '115', '122', '116'], # 物品稀疏特征 - 物品ID、类别、品牌等离散特征
'user_array': ['106', '107', '108', '110'], # 用户数组特征 - 用户兴趣标签、行为序列等数组特征
'item_array': [], # 物品数组特征 - 物品标签、属性列表等数组特征
'item_emb': self.mm_emb_ids, # 物品多模态特征 - 物品内容的预训练embedding向量 (默认['81'], 可选['81', '82', '83', '84', '85', '86'])
'user_continual': [], # 用户连续特征 - 用户活跃度、注册时长等数值特征
'item_continual': [] # 物品连续特征 - 物品热度、价格等数值特征
}
特征类型说明:
-
稀疏特征(sparse): 离散的分类特征，如ID、性别、类别等，需要embedding编码
-
数组特征(array): 多值的列表特征，如兴趣标签、属性列表等，需要特殊处理
-
多模态特征(emb): 预训练的embedding向量，如用户画像、物品内容等，直接使用
-
连续特征(continual): 数值型特征，如活跃度、价格等，直接作为数值使用
2. 模型架构 (model.py)
2.1 模型初始化
class BaselineModel(torch.nn.Module):
def __init__(self, user_num, item_num, feat_statistics, feat_types, args):
# 核心组件:
self.item_emb = nn.Embedding(item_num + 1, hidden_units) # [item_num+1, hidden_units] - 物品ID的embedding表，将物品ID映射为向量
self.user_emb = nn.Embedding(user_num + 1, hidden_units) # [user_num+1, hidden_units] - 用户ID的embedding表，将用户ID映射为向量
self.pos_emb = nn.Embedding(maxlen + 1, hidden_units) # [102, hidden_units] - 位置编码，为序列中的每个位置生成位置向量
self.sparse_emb = {} # 稀疏特征embedding表 - 存储所有稀疏特征(如性别、年龄、类别等)的embedding层
self.emb_transform = {} # 多模态特征线性变换 - 将预训练的多模态embedding转换为模型所需的维度
2.2 特征转Embedding（feat2emb函数）
核心功能: 将用户和物品的各种特征统一转换为向量表示。
关键函数:
def feat2emb(self, seq, feature_array, mask=None, include_user=False):
"""
将特征转换为embedding
seq中可能同时包含用户ID和物品ID，feat2emb会根据mask对用户ID和物品ID进行分别处理，这部分最好看具体代码
Args:
seq: [batch_size, seq_len] - 序列ID (用户ID或物品ID)
feature_array: [batch_size, seq_len] - 特征数组 (每个位置是一个特征字典)
mask: [batch_size, seq_len] - token类型掩码(1=item, 2=user)
Returns:
embs: [batch_size, seq_len, hidden_units] - 特征embedding
"""
关键变量详解:
-
seq: 序列ID，存储用户ID或物品ID，用于获取基础embedding -
feature_array: 特征字典数组，每个位置包含该用户/物品的所有特征信息- 例如:
{'103': 123, '104': 456, '81': [0.1,0.2]}(用户ID、性别、多模态特征等)
- 例如:
-
mask: token类型掩码，1=物品token，2=用户token，用于区分处理方式 -
include_user: 是否包含用户特征，推理时设为False
特征类型详解:
-
稀疏特征: 如用户ID(103)、性别(104)、物品类别(100)等离散特征
-
数组特征: 如用户兴趣标签(106)、物品属性列表等数组特征
-
多模态特征: 如物品内容embedding(81)等预训练向量特征
-
连续特征: 如用户活跃度、物品热度等数值特征
2.3 序列编码（log2feats函数）
核心功能: 将用户序列转换为统一的特征表示，通过Transformer捕获序列中的时序依赖关系。
编码过程说明:
-
特征转换: 调用 2.2 特征转Embedding 中
feat2emb函数将序列中的每个token（用户ID或物品ID）及其特征转换为向量表示 -
位置编码: 为序列中的每个位置添加位置信息，帮助模型理解token在序列中的位置
-
Transformer编码: 通过多层Transformer结构捕获序列中的长距离依赖关系，学习用户行为模式
-
输出: 返回每个位置的最终特征表示，用于后续的推荐任务
编码流程详解:
def log2feats(self, log_seqs, mask, seq_feature):
"""
将用户序列编码为特征表示
Args:
log_seqs: [batch_size, seq_len] - 用户序列ID (用户ID和物品ID的混合序列)
mask: [batch_size, seq_len] - token类型掩码(1=item, 2=user)
seq_feature: [batch_size, seq_len] - 序列特征 (每个位置的特征字典)
Returns:
log_feats: [batch_size, seq_len, hidden_units] - 序列特征表示
"""
# 1. 特征转embedding: 将序列中的每个token转换为向量表示
seqs = self.feat2emb(log_seqs, seq_feature, mask) # [batch_size, seq_len, hidden_units]
# 2. 位置编码: 为序列中的每个位置添加位置信息
positions = torch.arange(log_seqs.size(1), device=self.dev) # [seq_len]
seqs += self.pos_emb(positions) # 添加位置编码
# 3. Transformer编码: 通过多层Transformer捕获序列依赖关系
for i in range(len(self.attention_layers)):
# Multi-head attention: 计算注意力权重
mha_outputs, _ = self.attention_layers[i](seqs, seqs, seqs, attn_mask=attention_mask)
# Layer normalization + residual connection
seqs = self.attention_layernorms[i](seqs + mha_outputs)
# Feed forward + Layer normalization + residual connection
seqs = self.forward_layernorms[i](seqs + self.forward_layers[i](seqs))
# 4. 最终输出: 经过所有Transformer层后的序列表示
return self.last_layernorm(seqs)
3. 训练过程 (main.py)
3.1 训练循环
for epoch in range(epoch_start_idx, args.num_epochs + 1):
for step, batch in tqdm(enumerate(train_loader)):
# batch shapes:
seq = seq.to(args.device) # [batch_size, maxlen+1]
pos = pos.to(args.device) # [batch_size, maxlen+1]
neg = neg.to(args.device) # [batch_size, maxlen+1]
token_type = token_type # [batch_size, maxlen+1]
next_token_type = next_token_type # [batch_size, maxlen+1]
seq_feat = seq_feat # [batch_size, maxlen+1]
pos_feat = pos_feat # [batch_size, maxlen+1]
neg_feat = neg_feat # [batch_size, maxlen+1]
# 前向传播
pos_logits, neg_logits = model(seq, pos, neg, token_type, next_token_type,
next_action_type, seq_feat, pos_feat, neg_feat)
# pos_logits: [batch_size, maxlen+1] - 正样本logits
# neg_logits: [batch_size, maxlen+1] - 负样本logits
# 计算损失 (只对item token计算)
indices = np.where(next_token_type == 1)
loss = bce_criterion(pos_logits[indices], pos_labels[indices])
loss += bce_criterion(neg_logits[indices], neg_labels[indices])
3.2 模型前向传播（forward）
def forward(self, user_item, pos_seqs, neg_seqs, mask, next_mask, next_action_type,
seq_feature, pos_feature, neg_feature):
"""
训练时前向传播
Returns:
pos_logits: [batch_size, maxlen] - 正样本logits
neg_logits: [batch_size, maxlen] - 负样本logits
"""
# 1. 序列编码
log_feats = self.log2feats(user_item, mask, seq_feature) # [batch_size, seq_len, hidden_units]
# 2. 正负样本embedding
pos_embs = self.feat2emb(pos_seqs, pos_feature, include_user=False) # [batch_size, seq_len, hidden_units]
neg_embs = self.feat2emb(neg_seqs, neg_feature, include_user=False) # [batch_size, seq_len, hidden_units]
# 3. 计算相似度
pos_logits = (log_feats * pos_embs).sum(dim=-1) # [batch_size, seq_len]
neg_logits = (log_feats * neg_embs).sum(dim=-1) # [batch_size, seq_len]
# 4. 只对item token计算损失
loss_mask = (next_mask == 1).to(self.dev)
pos_logits = pos_logits * loss_mask
neg_logits = neg_logits * loss_mask
4. 推理过程 (infer.py)
4.1 用户表征生成（predict）
def predict(self, log_seqs, seq_feature, mask):
"""
生成用户序列表征
Args:
log_seqs: [batch_size, seq_len] - 用户序列
seq_feature: [batch_size, seq_len] - 序列特征
mask: [batch_size, seq_len] - token类型掩码
Returns:
final_feat: [batch_size, hidden_units] - 用户表征
"""
log_feats = self.log2feats(log_seqs, mask, seq_feature) # [batch_size, seq_len, hidden_units]
final_feat = log_feats[:, -1, :] # 取最后一个token作为用户表征
return final_feat
4.2 候选物品Embedding生成（get_candidate_emb, save_item_emb, process_cold_start_feat函数）
核心功能: 为推荐系统中的所有候选物品生成embedding向量，用于后续的向量检索和相似度计算。
作用说明:
-
离线预计算: 在推理阶段之前，预先为所有候选物品生成embedding，避免在线计算延迟
-
向量检索: 生成的embedding用于ANN（近似最近邻）检索，快速找到与用户query最相似的物品
-
存储优化: 将embedding保存到文件，支持大规模物品库的高效检索
候选库生成函数:
def get_candidate_emb(indexer, feat_types, feat_default_value, mm_emb_dict, model):
"""
生成候选库item的id和embedding
"""
# 1. 读取候选库数据文件
candidate_path = Path(os.environ.get('EVAL_DATA_PATH'), 'predict_set.jsonl')
# 2. 处理每个候选物品的特征
for line in f:
feature = line['features']
creative_id = line['creative_id']
retrieval_id = line['retrieval_id']
# 3. 处理冷启动特征和缺失特征
feature = process_cold_start_feat(feature)
for feat_id in missing_fields:
feature[feat_id] = feat_default_value[feat_id]
# 4. 处理多模态特征
for feat_id in feat_types['item_emb']:
if creative_id in mm_emb_dict[feat_id]:
feature[feat_id] = mm_emb_dict[feat_id][creative_id]
else:
feature[feat_id] = np.zeros(EMB_SHAPE_DICT[feat_id])
# 5. 调用模型生成embedding
model.save_item_emb(item_ids, retrieval_ids, features, save_path)
return retrieve_id2creative_id
冷启动特征处理:
def process_cold_start_feat(feature):
"""
处理冷启动特征。推理时遇到训练集未出现过的特征值，将字符串转换为0。
Args:
feature: dict - 物品特征字典，格式如 {'100': 123, '117': 'new_brand', '81': [0.1,0.2]}
其中key是特征ID，对应feat_types中的特征类型
"""
processed_feat = {}
for feat_id, feat_value in feature.items():
if type(feat_value) == list:
value_list = []
for v in feat_value:
if type(v) == str:
value_list.append(0) # 字符串特征转换为0
else:
value_list.append(v)
processed_feat[feat_id] = value_list
elif type(feat_value) == str:
processed_feat[feat_id] = 0 # 字符串特征转换为0
else:
processed_feat[feat_id] = feat_value # 数值特征保持不变
return processed_feat
冷启动处理说明:
-
feature含义: 每个物品的特征字典，key是特征ID（如'100', '117', '81'），对应feat_types中定义的特征类型
-
特征类型对应关系:
-
'100'→item_sparse(物品ID、类别等离散特征) -
'117'→item_sparse(品牌等离散特征) -
'81'→item_emb(多模态embedding特征)
-
-
问题: 推理时遇到训练集未出现过的特征值（如新品牌'new_brand'）
-
解决方案: 将字符串特征值统一转换为0，数值特征保持不变
-
应用: 在候选库特征处理阶段自动调用，确保模型输入的一致性
物品Embedding生成函数:
def save_item_emb(self, item_ids, retrieval_ids, feat_dict, save_path, batch_size=1024):
"""
生成候选库item embedding用于检索
Args:
item_ids: [num_items] - 候选item ID列表，包含所有需要生成embedding的物品ID
retrieval_ids: [num_items] - 检索ID列表，用于映射物品ID到检索ID
feat_dict: [num_items] - item特征字典，每个物品的所有特征信息，包含每个物品的稀疏特征、多模态特征等
save_path: str - 保存embedding的文件路径
batch_size: int - 批处理大小，用于内存优化
Returns:
无返回值，直接保存embedding到文件 save_emb(batch_emb, f"{save_path}/batch_{start_idx}.fbin")
"""
注意: retrieval_ids 是测试集json中的一个参数，可能是从0开始连续编号，便于高效的向量检索
4.3 检索流程
def infer():
# 1. 基于用户行为序列生成用户表征向量
for step, batch in enumerate(test_loader):
seq, token_type, seq_feat, user_id = batch
logits = model.predict(seq, seq_feat, token_type) # [batch_size, hidden_units]
all_embs.append(logits.detach().cpu().numpy())
# 2. 生成候选item embedding
retrieve_id2creative_id = get_candidate_emb(...)
# 3. 保存文件用于ANN检索
save_emb(all_embs, 'query.fbin') # 用户行为序列表征向量
save_emb(candidate_embs, 'embedding.fbin') # 候选库embedding
save_emb(candidate_ids, 'id.u64bin') # 候选库ID
# 4. 执行ANN检索
ann_cmd = "faiss_demo --query_vector_file_path=query.fbin --dataset_vector_file_path=embedding.fbin ..."
os.system(ann_cmd)
# 5. 读取检索结果
top10s_retrieved = read_result_ids("id100.u64bin") # [num_users, 10]
5. RQ-VAE模型架构 (model_rqvae.py)
参考资料：https://zhuanlan.zhihu.com/p/716658479
5.1 核心组件
K-means聚类:
def kmeans(data, n_clusters, kmeans_iters):
"""
标准K-means聚类算法
Args:
data: [num_samples, num_features] - 输入数据
n_clusters: 聚类数量
kmeans_iters: 最大迭代次数
Returns:
cluster_centers: [n_clusters, num_features] - 聚类中心
labels: [num_samples] - 聚类标签
"""
平衡K-means:
class BalancedKmeans(torch.nn.Module):
"""
平衡K-means聚类，确保每个聚类大小相对平衡
Args:
num_clusters: 聚类数量
kmeans_iters: 最大迭代次数
tolerance: 收敛容差
device: 计算设备
"""
5.2 编码器-解码器架构（代码中Encoder和Decoder只是简单的神经网络实现）
RQ编码器:
class RQEncoder(torch.nn.Module):
"""
RQ-VAE编码器
Args:
input_dim: 输入维度 (4096)
hidden_channels: 隐藏层通道数列表
latent_dim: 潜在空间维度
"""
def forward(self, x):
# x: [batch_size, input_dim] - 输入多模态embedding
# 返回: [batch_size, latent_dim] - 编码后的潜在表示
RQ解码器:
class RQDecoder(torch.nn.Module):
"""
RQ-VAE解码器
Args:
latent_dim: 潜在空间维度
hidden_channels: 隐藏层通道数列表
output_dim: 输出维度 (4096)
"""
def forward(self, x):
# x: [batch_size, latent_dim] - 潜在表示
# 返回: [batch_size, output_dim] - 重构的多模态embedding
5.3 向量量化嵌入
VQ嵌入层:
## Generate semantic id
class VQEmbedding(torch.nn.Embedding):
"""
向量量化嵌入层：将连续向量量化为离散的语义ID
功能：创建codebook → 计算距离 → 生成语义ID → 返回量化嵌入
"""
def __init__(
self,
num_clusters, # 聚类数量（codebook大小）
codebook_emb_dim: int, # codebook嵌入维度
kmeans_method: str, # kmeans初始化方法
kmeans_iters: int, # kmeans迭代次数
distances_method: str, # 距离计算方法
device: str, # 设备
):
super(VQEmbedding, self).__init__(num_clusters, codebook_emb_dim)
# 保存参数
self.num_clusters = num_clusters
self.codebook_emb_dim = codebook_emb_dim
self.kmeans_method = kmeans_method
self.kmeans_iters = kmeans_iters
self.distances_method = distances_method
self.device = device
def _create_codebook(self, data):
"""
创建codebook：根据数据初始化量化字典
支持三种初始化方式：
- kmeans: 标准kmeans聚类
- bkmeans: 平衡kmeans聚类
- random: 随机初始化
"""
if self.kmeans_method == 'kmeans':
# 标准kmeans聚类
_codebook, _ = kmeans(data, self.num_clusters, self.kmeans_iters)
elif self.kmeans_method == 'bkmeans':
# 平衡kmeans聚类（确保每个聚类大小相近）
BKmeans = BalancedKmeans(
num_clusters=self.num_clusters, kmeans_iters=self.kmeans_iters, tolerance=1e-4, device=self.device
)
_codebook, _ = BKmeans.fit(data)
else:
# 随机初始化
_codebook = torch.randn(self.num_clusters, self.codebook_emb_dim)
# 移动到指定设备并验证形状
_codebook = _codebook.to(self.device)
assert _codebook.shape == (self.num_clusters, self.codebook_emb_dim)
self.codebook = torch.nn.Parameter(_codebook)
@torch.no_grad()
def _compute_distances(self, data):
"""
计算数据与codebook的距离
支持两种距离计算：
- cosine: 余弦距离
- l2: 欧几里得距离
"""
_codebook_t = self.codebook.t() # 转置codebook
assert _codebook_t.shape == (self.codebook_emb_dim, self.num_clusters)
assert data.shape[-1] == self.codebook_emb_dim
if self.distances_method == 'cosine':
# 余弦距离：1 - cos(θ)
data_norm = F.normalize(data, p=2, dim=-1) # 数据归一化
_codebook_t_norm = F.normalize(_codebook_t, p=2, dim=0) # codebook归一化
distances = 1 - torch.mm(data_norm, _codebook_t_norm) # 计算余弦距离
else:
# L2距离：||x - c||² = ||x||² + ||c||² - 2<x,c>
data_norm_sq = data.pow(2).sum(dim=-1, keepdim=True) # ||x||²
_codebook_t_norm_sq = _codebook_t.pow(2).sum(dim=0, keepdim=True) # ||c||²
distances = torch.addmm(data_norm_sq + _codebook_t_norm_sq, data, _codebook_t, beta=1.0, alpha=-2.0)
return distances
@torch.no_grad()
def _create_semantic_id(self, data):
"""
生成语义ID：找到最近的codebook向量
Returns:
_semantic_id: 每个样本对应的codebook索引
"""
distances = self._compute_distances(data) # 计算距离
_semantic_id = torch.argmin(distances, dim=-1) # 找到最小距离的索引
return _semantic_id
def _update_emb(self, _semantic_id):
"""
根据语义ID获取对应的量化嵌入
Args:
_semantic_id: 语义ID
Returns:
update_emb: 对应的量化嵌入
"""
update_emb = super().forward(_semantic_id) # 使用父类的embedding查找
return update_emb
def forward(self, data):
"""
前向传播：完整的量化流程
Args:
data: 输入数据
Returns:
update_emb: 量化后的嵌入
_semantic_id: 语义ID
"""
self._create_codebook(data) # 创建codebook
_semantic_id = self._create_semantic_id(data) # 生成语义ID
update_emb = self._update_emb(_semantic_id) # 获取量化嵌入
return update_emb, _semantic_id
5.4 残差量化器
RQ模块:
## Residual Quantizer
class RQ(torch.nn.Module):
"""
残差量化器：通过多个量化步骤逐步减少输入数据的残差
Args:
num_codebooks: 量化器数量
codebook_size: 每个codebook的大小列表
codebook_emb_dim: codebook嵌入维度
shared_codebook: 是否共享codebook
kmeans_method: kmeans初始化方法
kmeans_iters: kmeans迭代次数
distances_method: 距离计算方法
loss_beta: RQ-VAE损失权重
device: 设备
"""
def __init__(
self,
num_codebooks: int,
codebook_size: list,
codebook_emb_dim,
shared_codebook: bool,
kmeans_method,
kmeans_iters,
distances_method,
loss_beta: float,
device: str,
):
super().__init__()
# 基本参数设置
self.num_codebooks = num_codebooks
self.codebook_size = codebook_size
assert len(self.codebook_size) == self.num_codebooks # 确保codebook数量匹配
self.codebook_emb_dim = codebook_emb_dim
self.shared_codebook = shared_codebook
# 量化相关参数
self.kmeans_method = kmeans_method
self.kmeans_iters = kmeans_iters
self.distances_method = distances_method
self.loss_beta = loss_beta
self.device = device
# 创建量化模块列表
if self.shared_codebook:
# 共享codebook：所有量化器使用相同的codebook大小
self.vqmodules = torch.nn.ModuleList(
[
VQEmbedding(
self.codebook_size[0], # 使用第一个codebook大小
self.codebook_emb_dim,
self.kmeans_method,
self.kmeans_iters,
self.distances_method,
self.device,
)
for _ in range(self.num_codebooks)
]
)
else:
# 独立codebook：每个量化器使用不同的codebook大小
self.vqmodules = torch.nn.ModuleList(
[
VQEmbedding(
self.codebook_size[idx], # 使用对应索引的codebook大小
self.codebook_emb_dim,
self.kmeans_method,
self.kmeans_iters,
self.distances_method,
self.device,
)
for idx in range(self.num_codebooks)
]
)
def quantize(self, data):
"""
关键函数：执行残差量化过程
量化流程：
- 第i步量化：输入[i] = VQ[i] + 残差[i]
- vq_emb_list: [vq1, vq1+vq2, vq1+vq2+vq3, ...] 累积量化结果
- res_emb_list: [res1, res2, res3, ...] 每步残差
- semantic_id_list: [vq1_sid, vq2_sid, vq3_sid, ...] 每步语义ID
Returns:
vq_emb_list: 累积的量化嵌入列表
res_emb_list: 残差嵌入列表
semantic_id_list: 语义ID张量 [batch_size, num_codebooks]
"""
# 初始化残差为输入数据的副本
res_emb = data.detach().clone()
vq_emb_list, res_emb_list = [], []
semantic_id_list = []
vq_emb_aggre = torch.zeros_like(data) # 累积量化结果
# 逐步量化
for i in range(self.num_codebooks):
# 使用第i个量化器量化当前残差
vq_emb, _semantic_id = self.vqmodules[i](res_emb)
# 更新残差：残差 = 残差 - 量化结果
res_emb -= vq_emb
# 累积量化结果
vq_emb_aggre += vq_emb
# 保存结果
res_emb_list.append(res_emb)
vq_emb_list.append(vq_emb_aggre)
semantic_id_list.append(_semantic_id.unsqueeze(dim=-1))
# 将所有语义ID拼接成张量
semantic_id_list = torch.cat(semantic_id_list, dim=-1)
return vq_emb_list, res_emb_list, semantic_id_list
def _rqvae_loss(self, vq_emb_list, res_emb_list):
"""
计算RQ-VAE损失
RQ-VAE损失包含两部分：
- loss1: 量化结果与残差的匹配损失（停止残差梯度）
- loss2: 残差与量化结果的匹配损失（停止量化梯度）
"""
rqvae_loss_list = []
for idx, quant in enumerate(vq_emb_list):
# 计算两个方向的匹配损失
loss1 = (res_emb_list[idx].detach() - quant).pow(2.0).mean() # 残差梯度停止
loss2 = (res_emb_list[idx] - quant.detach()).pow(2.0).mean() # 量化梯度停止
partial_loss = loss1 + self.loss_beta * loss2
rqvae_loss_list.append(partial_loss)
# 所有步骤损失求和
rqvae_loss = torch.sum(torch.stack(rqvae_loss_list))
return rqvae_loss
def forward(self, data):
"""
前向传播：执行量化并计算损失
Returns:
vq_emb_list: 累积量化嵌入
semantic_id_list: 语义ID
rqvae_loss: RQ-VAE损失
"""
# 执行量化
vq_emb_list, res_emb_list, semantic_id_list = self.quantize(data)
# 计算损失
rqvae_loss = self._rqvae_loss(vq_emb_list, res_emb_list)
return vq_emb_list, semantic_id_list, rqvae_loss
5.5 完整RQ-VAE模型
RQVAE类:
class RQVAE(torch.nn.Module):
"""
RQ-VAE模型：结合编码器、残差量化器和解码器的完整模型
工作流程：输入 → 编码 → 残差量化 → 解码 → 重构输出
"""
def __init__(
self,
input_dim: int, # 输入维度
hidden_channels: list, # 隐藏层通道数列表
latent_dim: int, # 潜在空间维度
num_codebooks: int, # 量化器数量
codebook_size: list, # 每个codebook的大小
shared_codebook: bool, # 是否共享codebook
kmeans_method, # kmeans方法
kmeans_iters, # kmeans迭代次数
distances_method, # 距离计算方法
loss_beta: float, # RQ-VAE损失权重
device: str, # 设备
):
super().__init__()
# 编码器：将输入映射到潜在空间
self.encoder = RQEncoder(input_dim, hidden_channels, latent_dim).to(device)
# 解码器：将潜在表示重构回原始空间（隐藏层顺序反转）
self.decoder = RQDecoder(latent_dim, hidden_channels[::-1], input_dim).to(device)
# 残差量化器：在潜在空间进行量化
self.rq = RQ(
num_codebooks,
codebook_size,
latent_dim,
shared_codebook,
kmeans_method,
kmeans_iters,
distances_method,
loss_beta,
device,
).to(device)
def encode(self, x):
"""编码：输入 → 潜在表示"""
return self.encoder(x)
def decode(self, z_vq):
"""解码：潜在表示 → 重构输出"""
# 如果输入是列表（量化结果列表），使用最后一个（累积结果）
if isinstance(z_vq, list):
z_vq = z_vq[-1]
return self.decoder(z_vq)
def compute_loss(self, x_hat, x_gt, rqvae_loss):
"""
计算总损失
Args:
x_hat: 重构输出
x_gt: 原始输入
rqvae_loss: RQ-VAE损失
Returns:
recon_loss: 重构损失
rqvae_loss: RQ-VAE损失
total_loss: 总损失
"""
# 重构损失：重构输出与原始输入的MSE
recon_loss = F.mse_loss(x_hat, x_gt, reduction="mean")
# 总损失：重构损失 + RQ-VAE损失
total_loss = recon_loss + rqvae_loss
return recon_loss, rqvae_loss, total_loss
def _get_codebook(self, x_gt):
"""
获取语义ID（用于推理）
Returns:
semantic_id_list: 语义ID列表
"""
z_e = self.encode(x_gt) # 编码
vq_emb_list, semantic_id_list, rqvae_loss = self.rq(z_e) # 量化
return semantic_id_list # 只返回语义ID
def forward(self, x_gt):
"""
前向传播：完整的RQ-VAE流程
Returns:
x_hat: 重构输出
semantic_id_list: 语义ID
recon_loss: 重构损失
rqvae_loss: RQ-VAE损失
total_loss: 总损失
"""
# 1. 编码：输入 → 潜在表示
z_e = self.encode(x_gt)
# 2. 量化：潜在表示 → 量化结果 + 语义ID
vq_emb_list, semantic_id_list, rqvae_loss = self.rq(z_e)
# 3. 解码：量化结果 → 重构输出
x_hat = self.decode(vq_emb_list)
# 4. 计算损失
recon_loss, rqvae_loss, total_loss = self.compute_loss(x_hat, x_gt, rqvae_loss)
return x_hat, semantic_id_list, recon_loss, rqvae_loss, total_loss
6. 关键变量含义
6.1 数据相关
-
seq: 用户交互序列ID，包含用户token和物品token -
token_type: token类型，1=物品token，2=用户token -
next_token_type: 下一个token类型 -
next_action_type: 动作类型，0=曝光，1=点击 -
seq_feat/pos_feat/neg_feat: 特征字典，key为特征ID，value为特征值
6.2 模型相关
-
hidden_units: 隐藏层维度，默认32 -
maxlen: 序列最大长度，默认101 -
num_heads: 注意力头数，默认1 -
num_blocks: Transformer层数，默认1
6.3 特征相关
-
sparse_emb: 稀疏特征embedding表 -
emb_transform: 多模态特征线性变换层 -
feat_statistics: 特征统计信息，包含各特征的数量 -
feat_types: 特征类型分类，分为sparse/array/emb/continual
浙公网安备 33010602011771号