diff --git a/cornac/models/__init__.py b/cornac/models/__init__.py index c4f9f03a..5b95a2aa 100644 --- a/cornac/models/__init__.py +++ b/cornac/models/__init__.py @@ -42,6 +42,7 @@ from .cvae import CVAE from .cvaecf import CVAECF from .dmrl import DMRL +from .diffgrm import DiffGRM from .dnntsp import DNNTSP from .ease import EASE from .efm import EFM diff --git a/cornac/models/diffgrm/README.md b/cornac/models/diffgrm/README.md new file mode 100644 index 00000000..9b8c3199 --- /dev/null +++ b/cornac/models/diffgrm/README.md @@ -0,0 +1,211 @@ +# DiffGRM + +Cornac implementation of **DiffGRM: Diffusion-based Generative Recommendation Model** (WWW 2026, [paper](https://arxiv.org/abs/2510.21805)). DiffGRM is a generative next-item recommender that represents each catalog item as a short Semantic ID and predicts all positions through a masked-diffusion process rather than autoregressive left-to-right generation. + +DiffGRM combines three mechanisms: + +1. **Parallel Semantic Encoding (PSE)** whitens item-content embeddings with PCA and uses OPQ/PQ to create four independent 8-bit Semantic-ID digits. +2. **On-policy Coherent Noising (OCN)** probes a fully masked target, ranks its digits by confidence, and constructs nested masked views that focus learning on the most uncertain digits. +3. **Confidence-guided Parallel Denoising (CPD)** lets every unfilled digit/code pair compete within a global beam, so the decoding order is selected from the model's confidence rather than fixed in advance. + +The implementation was independently written for Cornac from the published architecture and equations. The [released research repository](https://github.com/liuzhao09/DiffGRM) was audited at commit `ad7b971c7e525e9fea6fb8e362a5c49dccb2473c` to validate public behavior and resolve ambiguities between the paper and release (reproduce the results from the official repo and compare against the paper). That commit did not include a root repository license, so its source code was not copied into Cornac. + +## Requirements + +Install the optional DiffGRM dependencies listed in `requirements.txt`: + +```bash +pip install -r cornac/models/diffgrm/requirements.txt +``` + +DiffGRM requires `torch`, `faiss-cpu`, and `scikit-learn`. + +## Usage + +DiffGRM consumes precomputed item-content embeddings through Cornac's `FeatureModality`: + +```python +from cornac.data import FeatureModality +from cornac.eval_methods import NextItemEvaluation +from cornac.models.diffgrm import DiffGRM, DIFFGRM_SPORTS_CONFIG + +item_feature = FeatureModality(features=item_embeddings, ids=item_ids) +eval_method = NextItemEvaluation.from_splits( + train_data=train_data, + val_data=val_data, + test_data=test_data, + fmt="USIT", + item_feature=item_feature, +) + +model = DiffGRM(**{**DIFFGRM_SPORTS_CONFIG, "seed": 2026}) +``` + +`item_embeddings` must cover every mapped item, including validation and test items, because the fitted transform is applied to the full recommendation catalog. PCA and OPQ/PQ are fitted only on items exposed by the released training-row construction; held-out interactions do not enter that fit. Sentence-T5 encoding is an offline preprocessing step and is never downloaded implicitly by the model. + +For artifact-controlled experiments, pass an integer array with shape `(n_items, n_digit)` through `item_sids`. This bypasses PSE and is the recommended way to compare model behavior while keeping a published Semantic-ID artifact fixed. + +The released sliding augmenter creates prefix targets from `min_history` through `max_len` and stops after the first `max_len` target. It does not roll a length-`max_len` window over later targets in a longer sequence. Cornac follows this released behavior, and PSE fitting uses the union of items exposed by those exact training rows. + +## Training and inference controls + +### Masking and loss + +`masking_strategy="guided"` enables OCN. The alternatives are designed as explicit controls: `"random"` independently masks digits for the no-OCN comparison, `"coherent"` creates nested random-order masks without on-policy selection, and `"fixed"` always uses a fixed digit order. The released no-OCN recipe uses four independent views with `random_mask_prob=0.5`. + +The paper averages the masked loss within each view and then across views, exposed as `view_loss_reduction="view_mean"`. The released implementation pools all masked tokens across views instead, exposed as `view_loss_reduction="token_mean"`. + +### Decoding + +The `scoring` option separates the paper algorithm, released behavior, and controlled ablations: + +- `scoring="paper"` follows Equations 8--10 and performs global beam selection at every denoising step. +- `scoring="released"` reproduces the released CPD behavior, including greedy completion of the final digit, followed by complete-catalog filtering. +- `scoring="catalog"` applies paper CPD while constraining every partial assignment to catalog-compatible prefixes. +- `scoring="fixed"` uses a seeded fixed digit permutation and serves as the no-CPD control. + +Paper-style configurations use validation beam 32 through `val_beam_size` and the dataset-specific `beam_size` for test scoring. Validation decoding is batched according to `val_batch_size`. + +### Collisions and model selection + +PSE can map multiple items to the same Semantic ID. The default `collision_policy="all"` assigns a decoded path score to every item with that ID. `"last"` reproduces the released reverse-map behavior, where the last catalog item sharing an ID overwrites earlier items, while `"first"` retains only the lowest-index item. + +Checkpoint selection defaults to `model_selection="best"` and maximizes the SID-level validation objective `0.8 * NDCG@k + 0.2 * Recall@k`. SID-level selection ranks the target Semantic ID before expanding collisions to items; it does not regenerate IDs or use test data. Use `model_selection="last"` only when the final epoch is intentionally required or no validation split is available. + +After fitting, tokenizer and collision diagnostics are available through `tokenizer_time_`, `sid_hash_`, `sid_digit_utilization_`, `sid_digit_entropy_`, `sid_collision_count_`, `sid_collision_group_count_`, and `sid_max_collision_size_`. Training records `training_time_` and `loss_history_`; scoring records `last_decode_time_` and `last_decode_diagnostics_`. + +## Paper and released-code differences + +The paper and released repository differ in several consequential details. Cornac keeps these choices visible rather than silently blending them into one recipe. + +| Behavior | Paper | Released repository | Cornac control | +| -------------------------- | ---------------------------------------- | ---------------------------------------- | --------------------------------- | +| Multi-view loss | Mean within each view, then across views | Pool all masked tokens across views | `view_loss_reduction` | +| Final CPD step | Global beam selection through completion | Greedy final code per active branch | `scoring="paper"` or `"released"` | +| Test beam | Sports/Beauty/Toys: 128/256/128 | Shared default 256 | Dataset-specific `beam_size` | +| Maximum epochs | 100 | Commands inherit 200 | `n_epochs` | +| Beauty label smoothing | 0.1 | Command passes 0.2 | `label_smoothing` | +| Long-sequence augmentation | Described as all contiguous subsequences | Prefix targets only through `max_len=50` | Cornac follows the released rows | +| Semantic-ID collisions | Item resolution is underspecified | Last item overwrites earlier items | `collision_policy` | + +The backbone follows the audited released architecture: pre-normalized bias-free attention, normal initialization with standard deviation 0.02, a final normalization shared by the encoder and decoder, zeroed padded encoder states, and decoder cross-attention over those states without a padding mask. The Sports backbone contains 5,601,280 parameters. Cornac omits three behaviorally unused BOS/EOS/PAD embedding rows, accounting for the released model's additional 768 parameters. + +PSE uses the released 32-thread FAISS setting and FAISS's default clustering seeds; the model seed is not substituted for FAISS's defaults. PCA and FAISS artifacts are dependency-version sensitive, so algorithmically equivalent environments may still generate different code tables. + +## Sports reproducibility study + +### Scope and protocol + +The completed controlled study focuses on the Amazon Reviews 2014 Sports and Outdoors 5-core dataset used by the paper. The processed split contains 35,598 users, 18,357 items, 152,346 released-style training examples, and 35,598 validation and test cases each. Every controlled Cornac run uses the same split and frozen released Semantic-ID table, so differences after tokenization come from model initialization, minibatch and masking order, dropout, checkpoint selection, and decoding rather than regenerated item codes. + +The Cornac comparison uses declared seeds 2024, 2025, and 2026 on a single NVIDIA A40 per run. It keeps the released pooled-token loss, released CPD, validation beam 32, paper test beam 128, SID-level validation objective, and `collision_policy="all"`. Held-out test metrics are reported only after validation-based checkpoint selection and are not used to choose a seed or checkpoint. + +Metrics in the primary reproduction table are SID-level, matching the released evaluator. Item-expanded metrics are reported separately because collisions make SID retrieval and exact item recommendation different objectives. + +### Reproduction path + +The evaluation separated training changes from checkpoint-only rescoring: + +1. Run the released repository with its documented Sports command. +2. Rescore the same checkpoint with the paper beam to isolate beam width. +3. Retrain with the paper's per-view loss while holding the remaining released recipe fixed. +4. Regenerate embeddings and Semantic IDs without the released-only `Features` metadata field to test the paper metadata interpretation. +5. Run the Cornac adapter with the frozen released IDs, correct SID-level checkpoint selection, and the audited release-fidelity backbone. + +| Source | Recipe | Seed | Selected epoch | Test beam | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | +| ------------------- | ------------------------------------------------- | -----------: | -------------: | --------: | ----------------: | ----------------: | ----------------: | ----------------: | +| Paper | Reported DiffGRM | Not reported | Not reported | 128 | .0363 | .0245 | .0550 | .0305 | +| Released repository | As-released training and decoding | 2024 | 54 | 256 | .0329 | .0223 | .0502 | .0279 | +| Released repository | Same checkpoint, paper beam | 2024 | 54 | 128 | .0329 | .0223 | .0500 | .0278 | +| Released repository | Paper per-view loss; released decoding | 2024 | 44 | 128 | .0337 | .0225 | .0516 | .0283 | +| Released repository | Paper metadata fields; released pooled-token loss | 2024 | 28 | 128 | .0366 | .0242 | .0558 | .0303 | +| Cornac | Release-fidelity SID-selected | 2024 | 49 | 128 | .0324 | .0220 | .0501 | .0277 | +| Cornac | Release-fidelity SID-selected | 2025 | 26 | 128 | .0360 | .0244 | .0545 | .0304 | +| Cornac | Release-fidelity SID-selected | 2026 | 25 | 128 | .0325 | .0218 | .0518 | .0280 | +| Cornac | Release-fidelity mean $\pm$ sample SD | 2024--2026 | -- | 128 | .0337 $\pm$ .0020 | .0227 $\pm$ .0014 | .0521 $\pm$ .0022 | .0287 $\pm$ .0015 | + +### What explains the paper gap + +The as-released run is below the paper on all four Sports metrics; its Recall@10 and NDCG@10 gaps are `-8.73%` and `-8.65%`. Changing only the beam from 256 to the paper's 128 does not close the gap. Paper-style per-view loss improves the selected checkpoint, but Recall@10 and NDCG@10 remain `-6.12%` and `-7.36%` below the paper. + +The largest observed change comes from metadata preprocessing. The released data path includes a `Features` field that is absent from the paper's stated text fields. Removing that field before regenerating embeddings and Semantic IDs puts the single controlled run within `1.5%` of all four paper metrics: Recall@5 and Recall@10 are `+0.84%` and `+1.49%`, while NDCG@5 and NDCG@10 are `-1.42%` and `-0.49%`. This isolates metadata construction as the main observed source of the released-to-paper gap, but the conclusion is based on one seed. + +### SID selection and collision effects + +The frozen Sports catalog contains 18,357 items but only 15,448 unique Semantic IDs. There are 2,909 item-to-ID collisions across 1,394 collision groups, with a maximum group size of 47. The released last-item reverse map discards all but one item from each collided ID, whereas the controlled Cornac runs retain every collided item through `collision_policy="all"`. + +Early Cornac integration runs selected checkpoints after expanding decoded IDs to items. That item-level criterion is not equivalent to the released SID-level validation objective when IDs collide. Retraining the same three seeds with SID-level selection raises the mean of every SID test metric by `2.77%`--`3.96%` and reduces their observed sample standard deviations by `66.66%`--`74.43%`. The corresponding item-expanded means decrease by `3.87%`--`4.90%`, demonstrating a real objective tradeoff rather than a universally better checkpoint. + +This correction also illustrates why the README reports both endpoints explicitly. SID-level metrics measure recovery of the target code and are directly comparable with the released evaluator; item-level metrics measure exact catalog recommendation after resolving collisions. + +### Release-fidelity comparison + +The final architecture was chosen through a controlled comparison against the initial corrected Cornac backbone. Both variants used the same split, frozen IDs, seeds, optimization settings, released scoring, SID-level validation selector, and collision policy. The release-fidelity variant changed the attention and normalization layout, initialization, and cross-attention padding behavior to match the audited release. The frozen IDs intentionally bypassed PSE, while released FAISS behavior and last-item collision resolution were validated separately. + +| Metric | Corrected Cornac baseline | Release-fidelity implementation | Relative change | +| ------------------------ | ------------------------: | ------------------------------: | --------------: | +| SID validation objective | .033613 | .033842 | +0.68% | +| SID Recall@5 | .032745 | .033654 | +2.77% | +| SID NDCG@5 | .021557 | .022734 | +5.46% | +| SID Recall@10 | .049881 | .052147 | +4.54% | +| SID NDCG@10 | .027062 | .028676 | +5.96% | +| Item Recall@5 | .015844 | .015937 | +0.59% | +| Item NDCG@5 | .010761 | .011201 | +4.09% | +| Item Recall@10 | .024543 | .024571 | +0.11% | +| Item NDCG@10 | .013538 | .013955 | +3.07% | +| Valid-path fraction | 95.87% | 95.74% | -0.13% | +| Mean stop epoch | 69.0 | 48.3 | -29.95% | +| Mean training time | 107.8 min | 61.4 min | -42.99% | +| Mean test decode time | 191.5 s | 154.7 s | -19.21% | + +The release-fidelity implementation improves the mean validation objective and all eight SID/item ranking metrics while preserving the valid-path rate. It also stops earlier and completes sooner. The runtime difference is partly caused by earlier early stopping, so it is not evidence of a pure per-epoch optimization. The release-fidelity sample standard deviation is higher for all four SID test metrics, so the higher means should not be described as reduced seed sensitivity. + +Relative to the paper, the retained three-seed mean is close but does not match the reported Sports result: + +| Metric | Paper | Cornac mean | Mean gap | Seed 2025 | Seed gap | +| --------- | ----: | ----------: | -------: | --------: | -------: | +| Recall@5 | .0363 | .033654 | -7.29% | .036013 | -0.79% | +| NDCG@5 | .0245 | .022734 | -7.21% | .024388 | -0.46% | +| Recall@10 | .0550 | .052147 | -5.19% | .054497 | -0.91% | +| NDCG@10 | .0305 | .028676 | -5.98% | .030353 | -0.48% | + +Seed 2025 also has the highest validation score among the three declared seeds and is within 1% of every paper metric. The aggregate remains the primary reproduction result because the paper does not report its seed or variance, and test performance was not used to select among the Cornac seeds. + +### PSE artifact fidelity + +The PSE audit exactly reproduced all 18,357 released Sports Semantic-ID rows when it began from the released cached PCA matrix and ran the OPQ/PQ stage with FAISS 1.11.0 and scikit-learn 1.7.0. This establishes exact agreement for the changed Cornac PSE stage under the released input and dependency environment. + +Full regeneration from raw content embeddings through PCA did not reproduce the cached PCA-derived IDs, even with those dependency versions. The released determinism check also starts from the cached PCA matrix, so this is an unresolved artifact-provenance boundary rather than an end-to-end PSE reproduction. Newer supported versions of FAISS and scikit-learn can produce different code tables while following the same algorithm; use frozen `item_sids` whenever exact artifact identity matters. + +### Interpretation and remaining limits + +The completed Sports study supports the following conclusions: + +- The Cornac implementation is behaviorally close to the published Sports result, but its three-seed mean remains `5.19%`--`7.29%` lower across the four reported metrics. +- A validation-selected Cornac seed is within 1% of all four paper metrics, but that individual result does not replace the multi-seed aggregate. +- Metadata preprocessing is the main observed explanation for the released-to-paper gap in the controlled single-seed diagnostics. +- SID-level and exact-item evaluation answer different questions when the tokenizer has collisions; neither should be silently substituted for the other. +- The retained backbone is closer to the released architecture and improves the controlled validation objective and mean ranking metrics, but three seeds are insufficient for a precise variance claim. +- Exact PSE reproduction currently requires the released cached PCA matrix or frozen Semantic IDs; raw-to-PCA provenance remains unresolved. + +The study does not yet constitute a full reproduction of the entire paper. The released repository has only a one-seed matched reference, the central `random` and `fixed` ablations have not been run end to end in Cornac, and Beauty and Toys have not received the same three-seed evaluation. + +## Paper-reported Amazon-2014 results + +The following values are references from the paper, not Cornac reproduction claims. + +| Dataset | Model | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | +| ------- | ------- | -------: | -----: | --------: | ------: | +| Sports | RPG | .0314 | .0216 | .0463 | .0263 | +| Sports | DiffGRM | .0363 | .0245 | .0550 | .0305 | +| Beauty | RPG | .0550 | .0381 | .0809 | .0464 | +| Beauty | DiffGRM | .0603 | .0414 | .0876 | .0502 | +| Toys | RPG | .0592 | .0401 | .0869 | .0490 | +| Toys | DiffGRM | .0618 | .0455 | .0834 | .0524 | + +The exported paper configurations use four 256-way digits, one encoder layer, four decoder layers, 100 epochs, and dataset-specific learning rates, label smoothing, model dimensions, and beam widths. + +## References + +- Zhao Liu, Yichen Zhu, Yiqing Yang, Guoping Tang, Rui Huang, Qiang Luo, Xiao Lv, Ruiming Tang, Kun Gai, and Guorui Zhou. [DiffGRM: Diffusion-based Generative Recommendation Model](https://arxiv.org/abs/2510.21805). WWW 2026. +- [Released DiffGRM repository](https://github.com/liuzhao09/DiffGRM), behavior audited at commit `ad7b971c7e525e9fea6fb8e362a5c49dccb2473c`. diff --git a/cornac/models/diffgrm/__init__.py b/cornac/models/diffgrm/__init__.py new file mode 100644 index 00000000..3c8291db --- /dev/null +++ b/cornac/models/diffgrm/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +from .diffgrm_config import ( + DIFFGRM_BEAUTY_CONFIG, + DIFFGRM_CONFIG, + DIFFGRM_SMOKE_CONFIG, + DIFFGRM_SPORTS_CONFIG, + DIFFGRM_TOYS_CONFIG, +) +from .recom_diffgrm import DiffGRM diff --git a/cornac/models/diffgrm/diffgrm.py b/cornac/models/diffgrm/diffgrm.py new file mode 100644 index 00000000..2734a33e --- /dev/null +++ b/cornac/models/diffgrm/diffgrm.py @@ -0,0 +1,545 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""PyTorch modules for DiffGRM. + +This is an independent implementation of the architecture and equations in +the DiffGRM paper. The official research repository does not currently include +a license, so no source code from that repository is incorporated here. +""" + +import math + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class _MultiheadAttention(nn.Module): + """Pre-norm attention layout used by the released DiffGRM backbone.""" + + def __init__(self, d_model, n_head, dropout): + super().__init__() + self.d_model = d_model + self.n_head = n_head + self.head_dim = d_model // n_head + self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.projection = nn.Linear(d_model, d_model) + self.attention_dropout = nn.Dropout(dropout) + self.residual_dropout = nn.Dropout(dropout) + + def forward(self, query, key_value=None, key_padding_mask=None): + batch_size, query_len, _ = query.shape + query_projection = self.qkv(query) + q = query_projection[..., : self.d_model] + if key_value is None: + k = query_projection[..., self.d_model : 2 * self.d_model] + v = query_projection[..., 2 * self.d_model :] + else: + key_value_projection = self.qkv(key_value) + k = key_value_projection[..., self.d_model : 2 * self.d_model] + v = key_value_projection[..., 2 * self.d_model :] + + key_len = k.size(1) + q = q.view(batch_size, query_len, self.n_head, self.head_dim).transpose(1, 2) + k = k.view(batch_size, key_len, self.n_head, self.head_dim).transpose(1, 2) + v = v.view(batch_size, key_len, self.n_head, self.head_dim).transpose(1, 2) + attention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) + if key_padding_mask is not None: + attention = attention.masked_fill( + key_padding_mask[:, None, None, :], -torch.inf + ) + attention = self.attention_dropout(attention.softmax(dim=-1)) + hidden = torch.matmul(attention, v) + hidden = hidden.transpose(1, 2).contiguous().view( + batch_size, query_len, self.d_model + ) + return self.residual_dropout(self.projection(hidden)) + + +class _FeedForward(nn.Module): + def __init__(self, d_model, n_inner, dropout, activation): + super().__init__() + self.input = nn.Linear(d_model, n_inner) + self.output = nn.Linear(n_inner, d_model) + self.dropout = nn.Dropout(dropout) + self.activation = F.gelu if activation == "gelu" else F.relu + + def forward(self, hidden): + return self.dropout(self.output(self.activation(self.input(hidden)))) + + +class _EncoderBlock(nn.Module): + def __init__(self, d_model, n_head, n_inner, dropout, activation, norm_eps): + super().__init__() + self.attention_norm = nn.LayerNorm(d_model, eps=norm_eps) + self.attention = _MultiheadAttention(d_model, n_head, dropout) + self.feed_forward_norm = nn.LayerNorm(d_model, eps=norm_eps) + self.feed_forward = _FeedForward(d_model, n_inner, dropout, activation) + + def forward(self, hidden, padding_mask): + hidden = hidden + self.attention( + self.attention_norm(hidden), key_padding_mask=padding_mask + ) + return hidden + self.feed_forward(self.feed_forward_norm(hidden)) + + +class _DecoderBlock(nn.Module): + def __init__(self, d_model, n_head, n_inner, dropout, activation, norm_eps): + super().__init__() + self.self_attention_norm = nn.LayerNorm(d_model, eps=norm_eps) + self.self_attention = _MultiheadAttention(d_model, n_head, dropout) + self.cross_attention_norm = nn.LayerNorm(d_model, eps=norm_eps) + self.cross_attention = _MultiheadAttention(d_model, n_head, dropout) + self.feed_forward_norm = nn.LayerNorm(d_model, eps=norm_eps) + self.feed_forward = _FeedForward(d_model, n_inner, dropout, activation) + + def forward(self, hidden, memory): + hidden = hidden + self.self_attention(self.self_attention_norm(hidden)) + hidden = hidden + self.cross_attention( + self.cross_attention_norm(hidden), key_value=memory + ) + return hidden + self.feed_forward(self.feed_forward_norm(hidden)) + + +class DiffGRMBackbone(nn.Module): + """Encoder-decoder backbone with on-policy code masking.""" + + def __init__( + self, + n_digit, + codebook_size, + max_len, + d_model=256, + encoder_n_layer=1, + decoder_n_layer=4, + n_head=4, + n_inner=1024, + dropout=0.1, + activation="gelu", + layer_norm_eps=1e-5, + initializer_range=0.02, + masking_strategy="guided", + confidence_method="msp", + random_mask_prob=0.5, + n_views=None, + label_smoothing=0.1, + view_loss_reduction="view_mean", + ): + super().__init__() + self.n_digit = int(n_digit) + self.codebook_size = int(codebook_size) + self.max_len = int(max_len) + self.masking_strategy = masking_strategy + self.confidence_method = confidence_method + self.random_mask_prob = float(random_mask_prob) + self.n_views = self.n_digit if n_views is None else int(n_views) + self.label_smoothing = float(label_smoothing) + self.view_loss_reduction = view_loss_reduction + + self.code_embeddings = nn.Parameter( + torch.empty(self.n_digit, self.codebook_size, d_model) + ) + self.mask_embeddings = nn.Parameter(torch.empty(self.n_digit, d_model)) + self.item_projection = nn.Sequential( + nn.Linear(self.n_digit * d_model, d_model), + nn.ReLU(), + nn.Linear(d_model, d_model), + ) + self.history_positions = nn.Embedding(self.max_len, d_model) + self.embedding_dropout = nn.Dropout(dropout) + self.encoder_blocks = nn.ModuleList( + [ + _EncoderBlock( + d_model, + n_head, + n_inner, + dropout, + activation, + layer_norm_eps, + ) + for _ in range(encoder_n_layer) + ] + ) + self.decoder_blocks = nn.ModuleList( + [ + _DecoderBlock( + d_model, + n_head, + n_inner, + dropout, + activation, + layer_norm_eps, + ) + for _ in range(decoder_n_layer) + ] + ) + self.final_norm = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.register_buffer( + "item_codes", torch.zeros(1, self.n_digit, dtype=torch.long) + ) + self._reset_parameters(initializer_range) + + def _reset_parameters(self, initializer_range): + nn.init.normal_(self.code_embeddings, std=initializer_range) + nn.init.normal_(self.mask_embeddings, std=initializer_range) + nn.init.normal_(self.history_positions.weight, std=initializer_range) + for module in self.modules(): + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, std=initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + def set_item_codes(self, sid_table): + """Set catalog semantic IDs and append one row for history padding.""" + codes = torch.as_tensor(np.asarray(sid_table), dtype=torch.long) + if codes.ndim != 2 or codes.shape[1] != self.n_digit: + raise ValueError( + f"sid_table must have shape (n_items, {self.n_digit})" + ) + if codes.numel() and ( + codes.min().item() < 0 or codes.max().item() >= self.codebook_size + ): + raise ValueError( + f"semantic-ID digits must be in [0, {self.codebook_size})" + ) + pad = torch.zeros(1, self.n_digit, dtype=torch.long) + self.item_codes = torch.cat([codes, pad], dim=0).to( + self.code_embeddings.device + ) + + def encode_history(self, input_iids, attention_mask): + """Encode a right-padded item history.""" + item_codes = self.item_codes[input_iids] + digit_embs = [ + self.code_embeddings[d][item_codes[:, :, d]] + for d in range(self.n_digit) + ] + history = self.item_projection(torch.cat(digit_embs, dim=-1)) + positions = torch.arange(input_iids.size(1), device=input_iids.device) + history = history + self.history_positions(positions).unsqueeze(0) + history = self.embedding_dropout(history) + padding_mask = ~attention_mask.bool() + for block in self.encoder_blocks: + history = block(history, padding_mask) + history = self.final_norm(history) + history = history * attention_mask.unsqueeze(-1) + return history, padding_mask + + def decode_logits(self, memory, memory_padding_mask, partial_codes): + """Predict every digit; ``-1`` denotes a currently masked digit.""" + is_masked = partial_codes < 0 + visible = partial_codes.clamp_min(0) + digit_states = [] + for d in range(self.n_digit): + code_state = self.code_embeddings[d][visible[:, d]] + mask_state = self.mask_embeddings[d].expand_as(code_state) + digit_states.append( + torch.where(is_masked[:, d, None], mask_state, code_state) + ) + hidden = self.embedding_dropout(torch.stack(digit_states, dim=1)) + for block in self.decoder_blocks: + hidden = block(hidden, memory) + hidden = self.final_norm(hidden) + return torch.einsum("bnd,nkd->bnk", hidden, self.code_embeddings) + + def _confidence_order(self, memory, memory_padding_mask, targets): + batch_size = targets.size(0) + if self.masking_strategy == "fixed": + return torch.arange(self.n_digit, device=targets.device).expand( + batch_size, -1 + ) + if self.masking_strategy == "coherent": + noise = torch.rand( + batch_size, self.n_digit, device=targets.device + ) + return noise.argsort(dim=-1) + + fully_masked = torch.full_like(targets, -1) + was_training = self.training + self.eval() + with torch.no_grad(): + probabilities = self.decode_logits( + memory, memory_padding_mask, fully_masked + ).softmax(dim=-1) + if self.confidence_method == "entropy": + confidence = ( + probabilities + * probabilities.clamp_min(1e-12).log() + ).sum(dim=-1) + else: + confidence = probabilities.max(dim=-1).values + if was_training: + self.train() + return confidence.argsort(dim=-1, stable=True) + + def training_masks(self, memory, memory_padding_mask, targets): + """Return nested OCN masks, hardest digit first.""" + if self.masking_strategy == "random": + masks = ( + torch.rand( + targets.size(0), + self.n_views, + self.n_digit, + device=targets.device, + ) + < self.random_mask_prob + ) + no_mask = ~masks.any(dim=-1) + masks[:, :, 0] |= no_mask + return masks + + order = self._confidence_order(memory, memory_padding_mask, targets) + rank = torch.empty_like(order) + rank.scatter_( + 1, + order, + torch.arange(self.n_digit, device=targets.device).expand_as(order), + ) + counts = torch.arange( + 1, self.n_views + 1, device=targets.device + ).clamp_max(self.n_digit) + return rank[:, None, :] < counts[None, :, None] + + def forward(self, input_iids, attention_mask, target_iids): + """Average cross entropy over the masked digits in all OCN views.""" + memory, padding_mask = self.encode_history(input_iids, attention_mask) + masks = self.training_masks(memory, padding_mask, target_iids) + batch_size, n_views, _ = masks.shape + partial = target_iids[:, None, :].expand(-1, n_views, -1).clone() + partial[masks] = -1 + + memory = memory.repeat_interleave(n_views, dim=0) + padding_mask = padding_mask.repeat_interleave(n_views, dim=0) + partial = partial.reshape(batch_size * n_views, self.n_digit) + logits = self.decode_logits(memory, padding_mask, partial) + labels = ( + target_iids[:, None, :] + .expand(-1, n_views, -1) + .reshape(batch_size * n_views, self.n_digit) + ) + mask = masks.reshape(batch_size * n_views, self.n_digit) + token_losses = F.cross_entropy( + logits.reshape(-1, self.codebook_size), + labels.reshape(-1), + reduction="none", + label_smoothing=self.label_smoothing, + ).reshape(batch_size * n_views, self.n_digit) + if self.view_loss_reduction == "token_mean": + return token_losses[mask].mean() + return ( + (token_losses * mask).sum(dim=-1) + / mask.sum(dim=-1).clamp_min(1) + ).mean() + + +@torch.no_grad() +def cpd_decode_batch( + model, + memory, + memory_padding_mask, + beam_size, + catalog_codes=None, + valid_code_set=None, + constrained=False, + digit_order=None, + greedy_final=False, + return_diagnostics=False, +): + """Batched confidence-prioritized decoding over digit/code assignments. + + At every step, every still-masked digit competes globally. ``constrained`` + additionally removes partial assignments that cannot lead to a catalog + semantic ID. Complete duplicate IDs are collapsed by maximum path score. + """ + device = memory.device + batch_size = memory.size(0) + n_digit = model.n_digit + codebook_size = model.codebook_size + beams = torch.full( + (batch_size, 1, n_digit), -1, dtype=torch.long, device=device + ) + beam_scores = torch.zeros(batch_size, 1, device=device) + catalog = None + if constrained and catalog_codes is not None: + catalog = torch.as_tensor( + catalog_codes, dtype=torch.long, device=device + ) + + if digit_order is not None: + digit_order = tuple(int(d) for d in digit_order) + if sorted(digit_order) != list(range(n_digit)): + raise ValueError("digit_order must be a permutation of all digits") + + for step in range(n_digit): + n_beam = beams.size(1) + expanded_memory = ( + memory[:, None] + .expand(-1, n_beam, -1, -1) + .reshape( + batch_size * n_beam, + memory.size(1), + memory.size(2), + ) + ) + expanded_padding = ( + memory_padding_mask[:, None] + .expand(-1, n_beam, -1) + .reshape(batch_size * n_beam, memory_padding_mask.size(1)) + ) + logits = model.decode_logits( + expanded_memory, + expanded_padding, + beams.reshape(batch_size * n_beam, n_digit), + ).reshape(batch_size, n_beam, n_digit, codebook_size) + candidate_scores = logits.log_softmax(dim=-1) + candidate_scores.masked_fill_(beams[..., None] >= 0, -torch.inf) + if digit_order is not None: + allowed_digit = digit_order[step] + for digit in range(n_digit): + if digit != allowed_digit: + candidate_scores[:, :, digit] = -torch.inf + + if constrained: + if catalog is None: + raise ValueError("catalog_codes are required for constrained CPD") + for row in range(batch_size): + for branch in range(n_beam): + compatible = torch.ones( + catalog.size(0), dtype=torch.bool, device=device + ) + for digit in range(n_digit): + if beams[row, branch, digit] >= 0: + compatible &= ( + catalog[:, digit] + == beams[row, branch, digit] + ) + for digit in range(n_digit): + if beams[row, branch, digit] < 0: + allowed = catalog[compatible, digit].unique() + disallowed = torch.ones( + codebook_size, + dtype=torch.bool, + device=device, + ) + disallowed[allowed] = False + candidate_scores[ + row, branch, digit, disallowed + ] = -torch.inf + + candidate_scores = candidate_scores + beam_scores[:, :, None, None] + if greedy_final and step == n_digit - 1: + per_parent, per_parent_index = candidate_scores.reshape( + batch_size, n_beam, -1 + ).max(dim=-1) + keep = min(int(beam_size), n_beam) + top_scores, parent = per_parent.topk(keep, dim=-1) + remainder = per_parent_index.gather(1, parent) + digit = remainder // codebook_size + code = remainder % codebook_size + beams = beams.gather( + 1, parent[:, :, None].expand(-1, -1, n_digit) + ).clone() + beams.scatter_(2, digit[:, :, None], code[:, :, None]) + beam_scores = top_scores + continue + + flat = candidate_scores.reshape(batch_size, -1) + keep = min(int(beam_size), flat.size(1)) + top_scores, top_indices = flat.topk(keep, dim=-1) + parent = top_indices // (n_digit * codebook_size) + remainder = top_indices % (n_digit * codebook_size) + digit = remainder // codebook_size + code = remainder % codebook_size + beams = beams.gather( + 1, parent[:, :, None].expand(-1, -1, n_digit) + ).clone() + beams.scatter_(2, digit[:, :, None], code[:, :, None]) + beam_scores = top_scores + + catalog_set = valid_code_set + if catalog_set is None and catalog_codes is not None: + catalog_set = { + tuple(row) for row in np.asarray(catalog_codes).tolist() + } + batch_codes, batch_scores, batch_diagnostics = [], [], [] + for row in range(batch_size): + complete = {} + for codes, score in zip( + beams[row].cpu().tolist(), beam_scores[row].cpu().tolist() + ): + key = tuple(codes) + if -1 in key or not np.isfinite(score): + continue + complete[key] = max(complete.get(key, -float("inf")), score) + best = { + key: score + for key, score in complete.items() + if catalog_set is None or key in catalog_set + } + ranked = sorted(best.items(), key=lambda pair: pair[1], reverse=True) + batch_codes.append([codes for codes, _ in ranked]) + batch_scores.append([score for _, score in ranked]) + batch_diagnostics.append( + { + "complete_paths": len(beams[row]), + "unique_complete_sids": len(complete), + "valid_sids": len(best), + "invalid_sids": len(complete) - len(best), + "duplicate_paths": len(beams[row]) - len(complete), + } + ) + + if return_diagnostics: + return batch_codes, batch_scores, batch_diagnostics + return batch_codes, batch_scores + + +@torch.no_grad() +def cpd_decode( + model, + memory, + memory_padding_mask, + beam_size, + catalog_codes=None, + valid_code_set=None, + constrained=False, + digit_order=None, + greedy_final=False, + return_diagnostics=False, +): + """Single-history wrapper around :func:`cpd_decode_batch`.""" + if memory.size(0) != 1: + raise ValueError("cpd_decode expects one encoded history") + result = cpd_decode_batch( + model=model, + memory=memory, + memory_padding_mask=memory_padding_mask, + beam_size=beam_size, + catalog_codes=catalog_codes, + valid_code_set=valid_code_set, + constrained=constrained, + digit_order=digit_order, + greedy_final=greedy_final, + return_diagnostics=return_diagnostics, + ) + if return_diagnostics: + codes, scores, diagnostics = result + return codes[0], scores[0], diagnostics[0] + codes, scores = result + return codes[0], scores[0] diff --git a/cornac/models/diffgrm/diffgrm_config.py b/cornac/models/diffgrm/diffgrm_config.py new file mode 100644 index 00000000..ec990ef9 --- /dev/null +++ b/cornac/models/diffgrm/diffgrm_config.py @@ -0,0 +1,99 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Paper-style DiffGRM configurations for Amazon-2014 experiments.""" + +DIFFGRM_CONFIG = dict( + n_digit=4, + codebook_size=256, + pca_dim=256, + max_len=50, + min_history=2, + encoder_n_layer=1, + decoder_n_layer=4, + n_inner=1024, + dropout=0.1, + masking_strategy="guided", + confidence_method="msp", + n_views=4, + view_loss_reduction="view_mean", + scoring="paper", + n_epochs=100, + batch_size=1024, + weight_decay=0.0, + lr_schedule="cosine", + warmup_steps=10000, + max_grad_norm=1.0, + model_selection="best", + val_k=10, + val_batch_size=32, + val_beam_size=32, + val_eval_every=1, + early_stopping_patience=15, + val_sample=None, +) + +DIFFGRM_SPORTS_CONFIG = dict( + DIFFGRM_CONFIG, + d_model=256, + n_head=4, + learning_rate=0.003, + label_smoothing=0.1, + beam_size=128, + val_eval_start=20, +) + +DIFFGRM_BEAUTY_CONFIG = dict( + DIFFGRM_CONFIG, + d_model=256, + n_head=4, + learning_rate=0.01, + label_smoothing=0.1, + beam_size=256, + val_eval_start=20, +) + +DIFFGRM_TOYS_CONFIG = dict( + DIFFGRM_CONFIG, + d_model=1024, + n_head=8, + learning_rate=0.003, + label_smoothing=0.15, + beam_size=128, + val_eval_start=10, +) + +# Tiny architecture for unit/scheduler smoke tests. It requires precomputed +# ``item_sids`` because 4-way codes are not the paper's 8-bit PSE tokenizer. +DIFFGRM_SMOKE_CONFIG = dict( + n_digit=2, + codebook_size=4, + d_model=32, + encoder_n_layer=1, + decoder_n_layer=1, + n_head=4, + n_inner=64, + dropout=0.0, + max_len=5, + min_history=1, + n_views=2, + n_epochs=1, + batch_size=4, + learning_rate=1e-3, + lr_schedule="constant", + warmup_steps=0, + scoring="catalog", + beam_size=8, + model_selection="last", +) diff --git a/cornac/models/diffgrm/recom_diffgrm.py b/cornac/models/diffgrm/recom_diffgrm.py new file mode 100644 index 00000000..1841d891 --- /dev/null +++ b/cornac/models/diffgrm/recom_diffgrm.py @@ -0,0 +1,742 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import hashlib +import math +import time + +import numpy as np +from tqdm.auto import trange + +from cornac.models.recommender import NextItemRecommender + +from ...utils import get_rng + + +def _diffgrm_num_training_examples(train_set, min_history=2, max_len=50): + """Count prefix-to-next-item examples in a sequential training split.""" + return sum( + max(0, min(len(mapped_ids), max_len + 1) - min_history) + for mapped_ids in train_set.sessions.values() + ) + + +def _diffgrm_session_iter( + train_set, + pad_index, + batch_size=256, + max_len=50, + min_history=2, + rng=None, + shuffle=True, +): + """Yield one training row per eligible next-item prefix.""" + rng = rng if rng is not None else get_rng(None) + examples = [] + for sid, mapped_ids in train_set.sessions.items(): + target_stop = min(len(mapped_ids), max_len + 1) + for target_position in range(min_history, target_stop): + examples.append((sid, target_position)) + if shuffle: + rng.shuffle(examples) + + uir_tuple = train_set.uir_tuple + histories, masks, targets = [], [], [] + for sid, target_position in examples: + mapped_ids = train_set.sessions[sid] + items = np.asarray(uir_tuple[1][mapped_ids], dtype="int64") + history = items[max(0, target_position - max_len) : target_position] + input_iids = np.full(max_len, pad_index, dtype="int64") + input_iids[: len(history)] = history + attention_mask = np.zeros(max_len, dtype="float32") + attention_mask[: len(history)] = 1.0 + + histories.append(input_iids) + masks.append(attention_mask) + targets.append(items[target_position]) + if len(histories) == batch_size: + yield ( + np.asarray(histories, dtype="int64"), + np.asarray(masks, dtype="float32"), + np.asarray(targets, dtype="int64"), + ) + histories, masks, targets = [], [], [] + + if histories: + yield ( + np.asarray(histories, dtype="int64"), + np.asarray(masks, dtype="float32"), + np.asarray(targets, dtype="int64"), + ) + + +class DiffGRM(NextItemRecommender): + """DiffGRM: Diffusion-based Generative Recommendation Model. + + DiffGRM combines parallel OPQ semantic IDs (PSE), on-policy + nested masking (OCN), and confidence-prioritized decoding (CPD). Item + content embeddings are supplied through Cornac's ``FeatureModality``. + Precomputed semantic IDs may instead be passed through ``item_sids`` for + artifact-reproduction and tokenizer-controlled experiments. + + Parameters + ---------- + n_digit: int, default: 4 + Number of semantic-ID digits. + codebook_size: int, default: 256 + Number of codes per digit. PSE uses 8-bit FAISS PQ and therefore + requires 256. Smaller codebooks remain useful with ``item_sids``. + pca_dim: int, default: 256 + Whitened PCA dimension before OPQ. It is reduced, if necessary, to the + largest valid multiple of ``n_digit``. + faiss_omp_num_threads: int, default: 32 + CPU threads used by FAISS while fitting and applying OPQ/PQ, matching + the released configuration. + item_sids: array-like, optional + Precomputed un-offset semantic IDs with shape ``(n_items, n_digit)``. + d_model, encoder_n_layer, decoder_n_layer, n_head, n_inner: + Transformer architecture. Paper-style dataset configs are exported + next to this class. + max_len: int, default: 50 + Maximum number of history items. + min_history: int, default: 2 + Shortest prefix used as a training example. + masking_strategy: {'guided', 'random', 'coherent', 'fixed'}, default: 'guided' + ``guided`` is OCN. ``random`` independently masks digits (without + OCN), ``coherent`` uses nested random-order masks (without on-policy), + and ``fixed`` always masks lower digit indices first. + confidence_method: {'msp', 'entropy'}, default: 'msp' + Confidence statistic used to rank the hardest digits. + random_mask_prob: float, default: 0.5 + Per-digit probability for each independent no-OCN random view. + n_views: int, optional + Number of nested masks per target; defaults to ``n_digit``. + scoring: {'released', 'paper', 'catalog', 'fixed'}, default: 'paper' + ``released`` uses the official final-digit greedy completion and then + filters complete IDs to the catalog. ``paper`` follows Equations 8--10 + with global beam selection at every digit. + ``catalog`` additionally constrains every partial assignment. + ``fixed`` is the no-CPD, fixed-permutation beam control. + fixed_decode_order: sequence of int, optional + Digit permutation used by ``scoring='fixed'``. A seeded permutation is + generated once during fitting when omitted. + collision_policy: {'all', 'first', 'last'}, default: 'all' + Whether all items sharing a semantic ID receive its score, or only the + lowest- or highest-index item. ``last`` matches the released reverse + mapping's overwrite behavior; ``all`` preserves collisions explicitly. + model_selection: {'last', 'best'}, default: 'best' + ``best`` selects by SID-level + ``0.8 * NDCG@k + 0.2 * Recall@k``. + + References + ---------- + Liu et al. (2026). DiffGRM: Diffusion-based Generative Recommendation + Model. WWW. https://arxiv.org/abs/2510.21805 + """ + + def __init__( + self, + name="DiffGRM", + n_digit=4, + codebook_size=256, + pca_dim=256, + faiss_omp_num_threads=32, + feature_standardize=False, + normalize_after_pca=True, + item_sids=None, + d_model=256, + encoder_n_layer=1, + decoder_n_layer=4, + n_head=4, + n_inner=1024, + dropout=0.1, + activation="gelu", + layer_norm_eps=1e-5, + initializer_range=0.02, + max_len=50, + min_history=2, + masking_strategy="guided", + confidence_method="msp", + random_mask_prob=0.5, + n_views=None, + label_smoothing=0.1, + view_loss_reduction="view_mean", + n_epochs=20, + learning_rate=0.003, + weight_decay=0.0, + batch_size=256, + max_grad_norm=1.0, + lr_schedule="cosine", + warmup_steps=10000, + scoring="paper", + beam_size=128, + fixed_decode_order=None, + collision_policy="all", + model_selection="best", + val_k=10, + val_batch_size=32, + val_beam_size=None, + val_eval_start=1, + val_eval_every=1, + early_stopping_patience=None, + val_sample=2000, + device="auto", + trainable=True, + verbose=False, + seed=None, + ): + super().__init__(name=name, trainable=trainable, verbose=verbose) + if n_digit <= 0 or codebook_size <= 1: + raise ValueError("n_digit must be positive and codebook_size > 1") + if faiss_omp_num_threads <= 0: + raise ValueError("faiss_omp_num_threads must be positive") + if encoder_n_layer <= 0 or decoder_n_layer <= 0: + raise ValueError("encoder_n_layer and decoder_n_layer must be positive") + if d_model % n_head != 0: + raise ValueError("d_model must be divisible by n_head") + if max_len <= 0 or min_history <= 0: + raise ValueError("max_len and min_history must be positive") + if n_views is not None and not 1 <= n_views <= n_digit: + raise ValueError("n_views must be between 1 and n_digit") + if masking_strategy not in ("guided", "random", "coherent", "fixed"): + raise ValueError( + "masking_strategy must be 'guided', 'random', 'coherent', or 'fixed'" + ) + if confidence_method not in ("msp", "entropy"): + raise ValueError("confidence_method must be 'msp' or 'entropy'") + if not 0.0 < random_mask_prob <= 1.0: + raise ValueError("random_mask_prob must be in (0, 1]") + if view_loss_reduction not in ("view_mean", "token_mean"): + raise ValueError("view_loss_reduction must be 'view_mean' or 'token_mean'") + if scoring not in ("released", "paper", "catalog", "fixed"): + raise ValueError( + "scoring must be 'released', 'paper', 'catalog', or 'fixed'" + ) + if collision_policy not in ("all", "first", "last"): + raise ValueError("collision_policy must be 'all', 'first', or 'last'") + if fixed_decode_order is not None and sorted(fixed_decode_order) != list( + range(n_digit) + ): + raise ValueError("fixed_decode_order must be a permutation of all digits") + if lr_schedule not in ("constant", "cosine"): + raise ValueError("lr_schedule must be 'constant' or 'cosine'") + if model_selection not in ("last", "best"): + raise ValueError("model_selection must be 'last' or 'best'") + if val_batch_size <= 0 or val_eval_start <= 0 or val_eval_every <= 0: + raise ValueError( + "val_batch_size, val_eval_start, and val_eval_every must be positive" + ) + if early_stopping_patience is not None and early_stopping_patience <= 0: + raise ValueError("early_stopping_patience must be positive or None") + if n_epochs <= 0 or batch_size <= 0 or beam_size <= 0 or val_k <= 0: + raise ValueError( + "n_epochs, batch_size, beam_size, and val_k must be positive" + ) + if val_beam_size is not None and val_beam_size <= 0: + raise ValueError("val_beam_size must be positive or None") + + self.n_digit = n_digit + self.codebook_size = codebook_size + self.pca_dim = pca_dim + self.faiss_omp_num_threads = faiss_omp_num_threads + self.feature_standardize = feature_standardize + self.normalize_after_pca = normalize_after_pca + self.item_sids = item_sids + self.d_model = d_model + self.encoder_n_layer = encoder_n_layer + self.decoder_n_layer = decoder_n_layer + self.n_head = n_head + self.n_inner = n_inner + self.dropout = dropout + self.activation = activation + self.layer_norm_eps = layer_norm_eps + self.initializer_range = initializer_range + self.max_len = max_len + self.min_history = min_history + self.masking_strategy = masking_strategy + self.confidence_method = confidence_method + self.random_mask_prob = random_mask_prob + self.n_views = n_digit if n_views is None else n_views + self.label_smoothing = label_smoothing + self.view_loss_reduction = view_loss_reduction + self.n_epochs = n_epochs + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.batch_size = batch_size + self.max_grad_norm = max_grad_norm + self.lr_schedule = lr_schedule + self.warmup_steps = warmup_steps + self.scoring = scoring + self.beam_size = beam_size + self.fixed_decode_order = fixed_decode_order + self.collision_policy = collision_policy + self.model_selection = model_selection + self.val_k = val_k + self.val_batch_size = val_batch_size + self.val_beam_size = val_beam_size + self.val_eval_start = val_eval_start + self.val_eval_every = val_eval_every + self.early_stopping_patience = early_stopping_patience + self.val_sample = val_sample + self.device = device + self.seed = seed + self.rng = get_rng(seed) + + def _get_item_features(self): + item_feature = getattr(self.train_set, "item_feature", None) + features = getattr(item_feature, "features", None) + if features is None: + raise ValueError( + "DiffGRM requires item content embeddings unless item_sids are " + "provided. Attach FeatureModality through " + "NextItemEvaluation.from_splits(..., item_feature=...)." + ) + if features.shape[0] < self.total_items: + raise ValueError( + f"item_feature has {features.shape[0]} rows but " + f"{self.total_items} items are known" + ) + features = np.asarray(features[: self.total_items], dtype="float32") + if not np.isfinite(features).all(): + raise ValueError("item_feature contains NaN or infinite values") + return features + + def _pse_tokenize(self, features, train_mask): + """Whitened PCA followed by position-sensitive OPQ/PQ codes.""" + import faiss + from sklearn.decomposition import PCA + + if self.codebook_size != 256: + raise ValueError( + "PSE uses 8-bit PQ; codebook_size must be 256 unless " + "precomputed item_sids are supplied" + ) + train_features = features[train_mask] + if len(train_features) < 2: + raise ValueError("PSE needs at least two training items") + if len(train_features) < self.codebook_size: + raise ValueError( + "PSE needs at least codebook_size training items; provide " + "precomputed item_sids for smaller diagnostics" + ) + + if self.feature_standardize: + self.feature_mean_ = train_features.mean(axis=0) + self.feature_std_ = train_features.std(axis=0) + self.feature_std_[self.feature_std_ == 0] = 1.0 + features = (features - self.feature_mean_) / self.feature_std_ + train_features = features[train_mask] + + requested_dim = features.shape[1] if self.pca_dim <= 0 else self.pca_dim + n_components = min(requested_dim, features.shape[1], len(train_features) - 1) + n_components -= n_components % self.n_digit + if n_components < self.n_digit: + raise ValueError( + "PCA output dimension must be at least n_digit and divisible by it" + ) + self.pca_ = PCA( + n_components=n_components, + whiten=True, + random_state=self.seed, + ) + features = self.pca_.fit(train_features).transform(features) + features = np.asarray(features, dtype="float32") + if self.normalize_after_pca: + norms = np.linalg.norm(features, axis=1, keepdims=True) + features = features / np.maximum(norms, 1e-12) + + features = np.ascontiguousarray(features, dtype="float32") + train_features = np.ascontiguousarray(features[train_mask], dtype="float32") + factory = f"OPQ{self.n_digit},IVF1,PQ{self.n_digit}x8" + faiss.omp_set_num_threads(self.faiss_omp_num_threads) + index = faiss.index_factory(n_components, factory, faiss.METRIC_INNER_PRODUCT) + index_ivf = faiss.downcast_index(faiss.extract_index_ivf(index)) + index.train(train_features) + index.add(features) + + inverted = index_ivf.invlists + list_size = inverted.list_size(0) + if list_size != self.total_items: + raise RuntimeError( + f"FAISS encoded {list_size} items, expected {self.total_items}" + ) + code_size = inverted.code_size + codes = faiss.rev_swig_ptr( + inverted.get_codes(0), list_size * code_size + ).reshape(list_size, code_size)[:, : self.n_digit] + ids = faiss.rev_swig_ptr(inverted.get_ids(0), list_size).copy() + sid_table = np.empty((self.total_items, self.n_digit), dtype="int64") + sid_table[ids] = codes.astype("int64") + return sid_table + + def _prepare_semantic_ids(self): + if self.item_sids is not None: + sid_table = np.asarray(self.item_sids, dtype="int64") + expected = (self.total_items, self.n_digit) + if sid_table.shape != expected: + raise ValueError( + f"item_sids must have shape {expected}, got {sid_table.shape}" + ) + if sid_table.size and ( + sid_table.min() < 0 or sid_table.max() >= self.codebook_size + ): + raise ValueError( + f"item_sids digits must be in [0, {self.codebook_size})" + ) + return sid_table.copy() + + features = self._get_item_features() + train_mask = self._training_item_mask() + return self._pse_tokenize(features, train_mask) + + def _training_item_mask(self): + """Items exposed by the released prefix-augmentation training rows.""" + train_mask = np.zeros(self.total_items, dtype=bool) + item_indices = self.train_set.uir_tuple[1] + for mapped_ids in self.train_set.sessions.values(): + stop = min(len(mapped_ids), self.max_len + 1) + if stop <= self.min_history: + continue + items = np.asarray(item_indices[mapped_ids[:stop]], dtype="int64") + train_mask[items] = True + return train_mask + + def _build_model(self): + from .diffgrm import DiffGRMBackbone + + model = DiffGRMBackbone( + n_digit=self.n_digit, + codebook_size=self.codebook_size, + max_len=self.max_len, + d_model=self.d_model, + encoder_n_layer=self.encoder_n_layer, + decoder_n_layer=self.decoder_n_layer, + n_head=self.n_head, + n_inner=self.n_inner, + dropout=self.dropout, + activation=self.activation, + layer_norm_eps=self.layer_norm_eps, + initializer_range=self.initializer_range, + masking_strategy=self.masking_strategy, + confidence_method=self.confidence_method, + random_mask_prob=self.random_mask_prob, + n_views=self.n_views, + label_smoothing=self.label_smoothing, + view_loss_reduction=self.view_loss_reduction, + ).to(self.device_) + model.set_item_codes(self.sid_table) + return model + + def _make_scheduler(self, torch, optimizer): + if self.lr_schedule == "constant": + return None + n_examples = _diffgrm_num_training_examples( + self.train_set, self.min_history, self.max_len + ) + total_steps = max(1, math.ceil(n_examples / self.batch_size) * self.n_epochs) + + def lr_lambda(step): + if step < self.warmup_steps: + return step / max(1, self.warmup_steps) + progress = (step - self.warmup_steps) / max( + 1, total_steps - self.warmup_steps + ) + return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0))) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + def _val_sessions(self, val_set): + sessions = [] + for [_], [mapped_ids], [items] in val_set.si_iter(batch_size=1, shuffle=False): + if len(items) < 2: + continue + user_idx = int(val_set.uir_tuple[0][mapped_ids[0]]) + sessions.append((user_idx, [int(i) for i in items])) + if self.val_sample is not None and len(sessions) > self.val_sample: + indices = self.rng.choice(len(sessions), self.val_sample, replace=False) + sessions = [sessions[i] for i in sorted(indices)] + return sessions + + def _validation_score(self, sessions): + import torch + + from .diffgrm import cpd_decode_batch + + ndcg_values, recall_values = [], [] + self._ensure_device(torch) + self.model.eval() + for start in range(0, len(sessions), self.val_batch_size): + batch = sessions[start : start + self.val_batch_size] + input_iids = np.full( + (len(batch), self.max_len), self.pad_idx, dtype="int64" + ) + attention_mask = np.zeros((len(batch), self.max_len), dtype="float32") + for row, (_, items) in enumerate(batch): + history = items[:-1][-self.max_len :] + input_iids[row, : len(history)] = history + attention_mask[row, : len(history)] = 1.0 + inputs = torch.as_tensor(input_iids, dtype=torch.long, device=self.device_) + masks = torch.as_tensor( + attention_mask, dtype=torch.float32, device=self.device_ + ) + with torch.no_grad(): + memory, padding_mask = self.model.encode_history(inputs, masks) + batch_codes, _, batch_diagnostics = cpd_decode_batch( + self.model, + memory, + padding_mask, + beam_size=( + self.beam_size + if self.val_beam_size is None + else self.val_beam_size + ), + catalog_codes=self.sid_table, + valid_code_set=self.sid_to_items_, + constrained=self.scoring == "catalog", + digit_order=self.fixed_decode_order_ + if self.scoring == "fixed" + else None, + greedy_final=self.scoring == "released", + return_diagnostics=True, + ) + self.last_decode_diagnostics_ = batch_diagnostics[-1] + for (_, items), codes in zip(batch, batch_codes): + target = items[-1] + if target >= self.total_items: + continue + target_sid = tuple(int(digit) for digit in self.sid_table[target]) + rank = next( + ( + index + for index, sid in enumerate(codes) + if tuple(int(digit) for digit in sid) == target_sid + ), + None, + ) + hit = rank is not None and rank < self.val_k + recall_values.append(float(hit)) + ndcg_values.append(1.0 / np.log2(rank + 2) if hit else 0.0) + if not ndcg_values: + return 0.0 + return 0.8 * float(np.mean(ndcg_values)) + 0.2 * float(np.mean(recall_values)) + + def _fit_model(self, torch, val_set): + optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=self.learning_rate, + weight_decay=self.weight_decay, + ) + scheduler = self._make_scheduler(torch, optimizer) + select_best = self.model_selection == "best" and val_set is not None + val_sessions = self._val_sessions(val_set) if select_best else [] + best_state, best_value = None, -float("inf") + non_improving = 0 + self.loss_history_ = [] + + progress = trange( + 1, self.n_epochs + 1, disable=not self.verbose, desc="DiffGRM" + ) + for epoch in progress: + self.current_epoch = epoch + self.model.train() + total_loss, n_batches = 0.0, 0 + for input_iids, attention_mask, target_iids in _diffgrm_session_iter( + self.train_set, + pad_index=self.pad_idx, + batch_size=self.batch_size, + max_len=self.max_len, + min_history=self.min_history, + rng=self.rng, + shuffle=True, + ): + inputs = torch.as_tensor( + input_iids, dtype=torch.long, device=self.device_ + ) + masks = torch.as_tensor( + attention_mask, dtype=torch.float32, device=self.device_ + ) + targets = torch.as_tensor( + self.sid_table[target_iids], + dtype=torch.long, + device=self.device_, + ) + optimizer.zero_grad() + loss = self.model(inputs, masks, targets) + loss.backward() + if self.max_grad_norm is not None: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.max_grad_norm + ) + optimizer.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() + n_batches += 1 + progress.set_postfix(loss=total_loss / n_batches) + self.loss_history_.append(total_loss / n_batches) + + if ( + select_best + and epoch >= self.val_eval_start + and (epoch - self.val_eval_start) % self.val_eval_every == 0 + ): + self.model.eval() + value = self._validation_score(val_sessions) + if value > best_value: + best_value = value + non_improving = 0 + self.best_value = value + self.best_epoch = epoch + self.wait = 0 + best_state = { + name: value.detach().cpu().clone() + for name, value in self.model.state_dict().items() + } + else: + non_improving += 1 + self.wait = non_improving + if ( + self.early_stopping_patience is not None + and non_improving >= self.early_stopping_patience + ): + self.stopped_epoch = epoch + break + + if best_state is not None: + self.model.load_state_dict(best_state) + + def fit(self, train_set, val_set=None): + super().fit(train_set, val_set) + if not self.trainable: + return self + + import torch + + torch.manual_seed(0 if self.seed is None else self.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0 if self.seed is None else self.seed) + self.device_ = ( + "cuda" + if self.device == "auto" and torch.cuda.is_available() + else "cpu" + if self.device == "auto" + else self.device + ) + self.pad_idx = self.total_items + self.n_training_examples_ = _diffgrm_num_training_examples( + self.train_set, self.min_history, self.max_len + ) + if self.n_training_examples_ == 0: + raise ValueError( + "DiffGRM found no training prefixes; reduce min_history or " + "provide longer sessions" + ) + tokenizer_start = time.perf_counter() + self.sid_table = self._prepare_semantic_ids() + self.tokenizer_time_ = time.perf_counter() - tokenizer_start + self.sid_hash_ = hashlib.sha256(self.sid_table.tobytes()).hexdigest() + self.sid_to_items_ = {} + for item_idx, codes in enumerate(self.sid_table): + self.sid_to_items_.setdefault(tuple(codes.tolist()), []).append(item_idx) + self.unique_sid_count_ = len(self.sid_to_items_) + self.sid_collision_count_ = self.total_items - self.unique_sid_count_ + collision_sizes = [len(items) for items in self.sid_to_items_.values()] + self.sid_collision_group_count_ = sum(size > 1 for size in collision_sizes) + self.sid_max_collision_size_ = max(collision_sizes, default=0) + self.sid_digit_utilization_ = np.asarray( + [len(np.unique(self.sid_table[:, digit])) for digit in range(self.n_digit)], + dtype="int64", + ) + digit_entropies = [] + for digit in range(self.n_digit): + counts = np.bincount(self.sid_table[:, digit], minlength=self.codebook_size) + probabilities = counts[counts > 0] / counts.sum() + digit_entropies.append( + -float(np.sum(probabilities * np.log(probabilities))) + ) + self.sid_digit_entropy_ = np.asarray(digit_entropies) + self.fixed_decode_order_ = ( + tuple(int(d) for d in self.fixed_decode_order) + if self.fixed_decode_order is not None + else tuple(int(d) for d in get_rng(self.seed).permutation(self.n_digit)) + ) + + self.model = self._build_model() + training_start = time.perf_counter() + self._fit_model(torch, val_set) + self.training_time_ = time.perf_counter() - training_start + self.model.to("cpu").eval() + return self + + def _ensure_device(self, torch): + requested = torch.device(self.device_) + if requested.type == "cuda" and not torch.cuda.is_available(): + requested = torch.device("cpu") + self.device_ = "cpu" + if next(self.model.parameters()).device != requested: + self.model.to(requested) + + def _score_history(self, history_items): + import torch + + from .diffgrm import cpd_decode + + if not history_items: + return np.ones(self.total_items, dtype="float") + self._ensure_device(torch) + history = list(history_items)[-self.max_len :] + input_iids = np.full((1, self.max_len), self.pad_idx, dtype="int64") + input_iids[0, : len(history)] = history + attention_mask = np.zeros((1, self.max_len), dtype="float32") + attention_mask[0, : len(history)] = 1.0 + inputs = torch.as_tensor(input_iids, dtype=torch.long, device=self.device_) + masks = torch.as_tensor( + attention_mask, dtype=torch.float32, device=self.device_ + ) + self.model.eval() + decode_start = time.perf_counter() + with torch.no_grad(): + memory, padding_mask = self.model.encode_history(inputs, masks) + codes, path_scores, diagnostics = cpd_decode( + self.model, + memory, + padding_mask, + beam_size=self.beam_size, + catalog_codes=self.sid_table, + valid_code_set=self.sid_to_items_, + constrained=self.scoring == "catalog", + digit_order=self.fixed_decode_order_ + if self.scoring == "fixed" + else None, + greedy_final=self.scoring == "released", + return_diagnostics=True, + ) + self.last_decode_time_ = time.perf_counter() - decode_start + self.last_decode_diagnostics_ = diagnostics + return self._decoded_item_scores(codes, path_scores) + + def _decoded_item_scores(self, codes, path_scores): + scores = np.full(self.total_items, -1e10, dtype="float") + for sid, path_score in zip(codes, path_scores): + item_indices = self.sid_to_items_[tuple(sid)] + if self.collision_policy == "first": + item_indices = item_indices[:1] + elif self.collision_policy == "last": + item_indices = item_indices[-1:] + scores[item_indices] = path_score + return scores + + def score(self, user_idx, history_items, **kwargs): + return self._score_history(history_items) diff --git a/cornac/models/diffgrm/requirements.txt b/cornac/models/diffgrm/requirements.txt new file mode 100644 index 00000000..357fab86 --- /dev/null +++ b/cornac/models/diffgrm/requirements.txt @@ -0,0 +1,3 @@ +torch>=2.0.0 +faiss-cpu>=1.7.0 +scikit-learn>=1.2.0