[feat] SID: add SidRqvae (RQ-VAE) model — STE/Gumbel-Softmax + CLIP

[WhiteSwan1]

Summary

Adds SidRqvae, an end-to-end differentiable RQ-VAE semantic-ID model, on top of the SID foundation already in master (BaseSidModel / ResidualQuantizer from #538, SidRqkmeans from #539). It complements the encoder-free, FAISS-trained SidRqkmeans with a gradient-trained encoder→quantizer→decoder that learns the codebook jointly with a reconstruction objective, and optionally a pair-contrastive (masked InfoNCE) objective.

The model runs on both CPU and GPU (gradient training, no FAISS dependency at train time unless kmeans_init warm-start is requested).

What's in scope

Config-driven losses. SID losses are configured through ModelConfig.losses: a LossConfig carries one sid_loss oneof term per objective (recon_loss, commitment_loss, contrastive_loss), and BaseSidModel registers each as a _loss_modules entry and computes them centrally (mirroring RankModel). Each is a small _Loss module that takes its operands/mask and returns a scalar:

• SidReconLoss (tzrec/loss/sid_recon_loss.py) — per-row reconstruction distance (recon_type ∈ {l2, l1, cos}), masked-mean reduced (the mixed recon+contrastive path scores reconstruction rows only).
• SidCommitmentLoss (tzrec/loss/sid_commitment_loss.py) — VQ-VAE commitment between the encoder output and the per-layer cumulative quantized vectors (latent_weight [w1, w2], commitment_type).
• SidContrastiveLoss (tzrec/loss/sid_contrastive_loss.py) — masked InfoNCE over mixed reconstruction+pair batches, driven by a per-row pair flag. Modality-agnostic (no image/text assumption), with structural masked-logit fill (dtype finfo.min, AMP-safe), learnable per-group temperatures (clamped before exp), and built-in differentiable all-gather for DDP.

Models

• SidRqvae (tzrec/models/sid_rqvae.py): configurable-depth encoder/decoder MLPs (mirrored; framework MLP + a trailing bare Linear so the latent/reconstruction stay unbounded), an N-layer residual vector quantizer, and an optional contrastive dual-encoder path. Knobs: forward_mode ∈ {ste, gumbel_softmax}, normalize_residuals, distance_type, rotation_trick (arXiv:2410.06424), Sinkhorn balanced assignment, and FAISS kmeans_init warm-start (opt-in, default off).
• Inputs flow through the framework EmbeddingGroup / build_input, so SID models consume the same grouped feature tensor as every other model and support multiple content + side-info embeddings per FORGE/PLUM. The main input is the (sole, or first-declared) feature group — auto-detected, nothing to configure; the contrastive path names its paired + pair-flag groups in contrastive_config.

Modules

• ResidualVectorQuantizer (modules/sid/residual_vector_quantizer.py): gradient-trained N-layer residual VQ. STE returns the raw codebook vector and applies one aggregate straight-through on the encoder side (the codebook trains via the commitment loss); Gumbel-Softmax's soft assignment is differentiable directly. Includes the FAISS residual warm-start (faiss_residual_kmeans) with a DDP fail-together broadcast that avoids the rank-0-raises-while-others-hang deadlock.
• VectorQuantizeLayer (modules/sid/vector_quantize.py): single-layer QuantizeLayer with STE / Gumbel-Softmax / Sinkhorn balanced assignment (Sinkhorn auto-disabled under Gumbel; sinkhorn_epsilon guarded > 0 to avoid the exp-kernel overflow).
• Shared faiss_kmeans_fit primitive in modules/sid/kmeans_quantize.py, used by both the RQ-VAE warm-start and the RQ-KMeans offline fit (one-layer FAISS fit with an N >= K guard).

Protos

• loss.proto: SidReconLoss / SidCommitmentLoss / SidContrastiveLoss messages + the sid_loss oneof on LossConfig.
• models/sid_model.proto: SidRqvae, SinkhornConfig, and ContrastiveConfig messages. The main feature group is auto-detected (group_names()[0]), so there is no feature_group field.
• model.proto: sid_rqvae wired into the ModelConfig model oneof.

Test plan

Colocated unit tests for every new/changed module (distributed paths exercised by 2-rank cases inside the same *_test.py), plus an end-to-end integration test:

• sid_rqvae_test.py — model forward/loss/metric, STE + Gumbel, with/without the contrastive path, the fail-fast group/dim guards, and the logit-scale overflow clamp.
• sid_recon_loss_test.py, sid_commitment_loss_test.py, sid_contrastive_loss_test.py — the three loss modules (per-row distance + masked reduction, commitment directions, masked contrastive incl. empty-mask safety, 2-rank all-gather, AMP-safe fill).
• residual_vector_quantizer_test.py — Gumbel grad flow vs the STE codebook gradient (trained via commitment), FAISS warm-start, non-uniform codebooks.
• vector_quantize_test.py — STE / Gumbel / Sinkhorn assignment.
• residual_quantizer_test.py, kmeans_quantize_test.py, residual_kmeans_quantizer_test.py, quantize_layer_test.py — the shared quantizer primitives.
• sid_integration_test.py — RQ-VAE and RQ-KMeans end-to-end (train → eval → predict).

All CPU tests green via python -m tzrec.tests.run; GPU/DDP-only cases gate on CUDA. FAISS-dependent tests skip cleanly when faiss is absent.

[WhiteSwan1]

Thanks for the careful reviews @tiankongdeguiji — here's a consolidated summary of the changes. The smaller fixes are listed first, then the larger refactors and the one suggestion I'd handle a bit differently. All of these changes have been validated through end-to-end experiments and have no impact on model results.

Smaller fixes (done)

• Merged the *_dist_test.py files into the colocated *_test.py and deleted the separate dist files.
• Used config_to_kwargs for the sinkhorn kwargs so the proto and module parameters stay in sync.
• Dropped the incorrect is_train or is_eval phase check.
• Removed the train-time metrics for now; I'll make SID metrics configurable in a follow-up, disabled by default.
• Switched the manually-expanded test cases to parameterized.
• Moved the subclass tests out of the base-class test file.
• Added an end-to-end integration test (RQ-VAE and RQ-KMeans) in sid_integration_test.py.
• Renamed VectorQuantize → VectorQuantizeLayer.
• Inlined the single-use _train_gumbel helper.
• Made temperature a class-init parameter instead of a quantize() argument.
• Switched the encoder/decoder to the framework MLP (plus a bare trailing Linear for the unbounded latent/reconstruction projection, since MLP always activates its last layer); the rewrite is behavior-preserving.
• Replaced the hand-rolled squared-L2 distance with torch.cdist(x, y, p=2).pow(2). I checked it matches the previous form to fp32 noise, including a coincident zero-distance point (finite gradient, no NaN).

Config-driven losses

I added a oneof sid_loss to LossConfig (reconstruction / commitment / contrastive), merged the three reconstruction variants into a single one with a recon_type field, and centralized the loss handling in BaseSidModel (init_loss() + loss()), the way rank_model does. predict() now returns only the raw tensors and loss() computes the configured terms. I also extracted a standalone commitment-loss class and use it from the model, and moved the contrastive loss and its learnable temperatures into init_loss().

Input via build_input + feature groups

The SID models now consume the framework's build_input / EmbeddingGroup path like every other model instead of reading the embedding out of the batch directly, and input_dim is derived from the feature group's total dimension. This is the change you suggested for side-info / multiple content embeddings (FORGE/PLUM): the main feature group can now hold one content embedding or several content + side-info features concatenated. The paired embedding and the pair flag for the contrastive path are feature groups too. I updated the proto, the mock configs, and the unit tests accordingly.

Contrastive loss naming

Renamed MaskedCLIPLoss → MaskedInfoNCELoss and made it modality-agnostic (no image_embed / text_embed); the model now owns the pairing structure and passes the loss a generic pair mask.

On gathering the valid pairs outside the loss — keeping the masking inside now

Three reasons:

1. It's numerically identical. Gather-then-mask produces the exact same result as filter-then-gather in the forward: non-pair columns are masked with the dtype's most-negative finite value (so they drop out of the softmax, exp(...) = 0), and the pair rows are normalized by the valid-pair count. So there's no accuracy difference to gain.

2. The negatives are cross-rank. The loss builds its negatives by all_gather-ing the embeddings across ranks and scoring each local row against the global bank. If I filter the pairs before the gather, each rank gathers a different (ragged) row count — a size mismatch in the collective — and on a step where a rank has no pairs, logit_scale drops out of that rank's autograd graph, which hangs DDP's gradient sync. Keeping the full batch through the gather avoids both.

3. It keeps the loss torch.compile-friendly. The loss is deliberately written branch-free with fixed (B, B_global) shapes. Filtering by the pair mask makes the row count — and every downstream tensor shape — data-dependent, which triggers guard failures / recompiles (or graph breaks) under torch.compile; the masked form keeps the shapes static so the graph traces once.

Based on the above considerations, MaskedInfoNCELoss‘s masking did not make any changes in this PR. If necessary, this module will be carefully evaluated separately in the future PR to ensure that the filter-then-compute runs correctly in cross-rank scene.

To be completed in a follow-up PR

• Unifying the single-batch FAISS warm-start with the reservoir train_offline path: they currently sit in different classes with different responsibilities, so merging them is a larger refactor of where the FAISS fit lives — I'll handle it in a follow-up rather than expand this PR's scope.
• Replacing the single-batch FAISS fit with a proper multi-batch fit (as we discussed); a # TODO marks the spot.
• kmeans_init is kept as an opt-in for now, with documented caveats (it seeds from one batch, and gradient training refines it afterward).

[tiankongdeguiji]

not needed, auto detect group_names()[0]