STSCQ: Switchable Token-Specific Quantization

• Switchable Token-Specific Codebook Quantization is a method that employs token-specific codebooks and hierarchical grouping to achieve efficient compression of face images.
• It integrates with vector quantization models like VQ-VAE and VQGAN, delivering improved rate–distortion performance and superior identity preservation compared to global codebook methods.
• STSCQ uses a three-stage training process—including shared codebook initialization, token-specific refinement, and identity loss fine-tuning—to achieve significant gains in recognition accuracy at low bpp.

Switchable Token-Specific Codebook Quantization (STSCQ) is a codebook-based compression methodology optimized for face image compression at ultra-low bitrates. Distinguished from global codebook approaches, STSCQ introduces hierarchical image-level grouping and token-specific codebooks, significantly improving rate–distortion performance and face identity preservation. The design enables STSCQ to be integrated into any vector quantization-based representation learning pipeline, such as VQ-VAE, VQGAN, and TiTok, providing substantial gains over fixed-global codebook baselines at the same bits-per-pixel.

1. Model Structure and Pipeline

STSCQ operates atop the conventional "latent-space" autoencoder with vector quantization, consisting of the following principal components:

• Encoder ($\mathrm{Enc}$): Maps the input RGB face image $x \in \mathbb{R}^{256 \times 256 \times 3}$ to a grid of $T$ continuous latent tokens, $z_e = \mathrm{Enc}(x) \in \mathbb{R}^{T \times d}$.
• Routing and Codebook Selection: An image-level routing network $G_\theta$ selects a codebook group for the entire image: $k = G(z_e) \in \{1, \dots, M\}$.
• Token-Specific Quantization: For each token position $t$, a unique codebook $C_{k, t} \in \mathbb{R}^{K \times d}$ is assigned, yielding $z_q^t = \mathrm{Quant}_{C_{k, t}}(z_e^t)$.
• Decoder ($\mathrm{Dec}$): Reconstructs the image from quantized tokens: $\hat{x} = \mathrm{Dec}(\{z_q^t\}_{t=1}^T)$.

This architecture divides codebook capacity across $M$ image-level groups and $T$ token positions, yielding a total of $M \times T$ codebooks, each of size $K$, in contrast to a single monolithic codebook of size $N$.

2. Mathematical Formalism and Losses

The formulation introduces distinct codebooks $C^i_t$ for each group $i \in \{1, \dots, M\}$ and token position $t \in \{1, \dots, T\}$. For a given image, the active codebook group is selected, and each token is quantized via nearest-neighbor:

$$z_q^t = \mathrm{Quant}_{C^i_t}(z_e^t) = \arg\min_{e \in C^i_t} \|z_e^t - e\|_2^2$$

Bit allocation is explicitly determined:

$$\mathrm{bpp} = \frac{T \cdot \lceil \log_2 K \rceil + \lceil \log_2 M \rceil}{H \cdot W}$$

The training objective is a composite of several losses:

• Reconstruction and VQ Commitment Loss:

$$\mathcal{L}_{\mathrm{VQ}} = \|x - \hat{x}\|_2^2 + \sum_{t=1}^T \|\mathrm{sg}(z_e^t) - z_q^t\|_2^2 + \sum_{t=1}^T \|\mathrm{sg}(z_q^t) - z_e^t\|_2^2 + \lambda_p \mathcal{L}_{\mathrm{per}}(x, \hat{x})$$

where $\mathrm{sg}(\cdot)$ is the stop-gradient operator and $\mathcal{L}_{\mathrm{per}}$ denotes perceptual loss (e.g., LPIPS or VGG).

• Routing Loss ($\mathcal{L}_{\mathrm{router}}$):

Promotes balanced codebook usage, confident routing, and minimal quantization error through entropy, decision, and quantization-average losses.

• Identity Consistency Loss:

Applied in fine-tuning, encouraging ArcFace-based recognition feature consistency.

3. Training Procedure

STSCQ employs a three-stage progressive training schedule:

1. Stage 1: Switchable Token-Shared Codebooks
   • Replace global codebook with $M$ image-level sub-codebooks (shared across tokens).
   • Freeze encoder and decoder; optimize sub-codebooks and routing network.
   • Minimize reconstruction and routing loss.
2. Stage 2: Token-Specific Codebook Learning
   • For each group $i$ and token $t$, initialize $C^i_t$ from $C^i$.
   • Freeze encoder and decoder; optimize all $C^i_t$ and routing network.
   • Minimize sum of token-wise quantization errors and routing loss.
3. Stage 3: Decoder Fine-Tuning with Identity Loss
   • Freeze codebooks and routing; train decoder only.
   • Supplement loss with ArcFace-based identity regularization.

4. Quantization Capacity, Rate–Distortion, and BPP Efficiency

A global codebook of size $N$ provides high capacity but becomes infeasible at low bpp, where index sizes must be reduced. By setting $M = 2^s$, $K = N / 2^s$, effective total codebook capacity $M \cdot K \approx N$ is preserved, since each group specializes in a specific facial attribute distribution. Token-specific codebooks $C^i_t$ further reduce the per-token quantization variance, maximizing usage, and substantially decreasing quantization errors at a given bpp. The model’s bpp is dictated primarily by $T \log_2 K$, with the group overhead $\log_2 M$ negligible for typical $T$ and $K$.

5. Experimental Protocol and Evaluation

The method is trained and evaluated on standard face recognition datasets:

• Training Set: CASIA-WebFace ($\sim$500K images, 10,575 identities)
• Evaluation Sets: LFW, CFP-FP, AgeDB, CPLFW, CALFW
• Preprocessing: $256 \times 256$ crops, random horizontal flips
• Comparison Baselines: TiTok-S ($T=128$), TiTok-L ($T=32$), VQGAN/MaskGIT ($T=256$)
• Optimization: AdamW, learning rate $10^{-4}$, run on 8 NVIDIA V100 GPUs
• Training Schedule: Stage 1: 100K steps; Stage 2: 400K; Stage 3: 100K
• Metrics: Mean recognition accuracy (MeanAcc, ArcFace), identity similarity (IDS), bpp

6. Quantitative Results and Ablation Insights

At comparable or lower bpp, STSCQ yields substantial improvements over baselines:

Setting	bpp	Baseline MeanAcc	STSCQ MeanAcc	Baseline IDS	STSCQ IDS
MaskGIT-VQGAN+STSCQ, $T=256$	0.0503	90.70	93.51	0.6314	0.6659
TiTok-S, $T=128$, $K=4096$	0.0234	87.56	91.66	0.5764	0.6120
TiTok-L, $T=32$	0.0059	65.07	73.13	—	—

• Codebook Count ($M$): Performance increases with $M$ up to 16 (for $K=1024$), then degrades for $M \gg 16$ due to routing instability.
• Token-Specific vs. Shared: Per-token splitting yields +1.6–2.5% MeanAcc gain at fixed bpp.
• Inference Routing: Confident routing (CR) matches or surpasses nearest-neighbor quantization, with reduced latency and storage.
• Codebook Utilization: Mean per-token utilization improved from ~54% (global/shared) to ~74% (token-specific), decreasing dead embeddings.

7. Implications, Generalization, and Limitations

STSCQ achieves high-fidelity, identity-preserving reconstructions at bitrates previously unattainable for face compression methods based on vector quantization. The approach generalizes to any codebook-based representation learning scheme, as the architecture and training objectives are modular. A plausible implication is that similar hierarchical, token-specific codebook strategies may benefit other semantic-rich domains beyond faces, especially at ultra-low bpp. The primary limitation is routing instability as the number of codebook groups $M$ becomes large. Further research may focus on robust routing mechanisms, capacity scaling, and extending these insights to non-facial image modalities.

For in-depth methodology, experiments, and implementation specifics, refer to "Switchable Token-Specific Codebook Quantization For Face Image Compression" (Wang et al., 27 Oct 2025).