[doc] SID: add RQ-VAE / RQ-KMeans semantic-ID generation docs

docs/source/index.rst

@@ -26,6 +26,7 @@ Welcome to TorchEasyRec's documentation!
    models/rank
    models/multi_target
    models/generative
+   models/sid_model
    models/user_define
    models/feature_group
    models/loss

docs/source/models/sid_model.rst

@@ -0,0 +1,8 @@
+语义ID生成
+==========
+
+.. toctree::
+   :maxdepth: 2
+
+   sid_rqvae
+   sid_rqkmeans

docs/source/models/sid_rqkmeans.md

@@ -0,0 +1,106 @@
+# SID RQKMeans
+
+## 简介
+
+RQKMeans (Residual K-Means) 是一种语义ID (Semantic ID, 简称 SID) 生成模型, 和 [SID RQVAE](sid_rqvae.md) 一样把物品 embedding 量化成一串离散编码 `(code_0, code_1, ..., code_{n-1})`, 用作生成式推荐的物品token。区别在于: RQKMeans **用 FAISS K-Means 直接对残差做聚类来得到每层码本**——码本是离线一次性拟合出来的。
+
+拟合过程逐层进行: 第0层对原始 embedding 做 K-Means, 得到 `codebook_0` 个聚类中心作为该层码本; 每个样本减去其最近中心得到残差, 第 1 层再对残差做 K-Means, 依此类推。物品的 SID 即各层最近中心的下标元组。相比 RQVAE, RQKMeans 训练更快、无需调梯度超参, 适合在数据充足时快速产出码本。
+
+注意: RQKMeans 当前仅支持 CPU 且单进程——若检测到可用 CUDA 设备或 `world_size > 1` 会直接报错, 因此**务必使用 `--nproc-per-node=1`** 并在 CPU 环境运行。需要多卡/梯度训练时请改用 [SID RQVAE](sid_rqvae.md)。
+
+## 数据格式
+
+RQKMeans 只需要物品 embedding 一列, 通过`fg_mode: FG_DAG` 读取数组列:
+
+| 列名        | 类型           | 说明                                                               |
+| ----------- | -------------- | ------------------------------------------------------------------ |
+| `item_id`   | string         | 物品 ID (透传列, 模型不消费, 预测时可用 `--reserved_columns` 带出) |
+| `embedding` | array\<double> | 物品 embedding, 维度需与 `value_dim` 一致 (示例为 512)             |
+
+> 拟合要求样本数 `N >= max(codebook)`。示例使用较小的 `codebook: 256` 以适配小样本; 生产规模的 `codebook: 8192` 需要远多于 8192 行的数据。样本偏少时 FAISS 会打印 `please provide at least ... training points` 告警, 不影响跑通。
+
+## 配置说明
+
+```
+train_config {
+    sparse_optimizer { adagrad_optimizer { lr: 0.001 } constant_learning_rate {} }
+    dense_optimizer { adam_optimizer { lr: 0.00002 } constant_learning_rate {} }
+    num_epochs: 1
+    save_checkpoints_steps: 0
+    save_checkpoints_epochs: 0
+}
+
+feature_configs {
+    raw_feature { feature_name: "emb" expression: "item:embedding" value_dim: 512 }
+}
+
+model_config {
+    feature_groups { group_name: "deep" feature_names: "emb" group_type: DEEP }
+    sid_rqkmeans {
+        codebook: 256
+        codebook: 256
+        codebook: 256
+        normalize_residuals: true
+        faiss_kmeans_kwargs {
+            niter: 20
+            seed: 42
+            verbose: true
+            spherical: false
+        }
+    }
+}
+```
+
+- train_config: 训练流程配置。RQKMeans 不做梯度训练, 因此优化器与 epoch 数仅为框架训练循环所需的形式配置
+  - sparse_optimizer / dense_optimizer: 必填 (训练循环要构建优化器), 但模型只有一个 dummy 参数、损失恒为 0, **因此其学习率不影响最终码本**
+  - num_epochs: **设 `1` 即可**; 每个 epoch 只把 embedding 流式写入蓄水池采样 (reservoir), 真正的 FAISS 拟合在训练结束 (`on_train_end`) 时一次性完成, 多跑 epoch 不会迭代/精炼码本
+  - save_checkpoints_steps / save_checkpoints_epochs: **必须项，设 `0` 关闭周期性保存**; 拟合好的码本只随训练结束时的最终 checkpoint 持久化 (周期性保存可能会忽略checkpoint落盘, 故关闭)
+- feature_configs / feature_groups: 同 RQVAE, 但只需主物品 embedding 一组 (`deep`); 其拼接后的总维度即 K-Means 的向量维度
+- sid_rqkmeans: RQKMeans 模型参数
+  - codebook: 每层聚类中心数; **列表长度即残差层数 (= SID 的位数)**; 示例 `[256, 256, 256]` (生产常用 `[8192, 8192, 8192]`); 支持非均匀如 `[256, 512, 1024]` (每层独立拟合一个 faiss.Kmeans)
+  - normalize_residuals: 每层聚类前是否对残差做 L2 归一化, 默认 `false` (开启后效果更佳)
+  - faiss_kmeans_kwargs: 透传给 `faiss.Kmeans(D, K, **kwargs)` 的强类型参数, 未设置的字段回落到 FAISS 自身默认值
+    - niter: 每层 K-Means 迭代次数 (FAISS 默认 25)
+    - nredo: 重复聚类并取最优的次数 (默认 1)
+    - seed: 随机种子, 影响中心初始化/采样 (默认 1234)
+    - max_points_per_centroid: 每个中心的最大样本数, FAISS 会下采样到 `K * 该值` (默认 256); 也用于自动推导蓄水池容量
+    - min_points_per_centroid: 每个中心的最小样本数 (默认 39), 低于 `K * 该值` 仅告警不报错
+    - spherical: 是否做球面 (余弦) K-Means, 即每轮对中心归一化 (默认 false)
+    - verbose: 是否打印 FAISS 自身的逐轮日志 (默认 false)
+  - train_sample_size: 为拟合做蓄水池采样的目标样本数, 用于限制 host 内存; `0` (默认) 自动取 `max(codebook) * max_points_per_centroid` (即 FAISS 内部会下采样到的规模); 若显式设为小于 `max(codebook)` 的值, 模型在初始化时即报错 (需 `>= max(codebook)`)
+
+## 示例
+
+模型的训练和评估方式同[local_tutorial](../quick_start/local_tutorial.md)，示例数据和配置参数如下：
+
+### 数据
+
+[sid_generation_item_only_sample_4w.parquet](https://tzrec.oss-accelerate.aliyuncs.com/data/models/sid_generation_item_only_sample_4w.parquet)
+
+### 配置文件
+
+[sid_rq_kmeans.config](https://tzrec.oss-accelerate.aliyuncs.com/config/models/sid_rq_kmeans.config)
+
+### 训练参数
+
+```bash
+OMP_NUM_THREADS=$(nproc) \
+    torchrun --master_addr=localhost --master_port=32555 \
+    --nnodes=1 --nproc-per-node=1 --node_rank=0 \
+    -m tzrec.train_eval \
+    --pipeline_config_path sid_rq_kmeans.config \
+    --train_input_path "data/sid_example/item_only/*.parquet" \
+    --eval_input_path "data/sid_example/item_only/*.parquet" \
+    --model_dir experiments/sid_rqkmeans
+```
+
+`--nproc-per-node` 必须为 `1`。训练时每个 batch 只把 embedding 写入蓄水池并返回占位 (全 0) 编码; 训练结束时日志会打印 `[SidRqkmeans.on_train_end] fitting FAISS on N samples` 以及逐层聚类信息, 随后做一次 eval 输出 `mse` / `rel_loss` / `unique_sid_ratio`。
+
+**`OMP_NUM_THREADS=$(nproc)` 为必须配置项，否则默认设置为1，影响模型训练推理速度。**
+
+### 模型输出
+
+预测输出与输入 `dataset_type` 一致, 每行包含:
+
+- codes: `array<int64>`, 即该物品的 SID, 长度等于 `codebook` 层数, 每个元素为对应残差层的中心下标 (取值范围 `[0, codebook_i)`)。例如 `[8, 31, 26]`。
+- item_id: 由 `--reserved_columns` 透传的原始物品 ID。

docs/source/models/sid_rqvae.md

@@ -0,0 +1,150 @@
+# SID RQVAE
+
+## 简介
+
+RQ-VAE (Residual-Quantized Variational Auto-Encoder) 是一种语义ID (Semantic ID, 简称 SID) 生成模型, 用于把物品的内容/多模态 embedding 量化成一串离散的整数编码 `(code_0, code_1, ..., code_{n-1})`, 每个残差码本层产生一个 code。它是生成式推荐的上游 tokenizer: 先用 RQVAE 把物品 embedding 转成 SID, 下游再用 SID 作为物品 token 进行生成式建模。
+
+模型结构为 `编码器 MLP -> 多层残差向量量化 (Residual Vector Quantizer) -> 解码器 MLP`: 编码器把输入 embedding 压缩到 `embed_dim` 维的潜在向量, 残差量化器逐层把当前残差就近映射到该层码本 (一个可学习的 `nn.Embedding`) 的最近条目, 解码器再用各层量化向量之和重建原始 embedding。码本、编码器、解码器通过梯度联合训练, 量化的不可导通过直通估计 (Straight-Through Estimator, STE) 传递梯度。物品的 SID 即为各残差层最近邻码本下标组成的元组。
+
+![rqvae.png](../../images/models/rqvae.png)
+
+注意: RQVAE 通过梯度训练码本, 可多卡训练; 若希望用 FAISS 直接对残差做 K-Means 聚类得到码本 (无编码器/解码器, 单机 CPU), 请参考 [SID RQKMeans](sid_rqkmeans.md)。
+
+RQVAE 还支持双视图对比学习 (contrastive, 思路类似 CLIP 的图文对比): 除主物品 embedding 外, 再输入一个配对 (paired) embedding 以及一个 0/1 的配对标记, 在配对样本上额外施加 InfoNCE 对比损失, 使语义相近的两个视图被编码到相近的 SID 空间。本文示例即为带对比损失的配置; 关闭对比的用法见文末。
+
+## 数据格式
+
+RQVAE 的输入是离线预先算好的物品 embedding (例如多模态/文本/图像 embedding), 以 `fg_mode: FG_DAG` 直接读取数组列。带对比 (contrastive) 的样本单文件包含如下列:
+
+| 列名             | 类型           | 说明                                                                |
+| ---------------- | -------------- | ------------------------------------------------------------------- |
+| `item_id`        | string         | 物品 ID (透传列, 模型不消费, 预测时可用 `--reserved_columns` 带出)  |
+| `embedding`      | array\<double> | 主物品 embedding, 维度需与 `value_dim` 一致 (示例为 512)            |
+| `pair_item_id`   | string         | 配对物品 ID (透传列)                                                |
+| `pair_embedding` | array\<double> | 配对 embedding, 维度同 `embedding`                                  |
+| `is_pair`        | int32 (0/1)    | 配对标记: `1`=对比对样本, `0`=纯重建样本 (此时 `pair_*` 与主列相同) |
+
+非对比 (纯重建) 场景只需 `item_id` + `embedding` 两列。
+
+`is_pair` 在语义上是布尔标记, 但**以 0/1 整数存储而非 `bool`**，以 `> 0.5` 判定是否为对比对。
+
+## 配置说明
+
+```
+feature_configs {
+    raw_feature { feature_name: "emb" expression: "item:embedding" value_dim: 512 }
+}
+feature_configs {
+    raw_feature { feature_name: "pair_emb" expression: "item:pair_embedding" value_dim: 512 }
+}
+feature_configs {
+    raw_feature { feature_name: "is_pair" expression: "item:is_pair" value_dim: 1 }
+}
+
+model_config {
+    feature_groups {
+        group_name: "deep"
+        feature_names: "emb"
+        group_type: DEEP
+    }
+    feature_groups {
+        group_name: "pair"
+        feature_names: "pair_emb"
+        group_type: DEEP
+    }
+    feature_groups {
+        group_name: "pair_flag"
+        feature_names: "is_pair"
+        group_type: DEEP
+    }
+    sid_rqvae {
+        embed_dim: 64
+        hidden_dims: 256
+        hidden_dims: 256
+        codebook: 256
+        codebook: 256
+        codebook: 256
+        forward_mode: "ste"
+        kmeans_init: false
+        contrastive_config {
+            pair_feature_group: "pair"
+            pair_flag_feature_group: "pair_flag"
+        }
+    }
+    losses {
+        recon_loss {
+            recon_type: "l2"
+        }
+    }
+    losses {
+        commitment_loss {
+            latent_weight: 0.5
+            latent_weight: 0.5
+        }
+    }
+    losses {
+        contrastive_loss {}
+    }
+}
+```
+
+- feature_configs: 特征配置, 每个 embedding/标记列对应一个 `raw_feature`
+  - emb: 主物品 embedding, `value_dim` 须等于 embedding 维度 (示例 512)
+  - pair_emb: 配对 embedding, `value_dim` 须与主 embedding 相同
+  - is_pair: 0/1 配对标记, `value_dim: 1`
+- feature_groups: 特征组, 每组喂给模型的 `EmbeddingGroup`; **group_name 需与下方 sid_rqvae 的引用对应**
+  - deep: 主输入组; 模型自动取**第一个声明的 feature group** 作为主输入, 因此主输入组须列在最前; 其拼接后的总维度即编码器输入维度 `input_dim`; 类型 DEEP
+  - pair: 配对 embedding 组 (仅对比时需要), 须与主组同维; 类型 DEEP
+  - pair_flag: 配对标记组 (仅对比时需要); 类型 DEEP
+- sid_rqvae: RQVAE 模型参数 (主输入特征组无需单独配置, 取第一个声明的 feature group)
+  - embed_dim: 量化潜在维度, 即编码器输出维度与码本向量维度, 默认 64
+  - hidden_dims: 编码器各隐层大小 (解码器按反序镜像), 如 `[256, 256]`; 不配置时默认为 `[input_dim // 2]`
+  - codebook: 每层码本大小; **列表长度即残差量化层数 (= SID 的位数)**; 示例为 `[256, 256, 256]` (生产规模常用 `[8192, 8192, 8192]`); 支持非均匀如 `[512, 256, 128]`
+  - forward_mode: VQ 前向模式, `"ste"` (默认) 或 `"gumbel_softmax"`
+  - normalize_residuals: 每层量化前是否对残差做 L2 归一化, 默认 `false` (开启后重建度量 `mse`/`rel_loss` 不再可比)
+  - distance_type: 最近邻度量, `"l2"` (默认) 或 `"cosine"`
+  - rotation_trick: 是否使用 rotation-trick 形式的 STE 梯度, 默认 `false`
+  - kmeans_init: 是否在首个训练 batch 上用 FAISS 残差 K-Means 初始化码本, 默认 `false`; 需 `batch_size >= max(codebook)`
+  - sinkhorn_config: Sinkhorn 均匀分配子配置 (省略时默认启用, `iters=5`, `epsilon=10.0`; 设 `enabled: false` 可关闭)
+  - contrastive_config: 双视图对比结构 (不配置则关闭对比)
+    - pair_feature_group: 配对 embedding 特征组名, 须与主组同维 (示例 `pair`)
+    - pair_flag_feature_group: 配对标记特征组名 (维度 1, `>0.5` 视为对比对; 示例 `pair_flag`)
+- losses: 损失配置, 每个 `losses {}` 配一个 SID 损失项, 总损失为各项之和
+  - recon_loss: 重建损失 (解码输出 vs 输入 embedding), `recon_type` 可选 `"l2"`/`"l1"`/`"cos"`; 对比模式下仅在非配对行 (重建样本) 上计算
+  - commitment_loss: VQ commitment 损失; `latent_weight` 为 `[w1, w2]` 两个权重 (默认 `[1.0, 0.5]`, **必须长度为 2**), `commitment_type` 可选 `"l2"`/`"l1"`/`"cos"`
+  - contrastive_loss: 开启双视图对比 (masked InfoNCE) 损失; **须与 `contrastive_config` 同时设置**; 仅在配对行上生效
+
+> 评估指标自动输出 `mse` (重建均方误差)、`rel_loss` (相对 L1)、`unique_sid_ratio` (每个 batch 内不重复 SID 占比, 反映码本利用/多样性)。
+
+## 示例
+
+### 配置文件
+
+[sid_rqvae.config](https://tzrec.oss-accelerate.aliyuncs.com/config/models/sid_rqvae.config)
+
+### 数据
+
+对比 (contrastive) 混合数据：
+[sid_generation_contrastive_sample_4w.parquet](https://tzrec.oss-accelerate.aliyuncs.com/data/models/sid_generation_contrastive_sample_4w.parquet)
+
+非对比 Item 数据：
+[sid_generation_item_only_sample_4w.parquet](https://tzrec.oss-accelerate.aliyuncs.com/data/models/sid_generation_item_only_sample_4w.parquet)
+
+### 模型输出
+
+预测输出与输入 `dataset_type` 一致, 每行包含:
+
+- codes: `array<int64>`, 即该物品的 SID, 长度等于 `codebook` 层数, 每个元素为对应残差层的码本下标 (取值范围 `[0, codebook_i)`)。例如 `[241, 134, 78]`。
+- item_id: 由 `--reserved_columns` 透传的原始物品 ID。
+
+### 非对比 (纯重建) 用法
+
+如果不需要对比学习, 去掉 `pair` / `pair_flag` 两个特征组、对应的 `pair_emb` / `is_pair` 特征、`sid_rqvae.contrastive_config` 以及 `contrastive_loss`, 输入改用仅含 `item_id` + `embedding` 的 item-only 样本 (上方数据链接) 即可。
+
+## 参考论文
+
+[FORGE: Forming Semantic Identifiers for Generative Retrieval in Industrial Datasets](https://arxiv.org/abs/2509.20904)
+
+[Recommender Systems with Generative Retrieval (TIGER)](https://arxiv.org/abs/2305.05065)
+
+[Autoregressive Image Generation using Residual Quantization (RQ-VAE)](https://arxiv.org/abs/2203.01941)