Generative Recommender with End-to-End Learnable Item Tokenization

Enze Liu, Bowen Zheng, Cheng Ling, Lantao Hu, Han Li, Wayne Xin Zhao — Renmin University of China & Kuaishou Inc. ArXiv 2409.05546, 2024.

Relations

Builds on: TIGER (Rajput et al., NeurIPS 2023) (no note yet) Concepts used: Generative Recommender Systems — §4.1 (RQ-VAE Semantic IDs), codebook collapse and training objectives

Table of Contents

• 1. The Decoupling Problem
• 2. Architecture
  • 2.1 Item Tokenizer: RQ-VAE
  • 2.2 Generative Recommender: T5 Encoder-Decoder
• 3. Recommendation-Oriented Alignment
  • 3.1 Sequence-Item Alignment (SIA)
  • 3.2 Preference-Semantic Alignment (PSA)
• 4. Alternating Optimization
• 5. Experiments
  • 5.1 Overall Performance
  • 5.2 Ablation
  • 5.3 Generalization to Cold-Start Users
• 6. References

1. The Decoupling Problem 📐

In TIGER and its descendants, item tokenization proceeds in two fully independent stages:

1. Tokenizer training. A collaborative semantic embedding \(\mathbf{z} \in \mathbb{R}^d\) is computed for each item (typically from item co-occurrence statistics). An RQ-VAE is trained to minimize reconstruction loss \(\|\mathbf{z} - \hat{\mathbf{z}}\|^2\) over this embedding. The codebook is frozen after training.

2. Recommender training. The sequence model receives fixed semantic IDs — code tuples that were optimized to reconstruct \(\mathbf{z}\), not to facilitate next-item prediction.

The mismatch is fundamental. The RQ-VAE partitions item space to minimize reconstruction distortion of \(\mathbf{z}\); there is no guarantee that this partition aligns with the structure that the generative recommender finds useful for predicting user behavior. Two items may have very similar \(\mathbf{z}\) (and thus similar or identical codes) yet behave very differently in user sequences — and the recommender has no way to signal this to the tokenizer.

Definition (Suboptimality of Decoupled Tokenization). Let \(\mathcal{C}^*_{\text{rec}}\) denote the code assignment that minimizes next-item prediction loss, and \(\mathcal{C}^*_{\text{recon}}\) the code assignment that minimizes RQ-VAE reconstruction loss. In general:

\[\mathcal{C}^*_{\text{rec}} \neq \mathcal{C}^*_{\text{recon}}\]

because reconstruction loss is defined over the item embedding space, while recommendation quality is defined over user-item interaction patterns. Decoupling forces the recommender to work with \(\mathcal{C}^*_{\text{recon}}\), accepting its suboptimality.

ETEGRec’s answer is end-to-end training with alignment objectives that explicitly communicate recommendation signal to the tokenizer.

2. Architecture

2.1 Item Tokenizer: RQ-VAE

The item tokenizer follows the standard RQ-VAE construction. For item \(x\) with collaborative semantic embedding \(\mathbf{z} \in \mathbb{R}^{d_s}\):

Encoder: \(\mathbf{r}_0 = \text{EncoderT}(\mathbf{z}) \in \mathbb{R}^{d_c}\)

Residual quantization for levels \(l = 1, \ldots, L\):

\[c_l = \arg\max_{k \in [K]} P(k \mid \mathbf{v}_l), \qquad P(k \mid \mathbf{v}_l) = \frac{\exp(-\|\mathbf{v}_l - \mathbf{e}_k^l\|^2)}{\sum_{j=1}^K \exp(-\|\mathbf{v}_l - \mathbf{e}_j^l\|^2)}\]

\[\mathbf{v}_{l+1} = \mathbf{v}_l - \mathbf{e}_{c_l}^l, \qquad \mathbf{v}_1 = \mathbf{r}_0\]

Decoder: \(\tilde{\mathbf{z}} = \text{DecoderT}\!\left(\sum_{l=1}^L \mathbf{e}_{c_l}^l\right)\)

Note that ETEGRec uses a softmax distribution over codebook entries (rather than argmin) to define the quantization assignment \(c_l\). This is equivalent to the standard nearest-neighbor assignment but provides a differentiable approximation for gradient estimation.

Tokenizer loss:

\[\mathcal{L}_{\text{SQ}} = \underbrace{\|\mathbf{z} - \tilde{\mathbf{z}}\|^2}_{\text{reconstruction}} + \sum_{l=1}^L \left[\underbrace{\|\text{sg}[\mathbf{v}_l] - \mathbf{e}_{c_l}^l\|^2}_{\text{codebook}} + \beta \underbrace{\|\mathbf{v}_l - \text{sg}[\mathbf{e}_{c_l}^l]\|^2}_{\text{commitment}}\right], \quad \beta = 0.25\]

2.2 Generative Recommender: T5 Encoder-Decoder

The recommender is a T5-style encoder-decoder Transformer with 4 encoder and 4 decoder layers (hidden size 256, FFN 1536, 6 attention heads). Unlike HSTU’s causal decoder-only design, ETEGRec uses the full encoder-decoder architecture to process the user’s interaction history:

Encoder: \(\mathbf{H}^E = \text{EncoderR}(\mathbf{E}^x) \in \mathbb{R}^{N \times d}\)

where \(\mathbf{E}^x\) is the sequence of item token embeddings from the user’s last \(N\) interactions.

Decoder (autoregressive over target item’s \(L\) tokens):

\[\mathbf{H}^D = \text{DecoderR}(\mathbf{H}^E, \tilde{\mathbf{Y}}) \in \mathbb{R}^{L \times d}\]

Recommendation loss:

\[\mathcal{L}_{\text{REC}} = -\sum_{l=1}^L \log P(Y_l \mid X, Y_{<l})\]

where \(Y_l = c_l^{t+1}\) is the \(l\)-th codebook token of the target item, and \(X\) is the tokenized input sequence.

Figure 1 (Liu et al., 2024): The overall ETEGRec framework. The item tokenizer (bottom) performs residual quantization on item embeddings. The generative recommender (top) is a T5 encoder-decoder that processes the token sequence. SIA (left dashed box, panel c) aligns the encoder’s mean-pooled output \(\mathbf{z}^E\) with the target item’s codebook distributions via KL divergence at each quantization level. PSA (right dashed box, panel d) aligns the decoder’s first hidden state \(\mathbf{h}^D_0\) with the reconstructed item embedding \(\tilde{\mathbf{z}}\) via InfoNCE.

3. Recommendation-Oriented Alignment 🔑

The two alignment objectives are the core contribution. They differ in which representations they align and what signal they use.

3.1 Sequence-Item Alignment (SIA)

Motivation. The encoder’s hidden states \(\mathbf{H}^E\) summarize the user’s interaction history. If the encoder is well-trained, its mean-pooled output should represent “what kind of item the user wants next” — and this representation should produce similar codebook token distributions as the target item’s own embedding.

Construction. Apply mean pooling and an MLP to the encoder output:

\[\mathbf{z}^E = \text{MLP}(\text{MeanPool}(\mathbf{H}^E)) \in \mathbb{R}^{d_c}\]

Run \(\mathbf{z}^E\) through the same RQ-VAE residual quantization process to obtain predicted codebook distributions \(P^l_{\mathbf{z}^E}\) at each level \(l\). The target item’s embedding \(\mathbf{z}\) produces distributions \(P^l_{\mathbf{z}}\) via the same process.

Loss (symmetric KL divergence over all codebook levels):

\[\mathcal{L}_{\text{SIA}} = -\sum_{l=1}^L \left[D_{\text{KL}}\!\left(P^l_{\mathbf{z}} \,\|\, P^l_{\mathbf{z}^E}\right) + D_{\text{KL}}\!\left(P^l_{\mathbf{z}^E} \,\|\, P^l_{\mathbf{z}}\right)\right]\]

The symmetric KL (Jensen–Shannon-like) is used rather than one-directional KL to avoid the degenerate mode-seeking/mean-seeking failure modes of asymmetric divergences.

What SIA teaches the tokenizer. The gradient of \(\mathcal{L}_{\text{SIA}}\) with respect to codebook entries \(\mathbf{e}_k^l\) penalizes the tokenizer for assigning different code distributions to the target item and to the encoder’s prediction of “what the user wants.” Over many training steps, the codebook is pushed toward partitions where items that users actually transition to share similar codes — not just items that are similar in embedding space.

3.2 Preference-Semantic Alignment (PSA)

Motivation. The decoder’s first hidden state \(\mathbf{h}^D_0 \in \mathbb{R}^d\) represents the model’s prediction of user preference before generating any output token. The reconstructed item embedding \(\tilde{\mathbf{z}} \in \mathbb{R}^{d_s}\) is the RQ-VAE’s best approximation of the target item. PSA aligns these two representations via contrastive learning.

Loss (InfoNCE with in-batch negatives):

\[\mathcal{L}_{\text{PSA}} = -\left[\log \frac{\exp(s(\tilde{\mathbf{z}}, \mathbf{h}^D_0) / \tau)}{\sum_{\hat{\mathbf{h}} \in \mathcal{B}} \exp(s(\tilde{\mathbf{z}}, \hat{\mathbf{h}}) / \tau)} + \log \frac{\exp(s(\mathbf{h}^D_0, \tilde{\mathbf{z}}) / \tau)}{\sum_{\hat{\mathbf{z}} \in \mathcal{B}} \exp(s(\mathbf{h}^D_0, \hat{\mathbf{z}}) / \tau)}\right]\]

where \(s(\cdot, \cdot)\) is cosine similarity, \(\tau\) is a temperature, and \(\mathcal{B}\) is the set of all in-batch items. The bidirectional form (two InfoNCE terms) mirrors CLIP’s symmetric contrastive objective.

SIA vs. PSA. The two objectives operate at different levels of the architecture:

PSA provides a coarser but more direct semantic alignment: the decoder’s user-preference vector should be geometrically close to the target item in reconstructed embedding space. SIA provides a finer distributional alignment between the encoder’s prediction and the item’s quantization structure.

Figure 4 (Liu et al., 2024): t-SNE visualization of preference representations (\(\mathbf{h}^D_0\), circles) and semantic representations (\(\tilde{\mathbf{z}}\), stars) on Instruments (left) and Scientific (right). Each color denotes a distinct item group. After PSA training, preference points cluster near their corresponding semantic star, confirming that the decoder’s user-preference vector is geometrically aligned with the reconstructed target-item embedding.

4. Alternating Optimization ⚠️

ETEGRec alternates between two training phases in a cycle of 4 epochs:

Phase 1 — Tokenizer update (1 epoch, recommender frozen):

\[\mathcal{L}_{\text{IT}} = \mathcal{L}_{\text{SQ}} + \mu\, \mathcal{L}_{\text{SIA}} + \lambda\, \mathcal{L}_{\text{PSA}}\]

The RQ-VAE encoder, codebook entries, and decoder are updated. The recommender’s encoder-decoder weights are frozen. Gradients from \(\mathcal{L}_{\text{SIA}}\) reach the codebook via the tokenizer’s residual quantization; gradients from \(\mathcal{L}_{\text{PSA}}\) reach the RQ-VAE decoder via \(\tilde{\mathbf{z}}\).

Phase 2 — Recommender update (3 epochs, tokenizer frozen):

\[\mathcal{L}_{\text{GR}} = \mathcal{L}_{\text{REC}} + \mu\, \mathcal{L}_{\text{SIA}} + \lambda\, \mathcal{L}_{\text{PSA}}\]

The T5 encoder-decoder weights are updated. The RQ-VAE is frozen — semantic IDs are recomputed from the current tokenizer before each recommender epoch. Gradients from \(\mathcal{L}_{\text{SIA}}\) reach the encoder MLP; gradients from \(\mathcal{L}_{\text{PSA}}\) reach the decoder’s cross-attention and output layers.

The 1:3 ratio (one tokenizer epoch per three recommender epochs) reflects the different convergence rates: the tokenizer’s codebook is a discrete structure that converges quickly once the recommender provides a stable target; the recommender benefits from more gradient steps per tokenizer update.

Ablation on alternating training. Removing the alternating scheme (replacing it with simultaneous joint optimization) causes a −17.6% drop in Recall@5 on Instruments. This confirms that the sequential freeze-and-update discipline is necessary — not just a training heuristic — to prevent codebook destabilization.

5. Experiments 📊

5.1 Overall Performance

Evaluation on three Amazon product review datasets:

ETEGRec consistently leads, though margins over TIGER-SAS are modest (3–5% on most metrics). The largest gains appear on Instruments, the smallest on Video Games.

5.2 Ablation

Key observations: - Alternating training is by far the most important component. Removing it costs −17.6% Recall@5. This validates the core claim: end-to-end training requires the alternating discipline. - PSA contributes more than SIA (−2.1% vs. −0.5% on Recall@5). The geometric contrastive alignment provides a stronger gradient signal than the distributional KL. - Removing both alignment losses but keeping alternating training still underperforms the full model by −3.9% — confirming that the alignment objectives add value beyond simply enabling stable alternating updates.

5.3 Generalization to Cold-Start Users

On the 5% of users with the fewest historical interactions (coldest users), ETEGRec maintains its advantage over TIGER-SAS. The improvement is attributed to the decoder’s preference vector \(\mathbf{h}^D_0\) being aligned to item semantics via PSA, which provides a signal even when interaction history is sparse. A user with only 2–3 interactions has a noisy encoder output, but PSA ensures the decoder’s preference representation stays geometrically close to the relevant part of item space.

Figure 3 (Liu et al., 2024): Recall@10 for seen (training-set) and unseen (cold-start) users on Instruments (left) and Scientific (right). ETEGRec (red bars) leads on both seen and unseen splits. The gain on unseen users is larger in absolute terms on Scientific, consistent with the claim that PSA alignment provides a robust preference signal even for sparse interaction histories.

6. References