# RQ-VAE Semantic ID Tokenizer

> The tokenizer that turns an item content embedding into a short tuple of discrete codes, giving items a coarse-to-fine hierarchy a decoder can generate.

Pick as the front half of any semantic-ID recommender. The number of levels and codes sets the id length the downstream decoder has to emit.

- Category: Recommendation
- Paper: Rajput et al. (2023). Recommender Systems with Generative Retrieval. NeurIPS 2023 (https://arxiv.org/abs/2305.05065)
- Layers: 13
- Parameters: 1.25M
- Input shape (batchless): 768
- Output shape: 768
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/rqvae-semantic-id/model.json
- Open on the canvas: https://neurarch.com/?template=rqvae-semantic-id

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | Item content embedding | Input | shape=[768] | 768 |
| 2 | Encoder 768-512 | Linear | outFeatures=512, inFeatures=768 | 512 |
| 3 | ReLU | ReLU |  | 512 |
| 4 | Encoder 512-256 | Linear | outFeatures=256, inFeatures=512 | 256 |
| 5 | ReLU | ReLU |  | 256 |
| 6 | Encoder 256-128 | Linear | outFeatures=128, inFeatures=256 | 128 |
| 7 | Residual VQ (4 levels x 256 codes) | Residual VQ (RVQ) | embedDim=128 | 128 |
| 8 | Decoder 128-256 | Linear | outFeatures=256, inFeatures=128 | 256 |
| 9 | ReLU | ReLU |  | 256 |
| 10 | Decoder 256-512 | Linear | outFeatures=512, inFeatures=256 | 512 |
| 11 | ReLU | ReLU |  | 512 |
| 12 | Decoder 512-768 | Linear | outFeatures=768, inFeatures=512 | 768 |
| 13 | reconstructed embedding | Output |  | 768 |

## Verifier findings

- **info** `deep-no-norm`: 11 layers with no BatchNorm, LayerNorm, or GroupNorm. Without normalization, activations can explode or vanish across layers, causing slow or unstable training. Fix: Add BatchNorm after Conv2d (CV tasks), LayerNorm after attention/FFN (NLP/LLM), or GroupNorm for small batch sizes.

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class ResidualVQ(nn.Module):
    """Residual vector quantization: each level quantizes what the level above
    could not represent, which is what gives a semantic ID its coarse-to-fine
    hierarchy."""

    def __init__(self, num_quantizers: int = 8, codebook_size: int = 1024, embed_dim: int = 256):
        super().__init__()
        self.codebooks = nn.ModuleList(
            [nn.Embedding(codebook_size, embed_dim) for _ in range(num_quantizers)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual, out = x, torch.zeros_like(x)
        for cb in self.codebooks:
            dist = (residual.unsqueeze(-2) - cb.weight).pow(2).sum(-1)
            q = cb(dist.argmin(-1))
            out = out + q
            residual = residual - q
        return x + (out - x).detach()  # straight-through


class RQ_VAESemanticIDTokenizer(nn.Module):
    def __init__(self):
        super().__init__()

        self.linear_1 = nn.Linear(768, 512)
        self.linear_2 = nn.Linear(512, 256)
        self.linear_3 = nn.Linear(256, 128)
        self.residualVQ_1 = ResidualVQ(num_quantizers=4, codebook_size=256, embed_dim=128)
        self.linear_4 = nn.Linear(128, 256)
        self.linear_5 = nn.Linear(256, 512)
        self.linear_6 = nn.Linear(512, 768)

    def forward(self, x):
        # Item content embedding shape: [768]
        linear_e1 = self.linear_1(x)
        relu_a1 = F.relu(linear_e1)
        linear_e2 = self.linear_2(relu_a1)
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md