TorchEasyRec Documentation

RQ-VAE (Residual-Quantized Variational Auto-Encoder) is a Semantic ID (SID) generation model used as a tokenizer for generative recommendation. It quantizes item content or multi-modal embeddings into a sequence of discrete integer codes (code_0, code_1, ..., code_{n-1}).

Core Architecture

• Encoder MLP: Compresses input embeddings to an embed_dim latent vector.
• Residual Vector Quantizer (RVQ): A multi-layer quantizer that maps residuals to the nearest entry in a learnable codebook (implemented as nn.Embedding) at each layer.
• Decoder MLP: Reconstructs the original embedding using the sum of the quantized vectors from all layers.
• Training: The codebook, encoder, and decoder are trained jointly. Gradients are passed through the non-differentiable quantization step using the Straight-Through Estimator (STE).

Key Concepts

• Semantic ID (SID): A tuple of the nearest neighbor indices from each residual codebook layer.
• Contrastive Learning: Supports dual-view contrastive learning (similar to CLIP). By providing a paired embedding and an is_pair flag, the model applies an InfoNCE loss to ensure semantically similar views are mapped to similar SID spaces.
• Comparison with RQKMeans: Unlike SID RQKMeans (which uses FAISS K-Means on CPU), RQVAE is a gradient-based model that supports multi-GPU training.

UltraHSTU is a long-sequence recommendation architecture that implements four orthogonal efficiency optimizations to reduce computational and memory overhead while maintaining or improving accuracy. It is built upon the DlrmHSTU architecture.

Core Optimizations

• Semi-Local Attention (SLA): Limits causal attention to a local window of sla_k1 tokens and a global prefix of sla_k2 tokens. This reduces complexity from $O(L^2)$ to $O(L \cdot K1) + O(L \cdot K2)$. Requires kernel: CUTLASS or kernel: PYTORCH.
• Mid-stack Attention Truncation: Discards the UIH prefix after a specific layer (attn_truncation_split_layer) and keeps only the last attn_truncation_tail_len tokens for subsequent layers. This compresses KV cache and compute requirements.
• Mixture of Transducers (MoT): Runs $N$ parallel HSTUTransducer channels (e.g., click stream, view stream). Each channel has its own STU stack and SLA/truncation settings. Outputs are concatenated across the channel dimension before being fed to a FusionMTLTower.
• Selective Rematerialization: Reduces memory usage during backpropagation by recomputing intermediate tensors. Controlled by recompute_normed_x_in_backward and recompute_uvqk_in_backward (both default to true).

HSTU Match relies on specific feature_groups identified by group_name. The user_tower and item_tower use these names to index their inputs.

Required Groups:

• uih: User interaction history sequence. Type: JAGGED_SEQUENCE. Mandatory.
• candidate: Candidate item sequence (composed of positive and negative sampled items during training). Type: JAGGED_SEQUENCE. Mandatory.

Optional/Conditional Groups:

• contextual: User-side ID features (e.g., user_id). Type: DEEP.
• uih_action: Sequence of user action events (e.g., click, buy). Type: JAGGED_SEQUENCE. Mandatory if uih_preprocessor.action_encoder is configured.
• uih_watchtime: Sequence of user interaction durations. Type: JAGGED_SEQUENCE. Mandatory if action encoder requires watchtime.
• uih_timestamp: Sequence of user interaction timestamps. Type: JAGGED_SEQUENCE. Mandatory if positional_encoder.use_time_encoding=true.
• query_time: A scalar request time per row (must use same units as uih_timestamp). Type: DEEP. If configured, time encoding is calculated as ts_gap = query_time - behavior_timestamp. Otherwise, it defaults to the timestamp of the last UIH behavior.

Note: group_name must remain unchanged as it is used for indexing by the towers.

All feature types in TorchEasyRec share several common configuration parameters:

• feature_name: The name of the feature or its output name.
• embedding_dim: The dimension of the feature embedding. It should be a multiple of 4. A recommended heuristic is embedding_dim = 8 + x^{0.25}, where x is the number of unique values.
• embedding_name: The name of the embedding. If multiple features need to share the same embedding parameters, set their embedding_name to the same value.
• pooling: The pooling method for multi-value feature embeddings. Supported values: sum (default), mean.
• init_fn: The initialization function for feature embeddings. Can be any standard torch.nn.init function (e.g., nn.init.uniform_, a=-0.01, b=0.01).
• default_value: The default value for the feature. If set to "", no default is provided, and empty features will result in zero vectors. This value is applied before bucketize operations.
• separator: The delimiter used for multi-value string inputs. Defaults to \x1d. Using ARRAY types is recommended for better performance.
• fg_encoded_default_value: The default value for FG-encoded data when fg_mode=FG_NONE and data is not encoded using pai-fg.
• trainable: Whether the Embedding Variable is trainable. Defaults to true.
• stub_type: If true, the feature acts only as an intermediate result for FG and is not output as a feature. Note: Cannot be used when fg_mode=FG_NORMAL.
• data_type: The data type for the EmbeddingTable. Supports FP32 (default) and FP16.

The embedding_name_suffix (optional) allows you to append a string to the embedding names within a feature group. This enables different feature groups to use the same underlying features while maintaining independent embedding tables, which is useful for multi-tower architectures.

Suffix Behavior:

• No suffix or empty: Features with the same embedding_name across different groups will share the same embedding table.
• Same suffix across groups: Groups will still share the same embedding table.
• Different suffixes across groups: Each group will use its own independent embedding table.
• In WIDE groups: The final embedding table name follows the pattern <emb_name>_wide_<suffix>.
• In DEEP groups with nested sequences: The suffix is automatically inherited by nested sequence_groups. However, a sequence_group can explicitly set its own embedding_name_suffix to override the parent's value.