DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick
Authors: Mohammad Hassan Vali
Tom Bäckström
Arno Solin
Venue: ICLR 2026 Poster (2026)
Keywords: Vector Quantization
Differentiability
Backpropagation
Differentiable Vector Quantization
Gradient Collapse
Codebook Learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose DiVeQ and SF-DiVeQ, two differentiable vector quantization techniques that enable end-to-end training by preserving hard assignments in the forward pass while enabling meaningful gradient flow.

Abstract: Vector quantization is common in deep models, yet its hard assignments block gradients and hinder end-to-end training. We propose DiVeQ, which treats quantization as adding an error vector that mimics the quantization distortion, keeping the forward pass hard while letting gradients flow. We also present a space-filling variant (SF-DiVeQ) that assigns to a curve constructed by the lines connecting codewords, resulting in less quantization error and full codebook usage. Both methods train end-to-end without requiring auxiliary losses or temperature schedules. In VQ-VAE image compression, VQGAN image generation, and DAC speech coding tasks across various data sets, our proposed methods improve reconstruction and sample quality over alternative quantization approaches.

Published as a conference paper at ICLR 2026
DIVEQ: DIFFERENTIABLE VECTOR QUANTIZATION USING THE REPARAMETERIZATION TRICK
Mohammad Hassan Vali', Tom Bäckström² & Arno Solin'
'ELLIS Institute Finland & Department of Computer Science, Aalto University, Finland Department of Information and Communications Engineering, Aalto University, Finland (mohammad. vali, tom. tom. com.backstrom, arno. sol aalinn.a
ABSTRACT
Vector quantization is common in deep models, yet its hard assignments block gra- dients and hinder end-to-end training. We propose DiVeQ, which treats quantiza- tion as adding error vector that mimies the quantization distortion, keeping the forward pass hard while letting gradients flow. We also present a space-filling vari- ant (SF-DiVeQ) that assigns curve constructed by the lines connecting code- words, resulting in less quantization error and full codebook usage. Both methods train end-to-end without requiring auxiliary losses or temperature schedules. In VQ-VAE image compression, VQGAN image generation, and DAC speech cod- ing tasks across various data sets, our proposed methods improve reconstruction and sample quality over alternative quantization approaches.
INTRODUCTION
Vector quantization (VQ. Gersho Gray, 2012) a classical compression technique for discretizing continuous data distributions into finite set of representative vectors, resulting in codebook. Each data sample is assigned to its nearest codebook vector (codeword), enabling compact discrete rep- resentations. Within deep learning, VQ was popularized by the vector-quantized variational autoen- coder (VQ-VAE, van den Oord et al., 2017), where discretization of latent representations yields sub- stantial improvements in compression and generation quality. Since then, VQ has become a central component in architectures for images (Yu et al., 2021), video (Yan et al., 2021), and speech (Kumar al., 2023), serving as a fundamental module in the tool-stack for deep neural networks (DNNs).
Mapping a continuous data sample to its closest codeword is mathematically non-differentiable. Hence, when using VQ in the computational graph of DNN, the gradients will not pass through the VQ layer in the backward pass. For instance, in a VQ-VAE, the encoder parameters will not receive any gradients from the VQ layer. This issue known as the gradient collapse problem (Vali, 2025). There exists variety of solutions to this problem, each with its own challenges and drawbacks (see Table 1).
In this paper, we revisit differentiable vector quantization through the lens of the reparameteriza- tion trick (Kingma & Welling, 2013). We propose Differentiable Vector Quantization (DiVeQ), a method that models quantization the addition of simulated quantization error vector (see Fig. l). The direction this vector is aligned with the nearest codeword, while its magnitude equals the soput-codeword distance, thereby preserving hard assignments in the forward pass while enabling meaningful gradient flow. Unlike prior approaches such as NSVQ (Vali & Bäckström, 2022), which
E
0s
E
C: Codebook
C1
C2
C--
CK
DiveQ
d7z
05
2
2q
2q
||c. - x1|2-
D)
Va v tc -
0s
Er
Figure 1: We replace the non-differentiable VQ operation (2 = C₁+ = arg mine || - c, on the left with differentiable vector quantization (DiVeQ) on the right that lets the gradients flow.

---

Published as a conference paper at ICLR 2026
Table I: Summary of training properties for different VQ optimization methods in Sec. 2.
STE
EMA
RT
ST-GS
NSVQ
DIVeQ
SF-DiVeQ
No auxiliary loss terms
X
K
No complicated parameter tuning
/
Unbiased codebook gradients
X
N/A
No train-test mismatch
J
V
X
X
End-to-end training
K
X
V
Precise nearest-codeword assignment
V
V
Not prone to codebook misalignment
Avoiding codebook collapse
X
X
approximate the quantization error with limited directional fidelity, DiVeQ ensures that the differen- tiable surrogate remains geometrically consistent with the underlying nearest-neighhor operation.
Getting intuition from Vali & Bäckström (2023a) and building this foundation, further in- troduce Space-Filling DiVeQ (SF-DiVeQ), which extends quantization from discrete codewords to continuous line segments connecting neighboring codewords. This construction simultaneously mit- igates quantization error and alleviates codebook under-utilization without requiring replacement heuristics. By quantizing along structured manifolds within the codebook, SF-DiVeQ achieves full utilization and improves representational efficiency.
We evaluate DiVeQ and SF-DiVeQ on VQ-VAE (van den Oord al., 2017) for image compression, VQGAN (Esser et al., 2021) for image generation across CELEBA-HO. FFHQ. AFHQ, and LSUN Bedroom and Church, and the DAC (Kumar et al., 2023) model for speech coding on the VCTK data set. Our methods consistently improve image and speech reconstruction fidelity and maintain sample quality compared to existing quantization strategies. Importantly, DiVeQ and SF-DiVeQ act as drup-in replucements for standard VQ layers, requiring only minimal changes to model code. Our implementation code is publicly available at https: / github. com/
com/Aalt Aalt oML/DiVeO.
The contributions of this paper are as follows:
We propose DiVeQ, differentiable vector quantization technique that enables end-to-end train- ing with hard forward assignments, avoiding auxiliary losses and complicated tunings.
We further propose SF-DiVeQ, mpace-filling variant, which quantizes along codeword connec- tions, ensuring reduced quantization error and full codebook utilization without any auxiliary losses or codebook reinitialization. Contrary to all other methods, SF-DiVeQ avoids misalign- ment latent and codebook representations.
For VQ optimization methods prone to codebook collapse, we present an improved codebook replacement algorithm, achieving faster and more stable utilization than prior methods (see App. B. Also, we show that DiVeQ and SF-DiVeQ advantageously applicable other VQ variants like Residual VQ with superior performance to the other techniques (see App. C.9).
2 BACKGROUND AND RELATED WORK
Vector quantization (VQ) has been extensively studied in both signal processing and modern deep learning, and wide range strategies have been developed to overcome its non-differentiabiiiyy Since the core challenge addressed in this paper is precisely the integration of VQ into end-to-end trainable neural architectures, we provide structured review of existing solutions. This detailed discussion is necessary because, although many methods yield partial remedies, they differ substan- tially in terms of optimization objectives, gradient fidelity, codebook utilization, and computational overhead. Our goal in this section is therefore twofold: (i) to establish precise baseline under- standing the mechanisms underlying each class of methods, and (ii) to highlight their limitations in relation the desiderata of differentiable quantization.
VQ (Gersho & Gray, 2012) clusters a continuous distribution to limited set codewords, by map- ping latent vector RD the nearest codeword C;- R of codebook C (C1, (cI..... CK ) where arg min, I| - c,112. This mapping introduces a distortion &o = z - Ci+, known as quantization error. Since the nearest-neighhor assignment is non-differentiable, VQ blocks gradients
2

---

Published as a conference paper at ICLR 2026
during backpropagation, leading to the well-known gradient collapse problem (Vali, 2025). To explain existing techniques, we consider the case of training a VQ codebook in a vector-quantized variational autoencoder (VQ-VAE, van den Oord et al., 2017). In VQ-VAE (Fig. 1), the imput æ is fed to the encoder E to obtain the continuous latent variable z, which is then discretized by VQ to 2. The discrete latent variable E is fed to the decoder D reconstruct the input Wy:
((x) _
VQ(=); 2 = C₁- arg min ||2 - cyll2 - 2, D
(1)
6
Since arg min is not differentiable does not exist), VQ does not pass any gradients from æ,
02 +zz
Straight-Thrae Estimator (STE, Bengio et al., 2013) deals with this issue by copying the gra- dients through the non-differentiable VQ function, while assuming
1, such that it defines the final quantized latent as 2g = + sg( z]. The loss function that STE uses is
(2) where sg[:] refers to the stop gradient operator, and MSE is the mean squared error. The first term the reconstruction loss that optimizes the encoder and decoder parameters by copying the gradi- ents intact over VQ in backpropagation. The second term is the codebook loss that optimizes the selected codewords (i.e., c, and the third term is the commitment loss that optimizes the encoder parameters to make the encoder output z close to Ca
Loss MSE(x, z,) + ||sg [z] c--11 112 + B ||
sg [
Exponential Moving Averages (EMA, van den Oord al., 2017) also copies the gradients intact through VQ via 2n = 2 + sg 2 z]. However, EMA updates the codebook vectors as a function of moving averages of latent variables z. EMA keeps the reconstruction and commitment losses, and skips the codebook loss in Eq. (2) to update the selected codewords as
(t)
(1) gr- (1(0)
(1-1) +(1 7) S.L. 1C = 1- + (1--7).m4 (1) and g = 9:9
j=1
where at training iteration f, fz;-, 2i. Zir,nz mg- } is the set of latent variables that are closest the selected codeword ci+, and 7 E (0, 1) is the decay rate that controls how much past observations influence the updates.
Rotation Trick (RT, Fifty al., 2025) passes gradients through VQ by transforming the latent to its closest codeword C.- with combination of rotation and rescaling such that it defines
||c-|12
2q = Rg(pR]> s.t. =
and R RI 2rr T + 20;- )
(4)
|
T RT uses the same loss as in Eq. (2), while the codebook can be updated via gradients using codebook loss via EMA.
where 2 . C.-
Tc.. I and r
& 2+0
Gumbel-Softmax (GS, Jang et al., 2016) samples differentiable variables y, from a continuous distribution that approximates discrete K-class categorical distribution over K codewords
exp (flog (m;) + BT)/T)
softmax (llgg (T:) +g))/r) Σ exp (log (m₂) + +g))// /r)" E ,K}, (5)
where T21, kx) are class probabilities of the categorical distribution and {g1, .2........ gx
are i.i.d. samples drawn from Gumbel 1) distribution. The scalar - is the softmax temperature.
When approaches 0, the Gumbel-Softmax distribution converges to discrete categorical distribu- tion. To train a VQ-VAE, the loss function that Gumbel-Sofhnee
uses is
Loss MSE(#, xx,) + + y DKL [Q(= x)) P(=)]
(6)
where DKL. is the Hellbacck-eibble divergence that pushes the encoder's posterior distribution Q(z æ) to be close to the uniform prior distribution of P(z)
K The DKL term encour- ages all K codewords to be used for quantization. As stated in Jang et al. (2016), for VQ tasks it is better to use the Straight-Through version of Gumbel-Softmax (ST-GS), which y; is discretized using arg max in the forward pass, and gradients are computed from the continuous y; in Eq, (5).
Noise Substitution in Vector Quantization (NSVQ, Vali Bäckström, 2022) simulates the VQ distortion by adding noise to the latent vector to mimic the original quantization error Eq:
Zq + EQ + ||= =1|2
2 = C.+ arg min || c,ll2 and v~ ()( (7)
|
3

---

Published as a conference paper at ICLR 2026
In this way, 2g is differentiable function of the latent & and the selected codeword Cj.. As shown in Fig. 2. as NSVQ samples the noise from from.N((. (0, I), the quantized latent 2y can be mapped to any random direction on the surface of a hypersphere with the radius of quantization error that equals || - -1z---lll 1|2. Fig. shows the probability that 2, incurs higher quantization errors than the actual error (of || c. I12) equals 02/3600. For any radius, 82 = 240". Hence, NSVQ incurs higher quantization error than the true error with the probability of 240 /3600 0.67. In higher dimensions, this probability approaches 1, making NSVQ increasingly likely to overshoot the true nearest-codeword dis- tortion. Furthermore, due to high randomness in the direction of zq. is challenging for the codebook to converge to its optimum location. This randomness also causes train- mismatch, which leads to poor reconstructions, because NSVQ applies dif- ferent mappings for the same during training and testing.
02
C1
e = || 112
Figure 2: Illustration of NSVQ quantization. Input z is mapped to a random point on the circle. The mapping overshoots the true quantization error with proba- bility 02/3600 x 0.67, leading to a higher distortion than the nearest-codeword assignment.
Issues with the State-of-the-Art All of the methods mentioned above have their own drawbacks, which are highlighted in Ta- ble 1. To optimize the VQ codebook, they (i) change the hy- perplane of optimization by adding auxiliary loss terms (Eqs. (2) and (6)) to the main objective loss (STE, EMA, RT, ST-GS), (ii) need additional hyperparameter (i.e., a, a,,3 7, 9P, 7) tuning (STE, EMA, RT, ST-GS), (iii) do not train the codebook end-to-end (STE, EMA, RT), (iv) cause biased gradients via mismatch between forward and backward passes (STE, ST-GS), (v) cause rest mismatch as they apply different quantization for training and testing (ST-GS, NSVQ). (vi) suffer from codebook collapse (STE, EMA, RT, ST-GS, NSVQ), and (vi)) are prone to misaligned latent and codebook representations (STE, EMA, RT, ST-GS, NSVQ) (see App. C.8).
Other Related Work There exist many approaches that still use STE to pass gradients through the VQ layer (Chang et al., 2022; Rombach et al., 2022; Zhu et al., 2023; Huang et al., 2023; Dong et al., 2023; Lee et al., 2022; Huh et al., 2023), each with small modifications to improve something specific. For instance, some methods use STE while resolving codebook collapse with different codebook replacement and reinitialization strategies (Kolesnikov et al., 2022: Lañeucki al., 2020; Zheng & Vedaldi, 2023; Dhariwal et al., 2020; Zeghidour et al., 2021), or by using different distance metrics than Euclidean (Yu et al., 2021: Goswami al., 2024). Others use stochastic sampling (Maddison et al., 2017; Takida et al., 2022: Chen et al., 2024) while increasing codebook utilization by additional loss terms (Zhang et al., 2023; Yu et al., 2023), or soft quantization to backpropagate gradients through the VQ layer (Gautam et al.. 2023; Agustsson et al., 2017). Additionally, some recent methods use NSVQ (Gómez et al., 2023; Walsh et al., 2024; Lee et al., 2024; Wang et al., 2025; Ye et al., 2025; Zhu et al., 2025). and RT (Kim et al., 2025; Xue et al., 2025; Bae & Shin, 2025) techniques. In rather unique lines of work, Finite Scalar Quantization (FSQ, Mentzer et al., 2023), Random Projection Quantizer (RPQ, Chiu et al., 2022), Binary Spherical Quantization (BSQ, Zhao et al., 2024), and Lookup-Free Quantization (LFQ. Yu et al., 2023) constrain the codebook to follow predefined geometric structure and do not train the codebook together with the model.
METHODS
We address the problem of making vector quantization differentiable while preserving hard assign- ments in the forward pass. Our goal is to construct differentiable surrogate 2,(2, C) that (i) co- incides with the hard nearest-neighhor assignment in the small-variance limit, (ii) yields stable and geometrically faithful gradients for both encoder and codebook, and (iii) avoids auxiliary losses, hyperparameter tuning, or train-test test mismatches. We first introduce DiVe (Sec. 3.1,, which repa- rameterizes the quantization error to align its direction with the nearest codeword. We then extend this construction to SF-DiVeQ (Sec. 3.2), which quantizes along line segments between neighboring codewords reduce error and promote codebook utilization.
3.1 DIVEQ DIFFERENTIABLE VQ VIA DIRECTIONAL REPARAMETERIZATTOO
DiV models quantization adding simulated error vector whose magnitude equals the input-
4

---

Published as a conference paper at ICLR 2026
211
C1
2
(a) 0² = 10
(b)
(c) 0² 0.1
(d) o² = 0.01
Figure 3: Impact of a² in quantiiaatii accuracy. Each panel shows mappings of input 2; its closest codeword Ct- using our proposed DiVeQ (Eq. (8)) when sampling 1000 random directional trom.N(O,
vectors Va from MO o² I). DiVeQ quantization accuracy increases when o² 0 (see App. C.5).
codeword distance and whose direction is aligned with the nearest codeword. Similar to NSVQ (see Sec. 2), to obtain a differentiable quantized input Zq. DiVeQ approximates VQ as addition of the quantization eгтor εQ the input 2, i.e., 2g
+ Eq. However, difference NSVQ, Di VeQ models &Q such that 2y points to the closest codeword C;- precisely and thus, Eg results in the original quantization error of || c.-1l2 (see Fig. 3). To this end, by getting intuition from the reparameterization trick (Kingma & Welling, 2013), we define a directional noise V = v - d such that d = Ct- and v ~ _N00 r)) Then,
Zy z + Eo + ||d 2-12-- 8g
Va
x |c--- 2112-88 8g
v : C-- arg min ||
Vd
Cj 22,
[alla
6j
where sg[-] is the stop gradient operator. Here, by changing the variance o² from to zero (Fig. 3a to Fig. 3d), the quantization accuracy will increase, as 2g starts from pointing to a hypersphere (as in NSVQ and approaches pointing exactly to the selected codeword cr. The limits give:
lim Va lim [NO o2 dj N(e 2, 00) =
2g |cc--
l| sg [s]
-00
02-000
NN d = d - = 2q x ||c =112 sg
C₁+
lim Va=
lim
Ci+
-00
||c. - = =1|2
(9)
where s is a random variable almost uniformly distributed on the surface of the D-dim unit hyper- sphere of S (V's E RD ||s||| 1}. When using Eq. (8) during training, 2, is differentiable function of z and cr,, and thus it can be used in end-to-end training of VQ together with other trainable modules in neural network. Hence, the gradients can be calculated as
dzg 1 0z
C2+
dzg Ocr-
C₁-
C₁+
and
s.t.
sg
(10)
| 21|2
c;.- 21|2
|c,,
21|2
3.2 SF-DIVEQ: SPACE-FILLING DIFFERENTIABLE VECTOR QUANTIZATION
Space-Filling Vector Quantization (SFVQ, Vali & Bäckström, 2023a) is modification of VQ that maps an input to a piece-wise continuous curve (see Fig. 1 in Vali & Bäckström, 2023a). Contrary to VQ that quantizes the input exclusively the codewords, SFVQ twisted inside a D-dim distribution and quantizes the input on a continuous curve whose corner points are the codewords of the SFVQ base codebook C. SFVQ adopts dithering trick generate the curve. For each training iteration, by having an input E RD and a base codebook C with K codewords, SFVQ generates a dithered codebook Cd = {ct...... cx_ } with K - codewords by sampling at random places on the line connecting two subsequent codewords of cj and Cj+1. The input z is then quantized to the closest codeword from the dithered codebook Cdd as
curve that gets tumed and
2 arg min of
||= e/ et = (1 - Cc- + A Cr j.i E (1....,, K -1,, 1}, (11)
where cf. is the closest dithered codeword that is generated by the interpolation of two subsequent codewords of ci- and C1++1 from the base codebook C. Xy is the interpolation factor for the code- words ej and c3+1 that is sampled from uniform distribution of U(0, 1). Generating new dithered
5

---

Published as a conference paper at ICLR 2026
0.122
0.024
0.018
NIA
26-10--
5-10
39-10-
STE
EMA
RT
ST-GS
NSVQ
DiVeQ
SF-DiVeQ
Figure 4: Codebook misalignment: t-SNE plots of the learned codebook C2 (red crosses) and latent P- (gray points) representations for different VQ methods in VQ-VAE compression. The figure shows the misalignment between C₂ and P= (discussed in Sec. 4) for different methods. The plots refer to the cases highlighted in Fig. 26. The numbers report distortion per bit,. (see App. C.8).
codebooks by random interpolations (using Aÿ) on the line connecting Cj and Cj+1 over and over for different training batches will establish continuous topology between subsequent codewords of C and locates SFVQ codewords such that the lines connecting subsequent codewords to be inside the distribution space. The reason is that the points on the SFVQ curve are representatives of dithered codewords ef that should be valid quantization points for the input data.
Proposed SF-DiVeQ When training VQ, if a codeword is not selected for quantization, will not receive any gradients and, as result, it will not get updated. This problem is known as codebook collapse (Vali, 2025; Mentzer et al., 2023), in which a subset of codebook vectors remains inac- tive and will not get updated during training. To resolve this issue, NSVQ proposed a codebook replacement procedure such that after a specific number of training batches, inactive codewords will be replaced by perturbation of the active ones. To maximize the codebook vectors' usage during training. it is better and safer always to apply codebook replacement regardless of the VQ optimiza- tion technique (e.g., STE, EMA, RT, ST-GS, NSVQ. or DiVeQ), According Huh al. (2023), codebook replacement techniques work well and do not degrade the model performance. That is why we adopt codebook replacement in this paper. However, codebook replacement is heuris- tic method that adds to the complications of VQ training. By leveraging intuition from SFVQ, we propose SF-DiVeQ, a differentiable VQ technique that eliminates the need for codebook replace- ment. SF-DiVeQ quantizes the input to random location on the line connecting two subsequent codewords as
(1 - Vd.
Ap- Ve
2, = + ||.,- =212-89
+ +||c-+11 - s212.sg
1 E f1....,
K K-1} 1},
Vd
Vd
(12)
where C₁ and 41 are the two codewords that their interpolation the closest quantization point to the input z, and Vd,- = (c,. -2) and vd... - -2) such that v ~ N(0,0²I)) A is the interpolation factor that is sampled from the uniform distribution of (0, 1) for each training batch. As discussed earlier, since during training SF-DiVeQ quantizes = on the lines connecting subsequent codewords, these lines will be pulled inside the distribution space as they should be valid quantization points. As a result of this property, SF-DiVeQ prevents misalignment of codebook and latent representations (see Fig. 4), and it does not need any heuristic codebook replacement during training. Furthermore, as SF-DiVeQ quantizes z on a curve (rather than exclusively on the codewords), it has many more degrees of freedom for quantization than ordinary VQ methods. Therefore, it potentially results in smaller quantization errors than the other methods. Apart from DiV and SF-DiVeQ, we also propose two variants of them, which are discussed in App. B.2.
EXPERIMENTS
We compare the performance of our proposed DiVeQ and SF-DiVeQ with other approaches of STE, EMA, RT, ST-GS, and NSVQ in three different applications of image compression, image gener- ation, and speech coding. For image compression and speech coding tasks, we use the VQ-VAE of van den Oord et al. (2017) and the DAC model in Kumar et al. (2023), respectively, with minor modifications to have a basic compression model that can clearly reflect the performance differ- ence between various VQ methods, and for image generation, we use VQGAN (Esser et al., 2021). Experimental setup and relevant results are provided in the following. For more details on the im- plementation, model architectures, and hyperparameters, App. A.

---

Published as a conference paper at ICLR 2026
CELEBA --
1873
1962
338
13000
1-473
1855
0.259
Db246
0m2 78
ho260
194
0.2
28
AFHQ
0.224
0.256
FFHQ
LSUN BEDROOM
D. 182
U. 0.422
0.0001
D. 4939
0.30g
D,d7a
Ground truth
STE
EMA
RT
ST-GS
NSVQ
DiVeQ
SF-DiVeQ
Figure 5; DiVeQ and SF-DiVeQ improve image reconstruction. Qualitative comparison of recon- structed images VQ-VAE compression task for different VQ optimization methods with an 11-bit codebook (i.e,, codebook size of 2 2048). We report LPIPS., values in the left-hand corners.
Following the sensitivity experiments in App. C.5, we do not consider the variance a² an actual hyperparameter that needs tuning long it is small (o² I 10- Different values for the variance, when 6² < 10-2, only lead to marginal change in the performance of both DiVeQ and SF-DiVeQ in both VQ-VAE image compression and VQGAN image generation tasks (see Fig. 20, Fig. 21, and Table 10). For the results in the main paper, we use o² = 10 (for image compression and speech coding) and o² 10 (for image generation), with more results in App. C.5. Note that the values are not tuned on per-data set basis.
Loss Functions VQ- TQEE compression, to enhance the quality of reconstructions for 256x256 images, in addition to the MSE reconstruction loss in Eqs. (2) and (6), we add LPIPS (VGG-16) as the perceptual loss with a weighting coefficient of one. For VQGAN (Esser et al., 2021) and DAC (Kumar et al., 2023) models, we use the same loss functions as those used in the original models.
Evaluation Metrics During inference, the trained models with different VQ optimization tech- niques are used to reconstruct the test set images and speech samples for VQ-VAE compression and DAC-based codec, respectively, and to generate new images for VQGAN, For inference of all VQ methods, the input is mapped to the closest codeword C₁. with hard VQ using arg min, except SF-DiVeQ, in which the input is mapped to the nearest point on the space-filling curve (see Sec. in Vali & Bäckström (2023a)). Structural Similarity Index Measure (SSIM), Peak Signal to Noise Ra- tio (PSNR) (both from scikit-inage library), and Learned Perceptual Image Patch Similarity (LPIPS, Zhang et al., 2018) used evaluate the quality of the reconstructions for image compression, and Fréchet Inception Distance (FID) score (Parmar et al., 2022) is used to assess the quality of VQGAN generations. Log Spectral Distance (LSD), Mel-Frequency Cepstral Coefficients (MFCC) distance, Perceptual Evaluation of Speech Quality (PESQ), and Short-Time Objective Intelligibility (STOI) are used to evaluate the decompressed speech quality.
Data Sets In both image compression and generation tasks, we do the experiments over AFHQ (Choi al., 2020), CELEBA-HU HHQ (Karras al., 2018), FFHQ (Karras et al., 2019), LSUN Bedroom, and LSUN Church (Yu et al., 2016) data sets with resolution of 256x The data sets contain 15803, 30k, 70k, 70k, 70k images, respectively. For image compression, all data sets are divided into 80/20% train-test splits. For image generation, AFHQ and CELEBA-HQ. compute the FID score, we generate the same number of images used for training. In speech coding, we use the CSTR VCTK data set (Yamagishi et al.) divided into 80/20% train-test split with no overlapping speakers and speech files between the train and test sets (more details in App. A.3).
use the train sets for training except for which we use the full data sets, as we need more data for the task. To
Codebook Replacement and Initialization In the NSVQ paper (Vali & Bäckström, 2022), the NSVQ method is compared with STE and EMA, such that the codebook replacement is only used

---

Published as a conference paper at ICLR 2026
24
0.6
0.6
22
PSNR -
-PIPS
0.5
SSIM
0.5
20
STB
EMA
RT
ST-G8
NSVQ
DIVeQ (ours)
ST. DiVe (ours)
16
0.3
5
5
Codebook bits
Codebook bits
Codebook bits
Figure 6: Consistent improvement in image reconstruction. Quantitative comparison of recon- structed images from AFHQ test set for different VQ optimization methods over different codebook sizes. Each curve the average of the metric over all test set reconstructions and over three different individual runs. Results for FFHQ, CELEBA-HQ. and LSUN are in Figs. 12 to 15 in App. C.I.
for NSVQ. To have a fair comparison of different VQ training methods, our proposed codebook replacement (App. B.1) is used for all methods (except SF-DiVeQ, which does not require it). To evaluate how much codebook replacement contributes to the overall performance, in App. C.6, we skipped codebook replacement in the VQ- VQ-VAE compression task, where the results show that our proposed DiVeQ consistently outperforms other VQ optimization methods. Moreover, in all ex- periments, codebook initialization is the same for all methods except for SF-DiVeQ, which uses a custom initialization, in which the models are trained without VQ for two epochs, then SF-DiVeQ starts, such that the codebook is initialized by the mean of latent vectors of the last 20-50 training batches. In App. C.7, show that the performance of SF-DiVeQ with random initialization (i.e., in which the codebook is randomly initialized by the latent vectors of the first training batch) remains almost the same the custom initialization in the VQ-VAE compression task.
VQ-VAE Compression Results After training VQ-VAE with different VQ training methods, the trained codebook is used to reconstruct the test set images. Fig. shows ground truth images from different data sets and their corresponding reconstructions using learned codebooks from different VQ techniques. Apart from qualitative comparison, all test set images (from different data sets) are reconstructed, and then the SSIM, PSNR, and LPIPS metrics are computed for all of these recon- structions. Fig. 6 shows the quantitative comparison of all VQ methods over different codebook sizes for the AFHQ data set. The reported value for each metric is the average of that metric over all test set reconstructions and over three different individual runs. In the figure, our DiVeQ and SF-DiVeQ outperform other VQ optimization methods for different codebook sizes and for all three metries. Our proposed methods yield higher SSIM and PSNR values, and lower LPIPS with a big margin to RT and NSVQ, and with a smaller margin to STE, EMA, and ST-GS. In addition, the performance gap of our methods to STE and ST-GS grows with the increase in codebook size, while this gap decreases compared to the RT method. The last takeaway is that the RT performs poorly for small codebook sizes. App. C. provides similar quantitative comparisons for other data sets.
VQGAN Generation Results After training the VQGAN's generator and transformer models for each data set, we compute the FID between the original train set and generated samples. Ta- ble 2 presents the obtained FID val- ues for different VQ optimization meth- ods for CELEBA-HQ data set. Ta- ble in App. C.2 provides similar FID comparisons for other data sets. We trained the generator (i.e., VAE) with two different hyperparameter settings of HP, : (2r=2.5-1) -5, batch size= 8) and HP2: (Ir 2.5 10-1 batch size 32).
Table 2: DiVeQ SF-DiVeQ robustify training for challenging small codebooks, and they are less sitive to hyperparameter settings. FIDJ scores for VQ optimization methods on CELEBA-HQ. Values in red refer to the misalignment cases as in Fig. 4.
CELEBA- HQ
2.5-10-3
2.5- 10--
data set
batch size =
batch size
32
Approach bits
12
9
10
12
5.57
6.69
7.54
7.34
9.45
EMA
6.86
6.30
6.32
6.24
8.42
7.42
6.97
9.41
RT ST
9.32
7.55
6,40
5.44
12.3
9.33
6.53
6.58
8.47
6.81
5.48
5.47
309
41.1
197
155
NSVQ
81.5
70.4
59.2
48.9
78.4
70.1
62.1
49.6
DiveQ (ours)
6.69
7.69
8.01
9.54
SF-DiVeQ (ours)
6.24
5.21
5.57
6.00
8.46
6.66
7.02
7.4
8

---

Published as a conference paper at ICLR 2026
5.77
6.30
7.55
6.81
20.4
CHURCH (10 CELEBA HQQ (9
$.09
3.81
5.20
20-5
3.90
sturtters
3.36
635
6.2
17.1 83.8
FIRQ
STE
EMA
RT
ST-GS
NSVQ
DiVeQ
SF-DiveQ
Figure 7: Generation task. Qualitative comparison of randomly generated images in the VQGAN generation task for different VQ optimization methods. We report FIDJ values in the left-hand corners. More generations are provided in Figs. 30 to 34 in App. C.11.
The HP1 uses small batch size and learning rate to ensure convergence for all VQ methods while is less likely to highlight the robustness of our approaches to the use of different hyperparameter settings. Using larger batch size and higher learning rate (i.e., HP2) might cause misalignment between codebook and latent representations, where generations result in high FID values. The cases where misalignment happen are highlighted in red in Table 2 and Table 7 (i.e., cases where FID > 30). According to the tables, our proposed DiVeQ is less prone to misalignment compared to other VQ methods, whereas no misalignment happens for our proposed SF-DiVeQ.
The evaluations reveal that RT, ST-GS, and NSVQ reduce FID with the increase in codebook size, and (similar to STE and EMA) methods maintain the generation quality over different codebook sizes such that they achieve lower FIDs than the other methods for small codebooks (i.e., 8- and 9-bit codebooks) where the generation task becomes more challenging. Furthermore, by pairing up the FIDs for the RT method with the results in the VQ- VAE compression task, we can claim that RT performs poor codebook learning for small codebooks in both compression and generation tasks. Note that the FID metric reflects distributional similarity, not necessarily the reconstruction quality. Hence, the FID value should not necessarily improve with the increase codebook size because (i) the transformer might struggle with large K where the entropy of token indices is high, and (ii) sampling is inherently a random process that might end up in different FID values when sampling several times from a fixed pretrained transformer. Fig. 7 shows qualitative comparison of VQGAN generations for different VQ methods (see App. C.11 for more qualitative comparisons).
Misalignment of Codebook and Latent Representations When training a neural network (e.g., VQ-VAE), the latent representation P, of the network layers is constantly shifting. For high learning rates and large batch sizes, the update shift in P_ be large, causing misalignment between Pa and codebook representation C.. Misalignment means that C₂ does not fit well with Pz, and this phenomenon is also called internal codebook covariate shift (Huh al., 2023).
According to Shannon's 's rate-distortion theory (Shannon, 1959), the rate-distortion function D(R) continuous, strietly decreasing, convex function. Therefore, in lossy compression task, if the codebook size grows, then the distortion D must decrease, and as result, the quantitative metrics improve (assuming the metrics fully reflect the amount of distortion). After plotting the quantitative results of trained VQ-VAEs with different learning rates and batch sizes (see the ablation studies on batch size and learning rate in Apps. C.3 and C.4), we spot that the rate-distortion theory does not hold in some experiments, such that the increase in codebook size does not result in improvement of quantitative metries. Hence, we hypothesize that the reason could the misalignment of Pa and C.. In Fig. 26, we highlighted those cases that are potentially prone to misalignment with red circles. For illustration purposes, we provide t-SNE plots of these cases in Fig. 4, which confirms hypothesis that misalignment is the main culprit. It is important to note that by observing the quantitative results (e.g., Fig. 26), we spot the misalignment cases where the rate-distotion theory does not hold, or when the quantitative metrics do not improve by the increase in codebook
9

---

Published as a conference paper at ICLR 2026
Table 3: Consistent improvement in speech decompression. Quantitative comparison of decom-
pressed speech samples from the VCTK test set for different VQ optimization methods over different codebook sizes when batch size- 64. Each value shows the average of the metric over all test set decompressed samples. Results for other batch sizes are provided in Tables 8 and 9.
Log spectral distance.)
MFCC distance)
PESQT
STOM
Appraach's bits
13
11
13
11
13
STE
1.11
3.20
1.14
1.11
93.3
344
105
96.0
1.22
1.04
1.14
1.22
0.75
0.40
0.71
0.75
EMA
3.40
1.03
1.03
1.02
317
72.6
72.4
69.1
1.03
1.55
1.59
1.67
0.39
0.83
0.84
0.84
RT
1.09
1.05
1.04
1.05
92.6
84.0
82.8
80.8
1.27
1.35
1.43
1.41
0.76
0.80
0.80
0.81
ST-GS
1.13
3.32
3.45
1.13
97.9
349
333
97.9
1.03
1.04
1.21
0.76
0.41
0.39
0.76
NSVQ
1.11
1.09
1.10
1.07
108
103
101
93.6
1.35
1.43
1.49
1.56
0.79
0.81
0.82
0.83
DiveQ (ours)
1.04
1.02
75.5
72.6
1.55
1.64
0.83
0.85
SF-DIVeQ (ours)
1.05
1.02
1.01
1.01
74.6
71.8
68.3
66.8
1.49
1.52
1.62
1.75
0.83
0.84
0.85
0.85
size. Therefore, the 1-SNE plots should only interpreted as an intuitive visualization of codebook homogeneity with respect to the latents within each method separately. Furthermore, we consider case as misalignment even if all of the codewords are inside the latent representation Pg (like STE in Fig. 4), but they are not properly scattered within Pz.
Note that codebook replacement (App. B. can help reduce the risk of misalignment occurring, but it does not completely avoid it. Fig. 4 shows the learned P. and C. representations for the cases where misalignment happens (i.e., the cases highlighted in Fig. 26, where rate-distortion theory does not hold). Different types of misalignment happen for all methods (except for SF-DiVeQ), even if they use codebook replacement. For STE, RT, and NSVQ, the codewords are located inside Pa, but they are not homogeneously scattered. In ST-GS, Cg completely out of Pz. whereas in EMA, a large portion of codewords are outside Pz. A much milder misalignment happens for our proposed DiVeQ, as a small portion of codewords are scattered densely within P=, particularly in the corners. However, SF- VeQ scatters the codewords homogeneously within P.. In our experiments, we do not notice any sign of misalignment for SF-DiVeQ. We hypothesize the reason is SF-DiVeQ's s train- ing strategy, which pulls codewords inside Pa without requiring heuristic codebook replacement.
As final experiment, we evaluate speech decompression on the VCTK data set. We train a DAC-based speech coding model (cf. Kumar et al., 2023) with different VQ methods and compute the LSD, MFCC distance, PESQ, and STOI metrics between the original and decompressed speech. Table 3 shows the average of each metric over all VCTK test samples for different codebook sizes when batch size- = 64, and Tables 8 and 9 in App. C.3 show the same comparisons for batch size [32, 16]. In Tables 3, 8, 9, entries highlighted in red correspond to runs where codebook latent misalignment occurs, and the decoded speech is severely degraded and unintelligible. The tables show that our proposed DiVeQ and SF-DiVeQ consistently yield higher decomppreseed-speech quality than alternative VQ optimization methods. Moreover, in Tables 3, 8, STE, EMA, and ST-GS fail to converge some configurations and exhibit misalignment, whereas the remaining methods avoid this issue. For this experiment, we intentionally use a simplified DAC- based codec with a single codebook and smaller encoder and decoder networks than in the original setup, in order to isolate the effect of the VQ method (see App. A.3). Decompressed speech samples are provided in the supplementary material for subjective comparison.
Speech Coding Results
DISCUSSION AND CONCLUSION
In this paper, we introduced DiV and SF-DiVeQ, two new differentiable vector quantization meth- ods that can be trained end-to-end in neural network. We provided comprehensive overview of the limitations of existing approaches (Sec. 2) that motivated our approach. In Sec. 3.1, we showed how Di uses directional reparameterization of the quantization error to preserve hard assignments in the forward pass while allowing gradients to flow. We further extended this idea to Si-e (Sec. 3.2), which quantizes along codeword connections (i.e., space-filling) to reduce error and en- sure full codebook utilization. In Sec. 4, we demonstrated that both methods outperform existing quantization strategies in VQ-VAE, VQGAN, and DAC models across several benchmarks. Also, they are advantageously applicable to other VQ variants like Residual VQ (experiments in App. C.9).
Importantly, DiVeQ and SF-Di act as drop-in replacements for standard VQ layers and require no auxiliary losses, temperature schedules, or special heuristics. We advocate their use as choices for differentiable quantization in deep generative models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported by the Research Council of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAL We acknowledge funding from the Research Council of Finland (grants 339730 and 362408). We acknowledge the computational resources provided by the Aalto Science-IT project and CSC-IT Center for Science, Finland. We also gratefully acknowledge the support from Esteban Gómez in conducting the DAC-based speech coding experiments.
REPRODUCIBILITY STATEMENT
In App. A, we provide comprehensive implementation details our VQ-VAE image compression, VQGAN image generation, and DAC-based speech coding models. The section includes model architectures, hyperparameters, and the specifications how we implement other VQ optimization methods. For image compression and generation tasks, all the experiments done using the well- known benchmark data sets of AFHQ, CELEBA-HO, FFHQ, LSUN Bedroom, and LSUN Church obtained from https: / www. k kaggle cwm/ datasets..
We provide code implementation to train both VQ-VAR image compression and VQGAN image generation models (including generator, discriminator, and transformer) together with README md file explaining the structure of the contents in the supplementary material. In addi- tion, in relevance our DAC-based speech coding experiment, we add some speech samples the supplementary materials for subjective comparisons. Implementations of all VQ optimization meth- ods are also included in the supplementary materials. Upon acceptance of the paper, the reference implementation and examples will be made available as GitHub repository.
REFERENCES
Eirikur Agustsson, Fabian Mentzer, Michael Tschannen, Lukas Cavigelli, Radu Timofte, Luca Benini, and Luc Van Gool. Soft-to-hard vector quantization for end-to-end learning compress- ible representations. In Advances in Neural Information Processing Systems 31 (NeurIPS), pp- 1141-1151. Curran Associates, Inc., 2017.
Jungwoo Bae and Jitae Shin. Robust training framework via multi-stage feature rectification. In 2025 International Technical Conference on Cireuits/Systemss, Computers, and Communications (ITC-CSCC), pp. 1-5. IEEE, 2025.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: urX.v:/308.3433 3432, 2013.
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11315-11325, 2022.
Hang Chen, Sankepally Sainath Reddy, Ziwei Chen, and Dianbo Liu. Balance of number of embed- ding and their dimensions in vector quantization. arXiv preprint arXiv:2407.04993 2024.
Yongjian Chen, Tao Guan, and Cheng Wang. Approximate nearest neighbor search by residual vector quantization. Sensors, 11112) 11111112197771 11273, 2010.
Chung-Cheng Chiu, James Qin, Yu Zhang, Jiahui Yu, and Yonghui Wu. Self-supervised learning with random- ceadom-projection quantizer for speech recognition. In Proceedings of the International Conference Machine Learning (ICML), pp. 3915-3924. PMLR, 2022.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. StarGAN v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8188-8197, 2020.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: a generative model for music. arXiv preprint MxXiv:2005.00341 2020.
11

---

Published as a conference paper at ICLR 2026
Xiaoyi Dong, Jianmin Bao, Ting Zhang. Dongdong Chen, Weiming Zhang, Lu Yuan, Dong Chen, Fang Wen, Nenghai Yu, and Baining Guo. PeCo: Perceptual codebook for bert pre-training of vision transformers. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 552-560, 2023.
Patrick Esser, Robin Rombach, and Bjor Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recog- nilion (CVPR), pp. 12873-12883, 2021.
Christopher Fifty, Ronald Guenther Junkins, Dennis Duan, Aniketh Iyengar, Jerry Weihong Liu, Ehsan Amid, Sebastian Thrun, and Christopher Re. Restructuring vector quantization with the rotation trick. In International Conference on Learning Representations (ICLR), 2025.
Tanmay Gautam, Reid Pryzant, Ziyi Yang, Chenguang Zhu, and Somayeh Sojoudi. Soft con- vex quantization: Revisiting vector quantization with convex optimization. arXiv preprint arXiv:2310.03004, 2023.
Allen Gersho and Robert M Gray. Vector Quantization and Signal Compression, volume 159. Springer Science & Business Media, 2012.
Esteban Gómez, Mohammad Hassan Vali, and Tom Bäckström. Low-complexity real-time neural network for blind bandwidth extension of wideband speech. In Proceedings of the European Signal Processing Conference (EUSIPCO), pp. 31-35. IEEE, 2023.
Nabarun Goswami, Yusuke Mukuta, and Tatsuya Harada. HyperVQ: MLR-based vector quantiza- tion in hyperbolic space. arXiv preprint arXiv:2403.13015 2024.
Mengqi Huang, Zhendong Mao, Zhuowei Chen, and Yongdong Zhang. Towards accurate image coding: Improved autoregressive image generation with dynamic vector quantization. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22596-22605, 2023.
Minyoung Huh, Brian Cheung, Pulkit Agrawal, and Phillip Isola. Straightening the straight- through estimator: Overcoming optimization challenges in vector quantized networks. In Pro- ceedings of the International Conference on Machine Learning (ICML), pp. 14096-14113. PMLR, 2023.
Eric Jang, Shixiang Gu, and Ben Poole. Cateooriaal reparaaettttaaannneeiiiatiooo with Gunbel-softmax, arXiv preprint 1ixi::611.00 144, 2016.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for im- proved quality, stability, and variation. International Conference on Learning Representations (ICLR), 2018.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. Proceedings of the IÉEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4401-4410, 2019.
Nam-Gyu Kim, Deok-Hyeon Cho, Seung-Bi Kim, and Seong-Woin Whan Lee. Spotlight- -TSi: Spot- lighting the style via voiced-aware style extraction and style direction adjustment for expressive text-to-speech. arXiv preprint arXiv:2505 20868, 2025.
Diederik Kingma and Max Welling.
Auto-encoding variational Bayes.
arXiv preprint
arXhv:1312.6114.20 2013.
Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, and Neil Houlsby. UViM: A unified modeling approach for vision with learned guiding codes. In Advances Neural Information Processing Systems 35 (NeurIPS), pp. Associates, Inc., 2022.
Curran
Rithesh Kumar, Prem Seetharaman, Alejandro Luebs, Ishaan Kumar, and Kundan Kumar. High- fidelity audio compression with improved RVQGAN. In Advances in Neural Information Pro- cessing Systems 36 (NeurIPS), pp. 27980-27999, Curran Associates, Inc., 2023.
12

---

Published as a conference paper at ICLR 2026
Adrian Lańcucki, Jan Chorowski, Guillaume Sanchez, Ricard Marxer, Nanxin Chen, Hans JGA Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent. Robust training of vector quan- tized bottleneck models. In International Joint Conference on Neural Networks (IJCNN) pp. 1-7. IEEE, 2020.
Doyup Lee, Chiheon Kim, Sachoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 11523-11532, 2022.
Jun-Woo Lee, Min Lee, Dong-Joo Min, and Yongchae Cho. Reviving legacy seismic data via machine learning technique part I: Expanding 3D seismic survey coverage with gated convolution GAN. IEEE Transactions on Geoscience and Remote Sensing. 2024.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations (ICLR), 2017.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantiza- tion: VQ-VAE made simple. arXiv preprint arXiv:2309, 15505, 2023.
Gaurav Parmar, Richard Zhang, and Jun-Ye Yan Zhu. On aliased resizing and surprising subtleties in
GAN evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, 2022.
Claude E. Shannon. Coding theorems for discrete source with a fidelity criterion. IRE National Convention Record, 4:142-163, 1959.
Yuhta Takida, Takashi Shibuya, WeiHsiang Liao, Chieh-Hsin Lai, Junki Ohmura, Toshimitsu Ue- saka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. SQ-VAE: Variational bayes on discrete representation with self-annealed stochastic quantization. arXiv preprint arXiv:2205.0754 2022.
Mohammad Hassan Vali. Vector Quantization in Deep Neural Networks for Speech and Image Processing. PhD thesis, Aalto University, 2025.
Mohammad Hassan Vali and Tom Bäckström. NSVQ: Noise substitution in vector quantization for machine learning. IEEE Access, 10:13598-13611 2022.
Mohammad Hassan Vali and Tom Bäckström. Interpretable latent space using space-filling curves for phonetic analysis in voice conversion. In Proceedings of Interspeech, 2023a.
Mohammad Hassan Vali and Tom Bäckström. Stochastic optimization of vector quantization meth- ods in application to speech and image processing. In Proceedings of ICASSP, 2023b.
Aaron den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation leam- ing. In Advances in Neural Information Processing Systems 30 (NeurIPS), pp. 6309-6318. Curran Associates, Inc., 2017.
Harry Walsh, Abolfazl Ravanshad, Mariam Rahmani, and Richard Bowden. A data-driven repre- sentation for sign language production. In 18th International Conference on Automatic Face and Gesture Recognition, pp. 1-10. IEEE, 2024.
Haishan Wang, Mohammad Hassan Vali, and Aro Solin. Compressing 3D Gaussian splatting by noise-substituted vector guantization. In Scandinavian Conference on Image Analysis, pp. 338- 352. Springer, 2025.
Rongkun Xue, Yazhe Niu, Shuai Hu, Zixin Yin, Yongqiang Yao, and Jing Yang. HH-Codec: High compression fighliidelity discrete neural codec for spoken language modeling. arXiv preprint (tl1:2507.18891 2025.
13

---

Published as a conference paper at ICLR 2026
Junichi Yamagishi, Christophe Veaux, and Kirsten MacDonald. CSTR VCTK corpus: English multi-speaker corpus for CSTR voice cloning toolkit (version 0.92). The Centre for Speech Tech- nology Research, University of Edinburgh.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. VideoGPT: Video generation using VQ-VAE and transformers. arXiv preprint arXiv:2104. 10157, 2021.
Seonghyeon Ye, Joel Jang, Byeongguk Jeon, Sejune Joo, Jianwei Yang, Baolin Peng, Ajay Man- dlekar, Reuben Tan, Yu- We-Wei Chao, Yuchen Lin, Lars Liden, Kimin Lee, Jianfeng Gao, Luke Zettlemoyer, Dieter Fox, and Minjoon Seo. Latent action pretraining from videos. In Interna- tional Conference on Learning Representations (ICLR), 2025.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song. Thomas Funkhouser, and Jianxiong Xiao. LSUN: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 150.vv:506.000 530.03335, 2016
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. pector-quantized image modeling with improved VQ- GAN. arXiv preprint arXiv:21 1et::21100.4422 2021.
Lijun Yu, José Lezama, Nitesh B Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Vighnesh Birodkar, Agrim Gupta, Xiuye Gu, Alexander Hauptmann, Boqing Gong, Ming-Hsuan Yang, Irfan Essa, David Ross, and Lu Jiang. Language model beats diffusion- tokenizer key to visual generation. arXiv preprint arXiv:2310.05737, 2023.
Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Sound- Stream: an end-to-end neural audio codec. arXiv preprint arXiv:2107.03312. 2021.
Jiahui Zhang, Fangneng Zhan, Christian Theobalt, and Shijian Lu. Regularized vector quantization for tokenized image synthesis. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 18467-18476, 2023.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features a perceptual metric. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2018
Yue Zhao, Yuanjun Xiong, and Philipp Krähenbühl. Image and video tokenization with binary spherical quantization. arXiv preprint arXiv:2406.07544 075448 2024.
Chuanxia Zheng and Andrea Vedaldi. Online clustered codebook. Proceedings the IEEE/CVF International Conference Computer Vision (ICCV), pp. 22798-22807, 2023.
the
Lin Zhu, Weifeng Zhu, Shuowen Zhang, Shuguang Cui, and Liang Liu. Scalable transceiver design for multi-user communication in FDD massive MIMO systems via deep learning. arXiv preprint arXiv:2504.11 2025.
Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang. Yinpeng Chen, Lu Yuan, and Gang Hua. Designing a better asymmetric VQGAN for stablediffusion. arXiv preprint arXiv:2306.0440 04632, 2023.
14

---

APPENDICES
Here, we provide an overview of the contents covered the Appendices. In App. A, we provide comprehensive and detailed explanation of the implementations for training the models in VQ-VAE image compression, VQGAN image generation, and DAC-based speech coding tasks. In App. B, we present other proposals of the paper including (i) new proposed codebook replacement (App. B.1), and (ii) new variants of DiVeQ and SF-DiVeQ techniques (App. B.2). App. C provides additional qualitative and quantitative evaluations of all VQ optimization methods to substantiate the superior performance of our proposed DiVeQ and SF-Di techniques in the established VQ-VAE image compression, VQGAN image generation, and DAC speech coding tasks.
LIST OF APPENDICES
A
Implementation details
15
VQ-VAE compression
16
A.2
VQGAN generation
17
A.3
DAC speech coding
18
A.4
Gumbel-Softtax implementation
19
A.5
Rotation Trick implementation
19
A.6
Suggestions for training the codebook using SF-DiVeQ
19
B Other proposals
20
B.1
New proposed codebook replacement
20
B.2
Proposed variants of DiVeQ and SF-DiVeQ
21
C Additional results
23
C.1
Quantitative results on other data for VQ-VAE compression
23
C.2
Quantitative results on other data sets for VQGAN generation
25
C.3 Ablation batch size
25
C.4 Ablation on learning rate
27
C.5 Ablation on variance x² of DiVeQ and SF-DIVeQ
28
C.6 Ablation codebook replacement
31
C.7 Ablation on SF-DiVeQ initialization
32
C.8 Misalignment of codebook and latent representations
35
C,9 DiVeQ and SF-DiVeQ Residual Vector Quantization
36
C.10 Training logs
37
C.11 More samples generated by VQGAN
39
Disclosure of the use of large language models
39
A IMPLEMENTATION DETAILS
In this paper, establish the VQ- VAE image compression, VQGAN image generation, and DAC speech coding tasks to compare the performance of different VQ optimization methods within these three frameworks. This section provides the implementation details of these three tasks. There are some hyperparameters that are shared for all three of these tasks, which are mentioned in the following.
Mutual hyperparameters for different tasks In all VQ-VAE compression, VQGAN generation, and DAC speech coding tasks, we set the loss coefficients Eqs. (2) and (6) as a P and 3 = 0.25. The EMA decay factor is fixed at y 0.99 for all experiments in the paper. The ST-GS temperature exponentially annealed from Tuan to Tmin 0.1 over the training epochs (for more details see App. A.4).
15

---

Published as a conference paper at ICLR 2026
A.1 VQ- VAE COMPRESSION
Hyperparameters In VQ-VAE compression task (for Fig. 6 in Sec. 4 and Figs. 12 to 15 in App. C.I), we train all VQ methods for 100 epochs with the batch size of 32 using Adam opti- mizer with the initial learning rate of Ir = 5.5 10-1 that is halved after 40 and 70 epochs. All these experiments are done over cight different codebook bits of example. for B = 8, the number of codewords equals K 2 = 256. To demonstrate that our pro- posed De and SF-DiVeQ techniques perform consistently superior to the other VQ optimization methods in the image compression task, we do ablation studies over different batch sizes and learn- ing rates which are presented in App. C.3 and App. C.4, respectively. The variance o² of directional noise Va for DiVeQ and SF-DiVeQ is set to o² 10 for VQ-VAE compression (see the ablation study on d² in App. C.5).
{4. 5, 6, 8, 9, 10, 11). For
Our proposed codebook replacement (App. B.1) is actively applied during training for all VQ op- timization methods (except SF-DiVeQ, which does not require it). The replacement is done in two different phases of
iter S 2000
replacement after each 100 training iterations
2000 < iter S Na 1000 replacement after each 500 training iterations,
where iter is the training iteration number, and Niter is the total number of training iterations. The discarding threshold for the replacement equals 0.01, which means that the codebook replacement discards the codewords that are used less than 1 % during the period that the replacement is done.
Model architecture For the image compression task, we aim use the original VAEV proposed in van den Oord et al. (2017). The VQ-VAE implementation is provided in DeepMind's GitHub', and its PyTorch version in zalandoresearch GitHub repository?. These implementations for images of size x32. Hence, to make the VQ-VAE suitable for compressing images of size 256 x 256, we make some minor modifications these implementations and add the LPIPS (VGG- 16) as the perceptual loss to enhance the quality of reconstructions.
meant
Table shows the architecture of the VQ- VAE used in our image compression experiments. The encoder consists of four downsampling blocks, a stack of six residual blocks, and pre- VQ Conv2D layer. Each downsampling block is a strided 2D convolutional layer with kernel size 4 and stride = 2 followed by ReL.U activation function. Each residual block is comprised of ReL.U, 3x3 Conv2D (stride 1), ReI.U, and 1x1 Conv2D (stride = 1) in this order. The pre-VQ Conv2D layer is meant map the input channels to match the VQ embedding dimension. The decoder has a symmetric architecture to the encoder, but a reverse order, with difference that uses transposed 2D convolutions in the upsampling blocks.
Table 4: Architecture of the VQ-VAE model used in the image compression task.
Encoder
Decoder
E HX
Zq E RxxxxxC"
4 X (Downsampling Block] + Rhxuxo
Conv2D + RhxwxX
6 X [Residual Block) - cxxxx
6 x {Residual Block) + Axxxxx
Conv2D + RR x
4 x (Upsampling Block) + HXXCC
W 256, w 16,
3,
C' 256,
C" 512
httppp////hhuuuc.ccooooo/ogeeee/ example. ipynb httpps://githhine com hawaa//aaacceeerrr/////orrrooggo
eevvveee__ /vvaa__
16

---

Published as a conference paper at ICLR 2026
A.2 VQGAN GENERATION
Hyperparameters In VQGAN generation task, we train the generator model (i.e.. VQ-VAE) with different VQ optimization methods for 100 epochs over codebook bits of B {8,9, 10, 12) and over two different hyperparameter settings of HPI (tr= 2.5 10-3; batch size = 8) and HP22 (Ir= 2.5-10- -4 batch size= 32). During training. the imitial learning rate is halved after 50 and 75 epochs. The discriminator starts at epoch 50 when half of the training is passed. All GPT transformers (with different configurations in Table 5) are trained with a batch size of 32 and an initial learning rate of Ir = 4.5 10-3 with cosine decay that reduces the learning rate to I 1% of its initial value. For sampling from VQGAN, we use the temperature of = 1.0 with different top-k values over different codebook sizes (see Table 6). The variance x² of directional noise Va for DiVeQ and SF-DiVeQ is set to 0² = 10 2 for VQGAN generation.
When training the generator (VQ- VAE) model for all VQ optimization methods (except SF-DiVeQ), we apply our proposed codebook replacement (App. B.1) actively during training in two phases of
iter S 5000
replacement after each 50 training iterations
50000 < iter S Na 3000 replacement after each 300 training iterations,
where iter is the training iteration number, and Naer is the total number of training iterations. The discarding threshold for the replacement equals 0.01, which means that the codebook replacement discards the codewords that are used less than 1 % during the period that the replacement is done.
Models architecture For the VQGAN generation task, we use the same architectures as in Esser et al. (2021) for the generator (VQ- (A), discriminator, and transformer models. We adopt the implementation from dome272 GitHub repository" with some suggested minor modifications from this repository*.
Training transformers For all data sets, the GPT transformer is trained with a batch size of 32 and an initial learning rate of Ir = 4.5- 10-5 with a cosine decay that reduces the learning rate to 1% of its initial value over the course of training. In addition, we adopt a linear warm-up for the learning rate such that ir starts with really small value and linearly reaches its initial learning rate during the first 10k training iterations.
important variable to tune is the probability of keeping the ground truth indices, which is called Pkeep E (0, 1). In other words, the true indices of the tokens masked with the probability of 1 Pkeep- The reason to mask the true tokens is that it helps the transformer to become more robust against its own mistakes, such that after one wrong prediction during sampling, it does not propagate this error with snowball effect all other next predictions. In addition, masking true tokens helps the transformer not to memorize the token dependencies, which leads to better generalization and prevents overfitting.
When training the transformer,
For all transformers, we set a scheduled masking rate that increases the Pkeep from the initial value of 0.5 0.95 over training epochs. In the beginning, we mask half of the true token indices, foreing the transformer to learn large missing parts, i.e., the transformer learns dependencies of coarse structures in the images. As the training goes to the end, small portion of token indices will be masked, and as result, the transformer learns dependencies of finer details in the images.
To train the transformers for FFHQ and CALEBA-HQ data sets, we use the configurations suggested in Esser et al. (2021). Also, we use the same transformer configuration for the LSUN Bedroom and LSUN Church data sets. Since the AFHQ data set is much smaller than the other data sets (containing 15803 images), we adopted different and lighter transformer for it. Table 5 shows the transformers configuration for different data sets. Note that when experimenting with different codebook sizes, we keep the transformer configuration fixed.
is noteworthy to mention again here that since having more data favors for the VQGAN generation task, for CELEBA-HQ and AFHQ data sets, we use the full data sets to train the generator (VQ- VAF) and transformer, and to compute the FID scores. CELEBA-HQ and AFHQ full data sets contain 30k and 15803 images, respectively. However, for FFHQ and both LSUN data sets, we only consider the train sets that contain 56k images.
hhttps://gt thub. com 123422 1234221//42444n /vvgan
17

---

Published as a conference paper at ICLR 2026
Table 5: Configurations used to train the transformers for different data sets in the VQGAN gener- ation task. Mayer is the number of transformer blocks, Rhead refers to the number of attention heads, D Demeedding the dimensionality of the transformer embeddings, s is the length of the sequence, #params is the number of transformer parameters, dropout refers to the dropout rate used to train the transformer, and Dutem is the dimensionality of each codeword.
Data
Flayer
Thhaad
eembedding
length (s)
dropout
4params [M]
Duttm
CELEBA-HEE -HQ
28
16
1024
256
0.1
= 355
256
FFHQ
28
16
1024
256
0.1
&: 355
256
LSUN Bedroom
28
16
1024
256
0.1
R 355
256
LSUN Church
28
16
1024
256
0.1
R 355
256
AFHQ
20
16
768
256
0.1
RJ 143
256
Sampling from VQGAN When the transformers are trained, for each data set, sample the same number of images as the train set (full data sets for CELEBA-HQ and AFHQ) with the configuration mentioned in Table 6. To compute the FID score between the train set and new sampled images, use clean-fid fid from GaParmar GitHub repository?.
Table 6: Configuration used for generating images from trained VQGAN models for different code- book sizes. This configuration fixed for different data sets.
Codebook bits
Codebook size (K)
Sampling temperature (f)
top-k
256
1.0
75
512
1.0
150
1024
1.0
300
12
4096
1.0
300
A.3 DAC SPEECH CODING
Hyperparameters In the speech coding task (for Table Sec. 4), we train all VQ methods for 300 epochs with the batch size of 64 using the AdamW optimizer with learning rate of (r 10-4 and betas = (0.8, 0.99). All these experiments are done over four different codebook bits of = {10, 11, 12, 13). For example, for B 10 the number of codewords equals K 2 1024. To demonstrate that our proposed DiVeQ and SF-DiVeQ techniques perform consistently superior to the other VQ optimization methods in the speech coding task, we do ablation studies over different batch sizes, which are presented in Tables 8 and in App. C.3. The variance x² of directional noise Va for DiVeQ and SF-DiVeQ is set σ² 10-3 in the experiments.
Note that our proposed codebook replacement (App. B.1) is actively applied during training for all VQ optimization methods (except SF-DiVeQ, which does not require it). The replacement is done with the same schedule as in the VQGAN generation task, explained in App. A.2.
Model architecture For the speech coding task, we use a modified version of the Descript Audio Codec (DAC, Kumar et al.. 2023) called Pikku NAC', which consists of convolutional encoder, a residual vector quantization (RVQ) layer, and decoder. Both the decoder and the encoder have four layers that follow the same architecture as DAC, with a stride pattern of 2:4:8:: in the encoder and 8 8 4 2 in the decoder, resulting Pikku NAC' uses 32 encoder channels, 256 decoder channels, and a kernel size of 5, resulting in a significantly smaller model with only 5M encoder parameters and 2M decoder parameters. The 8-dim latent space is quantized with 9-layer RVQ (i.e., using 9 codebooks), where each codebook contains 1024 codewords, resulting in 156k parameters.
frame size of 512. However, unlike the DAC model,
In our experiments, change the residual VQ layer in Pikku NAC ordinary VQ layer, which contains only one codebook. To compensate for the loss in the number of codewords used for the quantization, did the experiments using large codebook sizes (i.e., K [2 221,211111 212 213)). Both the DAC and Pikku NAC models project the 512-dimensional latent space to 8 dimensions
http://hww...cmmmmaaaaaaa//aaaaa fid
'https: :/g/ thub. http://ttw.......o.///////////a/
18

---

Published as a conference paper at ICLR 2026
(D 8) before quantization and project the quantized latent vectors back into dimension 512. However, in order not to have a tight bottleneck and make the quantization task more challenging, keep the latent space to be of dimension 512. We achieve this by keeping the convolution layers that apply the projections, but making them project to the same dimension of 512.
Data As mentioned in Sec. 4, for the speech coding task, ishi et al.) with an 80/20% train-test split without any overlapping speakers and speech files between train and test sets. This data set contains 109 native English speakers with different accents, each reading about 400 sentences. As an additional curation step, we used signal-to-noise ratio (SNR) thresholding to remove silent segments between utterances and discarded samples with significant background noise. Since Pikku NAC is prepared for training wideband speech that is sampled at 16 kHz, we downsample the data during training to meet this sampling rate.
the CSTR VCTK data set' (Yamag-
A.4 GUMBEL-SOFTMAX IMPLEMENTATION
For VQ- VAE compression, VQGAN generation, and DAC speech coding, the implementation of Straight-Through Gumbel-Softman (ST-GS) is adopted from karpathy GitHub repository The ST- GS temperature r is exponentially annealed from Tstart 1 to Tmin = 0.1 over training epochs such that
max sttar * 17spuch Tmin
min
spocis
(13)
Tstart
where 17 is the decay rate for the temperature annealing, epoch is the current epoch number, and Nepochs is the total number of training epochs.
A.S ROTATION TRICK IMPLEMENTATION
We adopt the implementation of the Rotation Trick (RT) from lucidrains GitHub repository". Fol- lowing what is suggested in the RT paper (Fifty et al., 2025), for our VQ-VAE compression task, first trained the VQ-VAE by updating the codebook via Exponential Moving Averages (EMA) with the decay rate of 7 0.8. Since the results were not good enough, we changed the decay rate to y 0.99, which slightly improved the quality of reconstructions. However, we finally noticed that by skipping the EMA for updating the codebook and using codebook loss (2nd loss term in Eq. (2)), we achieve the best possible performance that we can get from the Rotation Trick in our VQ-VAR compression task. For the VQGAN generation task, the RT paper suggested updating the VQ codebook via gradients using codebook loss. Therefore, for both of our VQ-VAE compression and VQGAN generation tasks, we train the RT using the codebook loss.
A.6 SUGGESTIONS FOR TRAINING THE CODEBOOK USING SF-DIVEQ
As discussed Sec. 4, for training the codebook via SF-DiVeQ, it is recommended skip quan- tization for several initial epochs (or several thousand training iterations) in order to reach a more stable latent representation that already passed its initial big distribution shifts. After that, initial- ize the codebook vectors with the average of recent latent vectors and start discretizing the latent representation using the SF-DiVeQ approach.
For the initialization of the codebook, it is highly recommended follow these two guidelines; (i) take the very recent latent vectors from the last 20 to 50 training iterations as the set of vectors for initialization, (ii) according to the batch size, initialize each codebook vector to be the average of at least 20 to 40 recent latent vectors. For instance, suppose the encoder compresses the input image of size 256x256
latent of size 16. So, for batch size of 8, each training iteration would generate 8x 16x 16 = 2048 latent vectors. Hence, for a codebook of size K = 1024, to initialize each codebook vector with the average of 40 latent vectors, we need capture the latent vectors of the last 20 training iterations such that
Nittnts 40 K
Niatents 40 1024
N
20
(14)
2048
2048
19

---

Published as a conference paper at ICLR 2026
where Musennss is the number of recent captured latent vectors, and Niwes is the number of recent training iterations required for obtaining Nanents latent vectors. Therefore, for smaller batch sizes, better to capture more training iterations compared to larger batch sizes.
The most important point in training codebook via SF-DiVeQ is to track the codebook usage per- centage over training epochs (the plot at the top in Fig. 29). If quantization with SF-DiVeQ is started after big latent distribution shifts and by capturing enough recent latent vectors, the initial codebook usage would be high value (approximately close the maximum possible value that equals the number of codebook vectors K). Then, according to the SF-DiVeQ property, it pulls the codewords that are outside the latent representation inside during training iterations. Therefore, the codebook usage will reach its maximum value K after some epochs of starting to quantize the latent space, and will be kept its maximum until the end of training. Fig. 29 presents an example how codebook usage percentage looks for proper SF-DiVeQ training.
B
OTHER PROPOSALS
B.1 NEW PROPOSED CODEBOOK REPLACEMENT
According to Huh et al. (2023), codebook replacement is suitable solution to avoid codebook col- lapse (Vali, 2025; Mentzer al., 2023) for different VQ optimization techniques. In NSVQ (Vali & Bäckström, 2022), a codebook replacement method is proposed such that after a period of training iterations, unused codewords are replaced with permutation of some used codewords. Selection from the used codewords is completely random, regardless of the importance of the used codewords. However, in this paper, we propose a new codebook replacement technique that replaces unused codewords with the used ones based on the importance of the used codewords. In other words, use an importance sampling strategy that computes the probability of codewords occurrences (dur- ing period of training iterations), and replaces the unused codewords by selecting from the used codewords according to these probabilities. Therefore, a codeword with higher occurrence proba- bility is more likely to be selected for the replacement. This logical approach for replacement, as in quantization applications, the main goal is to reduce the average quantization error, and when the number of assigned inputs to codeword high, it is better define new codeword in that region to reduce the average quantization emror
Perplexity (average codebook usage)
Proposed codebook replacement NSVQ codebook replucement
200
400
600
800
1,000
1,200
1,400
Training iteration
Figure 8: Proposed codebook replacement achieves more stable codebook utilization faster than NSVQ's replacement method. Comparison the perplexity of our proposed codebook re- placement with NSVQ's replacement technique for the initial 1500 iterations when training the VQ- VAES on CELEBA-HQ data set using the DiVeQ approach. The codebook size equals K 2048.
Fig. 8 compares the perplexity (or average codebook usage) our proposed codebook replacement
and NSVQ's replacement technique (Vali & Bäckström, 2022) in the VQ-VAE compression of the CELEBA- -HEBA-HHQ&A data set using the DiVeQ approach with an 1 11-bt codebook (with K = 2 = 2048 codewords). For each training iteration, perplexity shows the number of selected codewords involved
20

---

Published as a conference paper at ICLR 2026
in quantizing that batch of data, and it is computed as
K perplexity - (C)(H((C) S.L. H(C) = Σ
log Pk.
(15)
where (C) the entropy of the codebook index distribution, and Pi: is the empirical usage prob- ability of the k-th codebook vector. Fig. 8 is plotted for the first 1500 training iterations, while the curves are smoothed via Blackman windowing with a window size of 11 and normalized by their maximum value. According to the figure, our proposed codebook replacement reaches higher perplexity values faster than NSVQ's 's replacement approach, and as a result, provides a longer training opportunity for the active codewords to be trained for more training iterations.
Fig. 9 shows the quantitative results of the reconstructions when using our proposed codehook re- placement and NSVQ's replacement method, shown in Fig. 8. We observe that our proposed code- book replacement results in slightly better performance by having higher SSIM, PSNR, and lower LPIPS values. Note that based on our investigations, our proposed codebook replacement can result in similar or better performance than NSVQ's replacement over different experiments. Therefore, we use our proposed codebook replacement for all VQ optimization techniques (i.e. STE, EMA, RT, ST-GS, NSVQ, and DiVeQ) and for all experiments in this paper.
0.74
0.3
0.72
-
LPIPS
0.28
SSIM
PSNR
0.7
0.26
0.08
NSVQ codehook
5
6
8
9
10 11
Codebook bits
Codebook bits
Codebook bits
Figure 9: Proposed codebook replacement improves the quality of reconstructions compared to NSVQ's replacement method. Quantitative comparison of reconstructed images from the CELEBA-HO -HQ test set when using our proposed codebook replacement vs. NSVQ's replacement method. The VQ-VAEs trained via the DiVeQ approach over different codebook sizes.
B.2 PROPOSED VARIANTS OF DIVEQ AND SF-DIVEQ
Apart from what is proposed for the formulation of our proposed DiVeQ (Eq. (8)) and SF-DiVeQ (Eq. (12)), also propose new variants of them by skipping the directional noise v. and using the stop gradient operator (i.e., detach in Py Porc))). We call these two variants as DiVeQ-detach and OF-DiVeQ-detachh and their formulation
C₁-
DiVeQ-detach: Zg = z + ||c.- 21122 sg
|c-- x||2
[1 \1- (Ci- - 2)
SF-DiVeQ-detachh 2q ||cc- - 2112
(16)
||: 2112
Ai- (c+++1 2)
||| =1|2 sg
| - 2 2112
21

---

Published as a conference paper at ICLR 2026
24
0.6
0.6
22
PSNR -
LPIPS
0.5
SSIM
STE
20
EMA
<-
ST
18
0.4
NSVQ
0.4
Diee-letuch (ours)
SF-DiNeQ-deteee (ours)
16
0.3
5
6
10
6
7
7
8
9
Codebook bits
Codebook bits
Codebook bits
Figure 10: Variants of our proposed methods: SF-DiVeO-detach consistently improves image reconstruction and DiVeQ-detach maintains the reconstruction quality. Quantitative compar- ison of reconstructed images from the AFHQ test set using variants of our proposed methods vs. other VQ optimization approaches. Each curve is the average of the metric over all test recon- structions and over three different individual runs.
We experiment with these two variants only in the VAE compression task to evaluate how they perform compared to DiVeQ and SF-DiVeQ, and also other existing techniques. We train the VQ- VAE compression model with learning rate of 5.5 10 and batch size of 32 over 100 epochs on the AFHQ data set. Fig. 10 shows the quality of reconstructions for DiVeQ-detach and SF-DiVeQ- detach compared to other approaches of STE, EMA, RT, ST-GS, and NSVQ. According to figure, similar to its original version, SF-DiVe consistently outperforms other methods for all three objective metrics and across all different codebook sizes. However, DiVeQ-detach performs comparable to other VQ methods. In addition, Fig. 11 compares the quality of reconstructions of -ieetddttach and SF-DiVee--dtaach with their original approaches of DiVeQ and SF-DiVeQ. According Fig. 11, DiVeQ performs clearly superior to its variant (i.e., DiVeQ-detack), while SF-DiVeQ and its variant perform almost similarly.
DiVeQ
0.38
SF-DiveQ
deVeC--etacch
ST-DiNeQ-ddttee
24
0.36
f- LPIPS
SSIM
0.6
PSNR
0.34
23
0.32
0.55
0.3
10
Codebook bits
Codebook bits
Codebook bits
Figure 11: Variants of proposed methods: Regarding reconstruction quality, SF-DiVeQ- detach performs almost similar and DiVeQ-detach performs worse than their original versions. Quantitative comparison of reconstructed images from the AFHQ test set using variants of our pro- posed methods vs. their original versions. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
22

---

Published as a conference paper at ICLR 2026
According to VQ-VAE compression experiments other data sets of CELEBA- -QQ, FFHQ. LSUN Bedroom, and LSUN Church, we find that DiVeQ and SF-DiVeQ always result in higher quality reconstructions compared to their variants. The performance gap is not significant, and even in some cases (like SF-DiVeQ in Fig. 11), DiVeQ-detach and SF-DiVeQ-detarde can perform close to their original versions. The main difference is that the detach variants do not need o² (in directional noise va). However, we do not consider this an advantage, since the variance in DiVeQ and SF-DiVeQ can be used to control the amount of generalization for the downstream task.
the variance
C ADDITIONAL RESULTS
C.I QUANTITATIVE RESULTS ON OTHER DATA SETS FOR VQ- VAE COMPRESSION
Apart from the results obtained for the AFHQ data set in Fig. 6, we also assess the quality of recon- structions in the VQ-VAE compression task for CELERA-HQ. FFHQ, LSUN Bedroom, and LSUN Church data sets. As mentioned earlier, for all data sets, we train the VQ-VAE models using the configurations mentioned in App. A.1. After training, we compress and reconstruct the test im- ages by doing hard VQ (using argmin) in the latent space. Then, we compute the SSIM, PSNR, and LPIPS metrics for the reconstructions by having the original ground-trute -tuth image available. Figs. 12 to 15 show the results for CELEBA-HO HQ. FFHQ, LSUN Bedroom, and LSUN Church data sets, re- spectively. Note that in the plots, the reported value for each metric is the average of that metric over all test set reconstructions and over three different individual runs.
0.6
0.7
24
0.5
PSNR ->
LPIPS
SSIN
22
0.4
0.0
EMA
20
ST-GS
0.3
NSVQ
DiVeQ
ST. DiVeQ (ours)
4
5
6
7
10
11
4
5
8
9
10
11
5
6
7 8
9
10
11
Codebook bits
Codebook bits
Codebook bits
Figure 12: The proposed DiVeQ and SF-DiVeQ lead to consistent improvement in image re- construction quality. Quantitative comparison of reconstructed images from the CELEBA-HO HQ test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
23

---

Published as a conference paper at ICLR 2026
0.7
0.6
u.5
SSIM -
PSNR +
LPIPS
STE
0.4
RT
ST&GS
NSVQ
0.33
DIVeQ (ours)
DI-DiVeQ (ours)
15
Codebook bits
Codebook bits
Codebook bits
Figure 13: The proposed DiVeQ and SF-DiVeQ lead to consistent improvement in image construction quality. Quantitative comparison of reconstructed images from the FFHQ test set for different VQ optimization methods different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
24
0.6
22
PSNR +
LPIPS
U.5
SSIM
EMA
RT
0.4
NSVQ
DIVe (ours)
0.3
5
6
9
10
11
5
6
7
Codebook bits
Codebook bits
Codebook
Figure 14: The proposed DiveQ and SF-DiVeQ lead to consistent improvement in image recon- struction quality. Quantitative comparison of reconstructed images from the LSUN Bedroom test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
0.35
0.6
20
0.5
PSNR ->
LPIPS
0.5
SSIM
0.45
STE EMA
0.4
ST-GS NSVQ
3
DIV.Q (ours)
Si-DiVeQ (ouns)
0.3
6
7
10
11
8
9
10
11
5
6
7
8
Codebook bits
Codebook bits
Codebook bits
Figure 15: The proposed DiveQ and SF-DiVeQ lead to consistent improvement in image re- construction quality. Quantitative comparison of reconstructed images from the LSUN Church test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
24

---

Published as a conference paper at ICLR 2026
C.2 QUANTITATIVE RESULTS ON OTHER DATA SETS FOR VQGAN GENERATION
we also evaluate the quality of generations for other data sets by computing the FID scores. As mentioned earlier, for all data sets, we train the VQGAN generator (i.e., VQ- TQ-VAE), discriminator, and transformer models using the configurations mentioned in App. A.2 over two different hyperparameter settings of HP1: (Ir= 2.5 100 batch size and HP2: 2: (Ir 2.5 -10 batch size=32) Then after the models are trained, we sample new generations from them and compute the FID value between the original training
Apart from the FID values obtained for CELEBA-HQ data set in Table 2,
and the generated samples (see Table 6 for sampling configuration). Table presents the obtained FID scores for FFHQ, AFHQ, LSUN Bedroom, and LSUN Church data sets using different VQ optimization methods and various codebook sizes. Figs. 30 to 34 provide the qualitative comparison of VQGAN generations for all data sets with 9-bit codebook.
The proposed DiVeQ and SF-DiVeQ robustify training and maintain the generation quality for challenging small codebooks, while being less sensitive to hyperparameeee setting, The table presents the FID., scores for different VQ optimization methods over different codebook sizes for FFHQ, LSUN Bedroom, LSUN Church, and AFHQ data sets. Values in red refer to the cases that misalignment happens (see Fig. 4).
Table 7:
2.5-10 batch size =
batch size 32
Codehook bits
Codebook bits
Data
Approach
8
10
5.36
356
EMA
6.53
6.57
6,15
619
8.65
9.00
8.35
11.5
KT
14.2
7.19
6.52
63.3
11.7
9.35
7.86
FFHQ
ST-GS
22.5
21.9
17.0
54.4
274
36.8
70.6
NSVQ
98.0
88.8
83.8
65.1
84.0
76.0
71.4
63.1
DiVeQ (ours)
SF SE-DiVeQ (ours)
7.91
6.21
5.93
8.60
10.3
9.41
8.04
9.74
STE
5.27
4.89
4.96
7.54
225
67.5
6.38
8.04
EMA
6.04
5.36
5.56
5.13
8.05
584
584
584
LSUN
RT
17.4
8.65
6.69
255
42.9
25.5
ST-GS
35.9
37.2
36.4
23.0
172
223
187
210
Bedroom
NSVQ
DiVeQ (ours)
5.71
4.87
5.40
7.94
7,34
347
6.58
9.03
SF-DiVeQ (ours)
5.96
6.01
5.40
7.79
7.16
6.69
6.36
7.62
STE
4.33
3.74
4.09
4,91
6.69
9.92
5.67
5.92
EMA
5.24
RT
9.33
6.37
5.20
4_31
78.5
142
134
259
LSUN
Church
26.5
NSVQ
107
92.8
71.7
50.9
85.6
76.9
61.6
47.8
DIVeQ (ours)
4.41
3.92
3,90
4.58
7.91
228
9.17
6.11
SF-DiVaQ (ours)
4.35
3.99
3,56
4,40
7.74
5.74
10.9
4,57
250
EMA
5.12
5.56
5.23
6.95
6.68
6,60
6.89
8.59
RT
6.31
5.55
250
248
288
56.3
AFHQ
ST-GS
31.3
26.5
20.9
14.4
302
50V
231
224
NSVQ
92.0
81.3
72.2
47.1
70.0
64.0
$3.0
45.0
Di (ours)
6.30
5.63
5.78
7.60
6.82
6.75
8,49
SF DiVeQ (ours)
5.63
5.40
5.69
6.86
7.95
7.04
6.87
8.39
C.3 ABLATION ON BATCH SIZE
In this section, we study the effect of the batch size on different VQ optimization methods. To this end, in the VQ-VAE compression task, we train all different methods with batch sizes of 64 and 128 at the learning rate of |r = 5.5 10-4 on the AFHQ data set for 100 epochs. The learning rate is halved after 40 and 70 epochs. Fig. and Fig. show the quality of reconstructions for batch sizes 64 and 128, respectively. Note that in the plots, the reported value for each metric is the average of that metric over all test set reconstructions and over three different individual runs. Taking Fig. (for batch size 32), Figs. 16 and 17 (for batch sizes 64 and 128) into account, can claim that by increasing the batch size, the performance gap between our proposed methods and other approaches (except EMA) becomes larger. For larger batch sizes, EMA performs comparably but slightly worse than our DiVeQ and SF-Di DiVe.. The reason is that EMA updates the codebook with the average of latent vectors, and as result, larger batch sizes yield less noisy and more representative updates.
25

---

Published as a conference paper at ICLR 2026
24
0.6
0.6
22
+ LPIPS
0.5
SSIM
03
PSNR
20
EMA
RT
0.
NSVQ
DeQ
16
ST. DiVeQ (ours)
0.3
5
6
5
10 11
10
Codebook bits
Codebook bits
Codebook bits
Figure 16: Ablation on batch size: The proposed DiVeQ and SF-DiVeQ lead to consis- tent improvement in image reconstruction quality over different batch sizes (this figure, batch size 64). Quantitative comparison of reconstructed images from the AFHQ test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
In the VQ-VAE compression task, it is expected that the increase in codebook size always enhances the quality of reconstructions. However, this does not happen for some in Fig. 16 and Fig. 17. The reason is that the codebook representation C₂ is not well fitted to the latent distribution P., and as result, a misalignment between codebook and latent representations happened (see Fig, 4, Sec. 4, and App. C.8 for more details). In Fig. 16, the misalignments happen for [ST-GS] codebook bits=10, 111, [RT codebook bits=5], and in Fig. 17 they arise for [STE | codebook bits=11], [EMA codebook bits=11), [ST-GS codebook bits=9, 11].
24
0.6
0.G
22
LPIPS
0.5
SSIM
PSNR
0.5
STR
EMA
0.
ST-GS NSVQ
DeQ (ouns)
ST DiVeQ (ours)
5
10
4
5
10 11
4
10 11
Codebook bits
Codebook bits
Codebook bits
Figure 17: Ablation on batch size: The proposed DiVeQ and SF-DiVeQ lead to consis- tent improvement image reconstruction quality over different batch sizes (this figure, batch size 128). Quantitative comparison of reconstructed images from the AFHQ test set for different VQ optimization methods over different codebook sizes. Each eurve the average of the metric over all test set reconstructions and over three different individual runs.
We also evaluate how our proposed DiVeQ and SF-DiVeQ perform compared to other methods when changing the batch size in the DAC speech coding task. We train the DAC-based speech coding model using different VQ methods with batch sizes 32 and 16 at the learning rate of Ir 10-4 on the VCTK data set for 300 epochs. Table 8 and Table 9 show the quality of decompressed speech from the VCTK test samples for batch sizes of 32 and 16, respectively. The reported value for each metric refers to the average of that metric over all test samples. In addition, entries highlighted in red correspond to runs where codebook-latens misalignment occurs, and the decoded speech is severely degraded and unintelligible. According to the quantitative comparisons Tables 3, 8, 9,
26

---

Published as a conference paper at ICLR 2026
proposed De and SF-DiVeQ consistently result in superior quality for the decompressed speech samples when changing the training batch size.
Table Ablation on batch size: The proposed Dive and SE-DiVeQ lead to consis- tent improvement in decompressed speech quality over different batch sizes (this figure, batch size 32). Quantitative comparison of decompressed speech samples from the VCTK test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set samples.
Log spectral distance
MFCC distance..
PFSQ1
STOM
Approach bits
10
11
12
13
10
11
12
13
10
11
12
13
10
11
12
13
STE
1.13
1.07
1.05
1.06
104
68.9
86.6
88.5
1.16
1.28
1.32
1.31
0.71
0.78
0.78
0.77
EMA
1.01
1.00
1.00
1.24
75.0
72.4
70.2
147
1.51
1.57
1.61
1.08
0.83
0.84
0.85
0.59
RT
85.5
ST-GS
1.09
3.50
3.42
3.35
93.6
346
335
328
1.23
1.03
1.04
1.04
0.77
0.39
0.42
0.42
NSVQ
1.09
1.06
1.05
1.04
109
104
99.2
95.4
1.36
1.44
1.53
1.60
0.80
0.82
0.83
0.84
DiVeQ (ours)
1.03
1.01
1,00
0.99
79.5
74.8
72.1
68.9
1.41
1.54
1.66
1.74
0.81
0.83
0.85
0.85
SF SF-DIVeQ (nurs)
1.01
1.00
1.00
0.98
73.6
71.5
69.5
65.5
1.47
1.65
1.69
1.78
0.83
0.84
0.85
0.86
Table 9: Ablation batch size: The proposed DiVeQ and SF-DiVeQ lead to consis- tent improvement in decompressed speech quality over different batch sizes (this figure, batch size 16). Quantitative comparison of decompressed speech samples from the VCTK test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set samples.
Log spectral distance
MFCC distance..
PESQT
STOIT
Approach bits
10
11
12
13
10
11
12
13
10
11
12
13
10
11
12
STE
1.04
1.07
1.03
1.04
84.2
91.7
83.2
83.3
1.42
1.20
1.40
1.38
0.81
0.75
0.80
0.80
EMA
0.99
7.70
73.8
1.52
1.56
1.64
RT
1.07
1.04
1.02
1.01
94.8
86.9
79.8
77.1
1.24
1.34
1.36
1.49
0.77
0.79
0.81
0.82
ST-GS
3.36
3.45
316
328
347
1.08
1.06
1.04
1.05
0.45
0.41
0.43
NSVQ
1.06
1.05
1.04
1.02
112
108
103
95.6
1.36
1.45
1.54
1.62
0.81
0.81
0.83
0.84
DIVeQ (ours)
1.04
1.01
0.99
0.98
80.6
75.4
71.0
69.4
1.50
1.57
1.68
1.76
0.82
0.83
0.85
0.85
SF-DiVoQ (ours)
1.00
0.99
74.3
73.2
1.56
1.50
1.65
0.83
0.84
0.84
C.4 ABLATION ON LEARNING RATE
In this section, we evaluate the performance of different VQ optimization techniques in the VQ- TQ-VAE compression task on the CELEBA-HHE HHQ data set using different learning rates. The ablation is per- formed over two new learning rates of {10-3, 3,003,1 1000-) while training the models for 100 epochs with the batch size of 32. During training, the learning rates are halved after 40 and 70 epochs. Fig. 18 and Fig. 19 show the quality of reconstructions when Ir = 10-3 and Ir = 10-4, respec- tively. Note that in the plots, the reported value for each metric is the average of that metric test set reconstructions and over three different individual runs. According to Fig. 18, DiVeQ and SF-DiVeQ outperform other methods, achieving higher SSIM, PSNR values, and obtaining lower LPIPS values. According to Huh et al. (2023), when increasing the learning rate in VQ-VAEs, the model is highly prone to misaligned codebook and latent representations. Therefore, in Fig. 18, we find that some misalignments are happened for (ST-GS codebook bits-6, 8, 9, 10, 11], [EMA] codebook bits=7], and [DiVeQ] codebook bits=11). The misalignments are visible as sudden jumps in the objective metrics, and happen in cases where the increase in code- book size does lead to enhancement in the objective metries (more details in Fig, 4. Sec. and App. C.8).
all
In Fig. 19, again, our proposed methods obtain better objective metrics than other methods, while performing comparable to EMA. The reason is that EMA updates the codebook without gradient descent, and thus, there is no explicit learning rate defined for its codebook learning. The decay rate is the hyperparameter responsible for the speed of the codebook updates. As we do not have a proper definition on how to set the decay rate corresponding to a specific learning rate, we use the fixed decay rate of y 0.99 for EMA in all experiments in the ablation study learning rates (i.e., Figs. 18 and 19). According to Fig. 19, a misalignment happened for [RT I codebook bi bits==).
27

---

Published as a conference paper at ICLR 2026
25
0.7
0.5
PSNR -
+ LPIPS
SSIM
0.6
20
0.4
STB
EMA
0.
ST-GS
0.3
NSVO
DIV&Q (ours)
Si-DiveQ (ouns)
5
6
10
11
5
7
10 11
4
5
10 11
Codebook bits
Codebook bits
Codebook bits
Figure 18: Ablation on learning rate: The proposed DiVeQ and SF-DiVeQ lead to consis- different learning rates (this figure,
tent improvement in image reconstruction quality
Ir = 10 *). Quantitative comparison of reconstructed images from the CELEBA-HH HH test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
28
0.7
0.5
24
+- LPIPS
SSIM
PSNR
22
0.6
STE
FMA
ST-G8
0.3
DiVeQ (ours)
0.5
SF-DiVeQ
5
10 11
10 11
Codebook bits
Codebook bits
Codebook bits
Figure 19: Ablation on learning rate: The proposed DiVeQ and SF-DiVeQ lead to consis- tent improvement in image reconstruction quality over different learning rates (this figure, Ir 10 1). Quantitative comparison of reconstructed images from the CELEBA-HO test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
C.5 ABLATION ON VARIANCE d² OF DIVEQ AND SF-DIVEQ
To assess the impact of variance 2 in our proposed DiVeQ and SF-DiVeQ, we did ablation studies using different variances in both VQ-VAE compression and VQGAN generation tasks. This section aims to show that the performance of our proposed DiveQ and SF-DiVeQ is largely insensitive to the selection of the variance g². In other words, methods do not require tuning the variance, and using a small-enough variance (o < 100 is the only requirement to ensure DiVeQ and SF-DiVeQ perform better than the other VQ optimization methods.
28

---

Published as a conference paper at ICLR 2026
In the VQ- (QE compression task, did an ablation study with the CELEBA-HO data set using 10- }. Fig. 20 and Fig. 21 show the quality re-
different variances of o² = {10-1, 10-2, 2, 1003, 33 10 constructions over different variances for the DiVeQ and SF-DiVeQ techniques, respectively. Note that in the plots, the reported value for each metric is the average of that metric over all test set re- constructions and over three different individual runs. Conforming with what presented in Fig. 3, results in Figs, 20 and 21 confirm that to achieve higher reconstruction quality, the variance should be small to have more precise nearest-codeword mappings of inputs the selected codewords Cs,. That is why in Figs. 20 and 21, when o² 10-1 (which is still considered a high value for the variance), the reconstructions are of low quality.
0.75
28
0.5
PSNR -
LPIPS
24
SSIM
0.03
10
0.3
22
7
8
10 11
4
5
6
8
9
10
11
5
6
7
8
10
11
Codebook bits
Codebook bits
Codebook bits
Figure 20: Ablation on variance o² for DiVeQ: Quality of image reconstruction is improved when reducing the variance σ². Quantitative comparison of reconstructed images from the CELEBA-HQ test set for DiVeQ optimization technique over different variances. Each curve is average of the metric over all test set reconstructions and over three different individual runs.
On the other hand, we observe comparable quality for the reconstructions when o² (10-2, 10-3 10-1). Based the results, choose 0² 10-3 for the VQ-VAE compression task. Note that we choose the variance value only based on the results obtained from the CELEBA- HQ data set, and we blindly use this value (i.e., o² 10-3 for other data where our proposed techniques outperform other VQ methods (see App. C.I). From a different perspective, we can think of o² hyperparameter that marginally controls the amount of generalization for the VQ weerest-codewer assignments, but does not need proper tuning as long as o² < 10-2, The larger the variance o² is, the more generalization will be injected in nearest-codeword assignments.
0.78
20
0.6
0.7
0.5
SSIM -
PSNR +
24
LPIPS
0.00
0.3
10
22
10
0.6
10
7
8
10
5
10
11
6
7
11
Codebook bits
Codebook bits
Codebook bits
Figure 21: Ablation on variance 2 for SF-DiVeQ: Quality of image reconstruction im- proved when reducing the variance ² Quantitative comparison of reconstructed images from the CELEBA-HO HQ test set for SF-DiVeQ optimization technique over different variances. Each curve the average of the metric over all test reconstructions and over three different individual runs.
29

---

Published as a conference paper at ICLR 2026
In addition, did another ablation study in the VQGAN generation task with the FFHQ, LSUN
Bedroom, and LSUN Church data sets using different variances of = {10 10.00. 10 111--) *}. Ta- ble 10 shows the obtained FID values over different variances for Di and SF-DiVeQ techniques. Following the results in the table, for a specific codebook size, the performance of both Di and Se-DiVee remains almost similar over different variances. This shows that the performance of DiVeQ and SF-DiVeQ is not sensitive to the variance o².
Ablation on variance o² for DiVeQ and SF-DiVeQ: Quality of generations remains quite the same over different variances. The table presents the FID scores for the proposed Di and SF-DiVeQ over different variances for FFHQ, LSUN Bedroom, and Church data sets.
Table 10:
Codebook bits
Data set
Approach
Variance
12
10 2
6.70
5.61
5.11
8.22
DiVeQ (ours)
10-3
7.12
6.22
5.37
8.47
FFHQ
10
7.32
5.86
5.45
8.59
10 2
7.91
6.21
5.93
8.60
Sr-DiVeQ (ours)
10
6,91
7.16
5.72
7.81
10
7.01
6.24
5.69
7.66
10
5.71
4.87
5.40
7.94
DiVeQ (ours)
10-3
5.53
5.39
5.83
8.32
LSUN Bedroom
10
5.60
5.81
5.42
7.69
10 2
5.96
6.01
5.40
7.79
SF-DiVeQ (ours)
10
5.79
5,28
5.21
7.02
10
5.75
5.34
4.91
6.55
10
4.41
3.92
3.90
4.58
DiVeQ (ours)
10 10
4.64
4.36
4.04
5.85
LSUN Church
10-4
4.37
4.14
4.29
6.48
10-2
4.35
3.99
3.56
4.40
SF-DiVeQ (ours)
10
4.32
4.21
3.72
3.96
10
4.46
4.23
3.26
3.75
30

---

Published as a conference paper at ICLR 2026
C.6 ABLATION ON CODEBOOK REPLACEMENT
codebook replacement (App. B.1) for all VQ methods when optimizing the codebook (except for SF-DiVeQ. which does not need codebook replacement). This raises the questions that (i) how much codebook replacement contributes to the overall performance for all different VQ methods, (ii) whether codebook replacement is more in favor of our proposed DiVeQ than the other methods, and (iit) whether codebook replacement actually improves the per- formance when training VQ codebook. In this section, we provide a set of experiments that show that by skipping the codebook replacement, our proposed DiVeQ still obtains higher quality recon- structions and faster convergence than other VQ optimization methods.
In all experiments in the paper,
25
0.6
0.6
20
LPIPS
03
0.5
SSIM
STE
PSNR
EMA
RT
0.4
SE-S
0.1
NSVQ
13
DiveQ (ours)
SF-DDVVO (ours)
0.3
0.3
11
Codebook bits
Codebook bits
Codebook bits
1,500
Perplexity +
500 1,000
Codehook bits
II
10
20 30
40
50
60
70
80
90
100
Training epoch
Figure 22: Ablation on codebook replacement: By skipping codebook replacement, the proposed DiveQ still leads to consistent improvement in image reconstruction quality and achieves the highest perplexity. (top) Quantitative comparison reconstructed images from the AFHQ test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs. (bottom) Perplexity (or average codebook usage) during training for individual with 11- bit codebook.
In the VQ-VAE compression task, we trained all VQ optimization methods by skipping codebook replacement. All models are trained with batch size of 32 learning rate of ir 5.5 10-4 on the AFHQ data set for 100 epochs. Fig. 22 shows the quality of reconstructions along with the perplexity (Eq. (15)) tracked over training epochs for different VQ methods while the codebook replacement is skipped. Note that in the plots, the reported value for each metric is the average of that metric over all test set reconstructions and over three different individual runs. Based on the results in the figure,
if we skip the codebook replacement, proposed DiVeQ achieves higher quality reconstructions than other VQ methods, while the performance gap increases with the increase in codebook size.
Comparison of the results in Fig. 22 (without codebook replacement) with Fig. (with codebook re- placement) reveals that the performance of STE, RT, and especially ST-GS becomes worse without applying codebook replacement, while the performance remains quite the same for EMA. NSVQ, and our proposed DiVeQ In the perplexity plot in Fig. 22, proposed DiVeQ achieves the high-
31

---

Published as a conference paper at ICLR 2026
est perplexity among other VQ methods (except SF-DiVeQ) in the early stages of training, and consistently keeps the perplexity high until the end of training. However, it takes the whole 100 training epochs for the EMA approach to reach comparable perplexity as proposed DiVeQ, which means that DiV converges faster than EMA. Comparing EMA perplexity in Fig. 22 and Fig. 29 shows that codebook replacement helps EMA to converge faster, such that reaches similar perplexity as our DiVeQ at epoch 45. In Fig. 23, we provide the quality of reconstructions over different training epochs for the individual experiment for which the perplexity is shown in Fig. 22 with an 1 -bt codebook. Fig. 23 demonstrates that by initializing VQ with a well-fitted codebook and by skipping the codebook replacement, our proposed DiVeQ can converge faster than other VQ methods.
In Fig. 22, we also plotted the results for the SF-DiVeQ method by copying them directly from Fig. 6 with the aim of having them as baseline for further comparisons. As discussed in Sec. 4, initialize the SF-DiVeQ codebook, the quantization skipped for the first two epochs, and then the codebook is initialized by the average of the latest latent vectors at the end of the second epoch. Note that for the experiments this section (that skip the codebook replacement), we initialize the codebooks of all VQ methods in similar way the SF-DiVeQ approach because of two reasons. First, we can quantitatively compare all VQ methods with SF-DiVeQ when all methods have similar initialization, and none of them use codebook replacement (SF-DiVeQ does need codebook replacement at all). Second, to evaluate how different VQ methods optimize the codebook and how good their gradients are for updating the codebook, a proper experiment to initialize all of them with a codebook that is well-fitted to the current latent distribution. In this way, all codewords should be active when VQ is started, and then afterwards, everything will be dependent on how each technique treats this well-fitted codebook. The metrics and perplexity in Fig. 22 show that our proposed DiVeQ outperforms other methods by introducing well-defined gradients for the codebook and keeping the perplexity (or codebook usage) high throughout the training.
25
0.6
0.6
1
LPIPS
0.5
SSIM
PSNR
EMA
0.4
RT
0.4
GT
NSVQ
15
DiVeQ (ouns)
0.3
0.3
25 35
45 55
85 95
5 15
25
45
55
65
75
95
15 25 35
45
75 85 95
Training epoch
Training epoch
Training epoch
Figure 23: Convergence speed: By skipping codebook replacement, the proposed DiVeQ con- verges faster than other VQ methods by reaching high image reconstruction qualities the early stages of training, Quantitative comparison of reconstructed images from the AFHQ test set for different VQ optimization methods over different training epochs. The curves are for the indi- vidual run shown as the perplexity plot in Fig. 22, with 11 1-bit codebook.
C.7 ABLATION ON SF-DIVEQ INITIALIZATION
In all experiments in the paper. we initialize SF-DiVeQ with custom initialization in which the quantization is skipped for the first two epochs, and then SF-DiVeQ starts quantizing the latent space such that its codebook is initialized with the average of latent vectors obtained from the latest 20-50 training batches. This raises the question that (i) how much this custom initialization contributes to the overall performance of SF-DiVeO (ii) how other VQ methods perform with similar custom initialization (see the answer in Fig. 22), and (iii) whether SF-DiVeQ is still capable of pulling inactive codewords inside the latent distribution using random initialization. In this section, provide a set of experiments showing that SF-DiVeQ with random initialization can reach similar reconstruction quality and perplexity compared to the custom initialization.
32

---

Published as a conference paper at ICLR 2026
0.65
0.38
24
0.36
+ LPIPS
SSIM
0.6
PSNR
0.31
23
0.32
DiVeeicustom
0.33
SF-DI Dww.DieeOrando/110
SF- 30-Dii-Qrmmmmm200
0.3
Codebook bits
Codebook bits
Codebook bits
800
600
Perplexity
200 400
Codchook bits 10
190
Training epoch
Figure 24: Ablation on SF-DiVeQ initialization: By using random initializations, the proposed SF-DiVeQ still performs comparable to custom initialization for small and moderate codebook sizes, while performing worse with the increase the codebook size. (top) Quantitative com- parison of reconstructed images from the AFHQ test set for SF-DiVeQ, such that it uses different initializations (custom and random) and is trained for different epochs (f 1000, 150, 200)). Each curve the average of the metric over all test set reconstructions and over three different individual runs. (bottom) Perplexity (or average codebook usage) during training for one individual run with a 10-bit codebook.
In the VQ-VAE compression task, we train SF-DiVeQ by random initialization for different numbers of training epochs (f100, 150, 200)). All experiments are trained with batch size of 32 at the initial learning rate of [r = 5.5 10 on the AFHQ data set, such that the learning rate is halved after 50% and 75% of training epochs. Fig. 24 shows the quality of reconstructions along with the perplexity (Eq. (15)) tracked over training epochs for different SF-DiVeQ experiments. Note that in the plots, the reported value for each metric the average of that metric over all set reconstructions and over three different individual runs. We also provided the SF-DiVeQ results with custom initialization (captured from Fig. 6) to have them as a baseline for comparisons.
According to Fig. 24, our proposed SF-DiVeQ with random initializations still performs comparable the with custom initialization for small and moderate codebook sizes. However, SF-DiVeQ performance degrades with the increase in codebook size, and the reason stems from the perplexity plot, which shows the average codebook usage for one individual experiment with 10-bit codebook. The perplexity values reveal three important facts; (i) in random initializations, the perplexity is regularly increased during the training, which confirms the fact that SF-DiVeQ consistently pulls inactive codewords inside the latent distribution, (ii) for big codebook sizes, the custom initialization helps SF-DiVeQ to perform better such that SF-DiVeQ starts with high codebook usage from the beginning of training while keeping most (or all of (see Fig. 29)) the codewords active until the end of training. Whereas, this does not happen for the cases random initialization, and (iii) in random initializations, the more SF-DiVed is trained, the higher perplexity it reaches, and thus, the better performance achieves. The reason is that for longer training epochs, SF-DiVeQ has more time to pull more inactive codewords inside the latent distribution.
33

---

Published as a conference paper at ICLR 2026
0.65
0.38
21
0.36
+ LPIPS
SSIM
PSNR
0.6
0.34
SF-DI VeQeeto
SF- DEQi
33
0.32
SE-DiNeQirr coef 30
SF-DINeUHr_oef-44
0
SF-DiNo@fr.owf=33
0.3
Codebook bits
Codebook bits
Codebook bits
1,300
1,000
Perplexity
500
Codehook bits
0
10
20
30
40
50
70
80
90
100
Training epoch
Figure 25: Ablation on SF-DiVeQ initialization: By using learning rates that are multiple times higher than other model parameters, SF-DiVeQ with random initialization can reach similar reconstruction quality and perplexity as SF-DiVeQ with custom initialization. (top) Quantitative comparison of reconstructed images from the AFHQ test set for SF-DiVeQ while its codebook learning rate is Ir coef times higher than other model parameters. Each curve is the average of the metric over all test set reconstructions and over three different individual runs. (bottom) Perplexity (or average codebook usage) during training for one individual run with an 11-bit codebook.
According to the perplexity curves for random initializations in Fig. 24, we notice that the perplexity not increased anymore when SF-DiVeQ training reaches its halfway point. We hypothesize reason could be that the learning rate is halved after 50% and 75% of training epochs, and lower learning rate results in smaller updates in the codebook, and hence, the remaining inactive codewords will not be pulled inside the latent distribution anymore. Therefore, the codebook learning rate SF-DiVeQ with random initialization. Motivated by train SF-DiVeQ with random initializations using different learning rates for the codebook. Fig. 25 shows the quality of reconstructions along with the perplexity tracked over training epochs when the learning rate of the codebook higher than the model learning rate by different coefficients (i.e., cr_coef {10, 20, 30, 40, 50}).. Note that in the plots, the reported value for each metric is the average of that metric over test set reconstructions and over three different individual runs. We also provided the SF-Di SF-DiVeQ results with custom initialization (captured from Fig. 6) to have it as a baseline comparisons, which shown as SF-DiVeQ/custom..
could have a big effect on the performance this hypothesis,
Fig. 25 demonstrates that even by using random initializations, SF-DiVeQ can achieve similar per- formance to the case of custom initialization for all different codebook sizes, if we adopt higher learning rate for the codebook. The perplexity plots in Fig. 25 substantiate our hypothesis that higher learning rate for the codebook encourages SF-DiVeQ to pull the inactive codewords faster inside the latent space, and thus, quickly reach the same (or even higher) perplexity as in custom initialization (i.e., SF-DiVee/custom) Furthermore, the higher the learning rate is, the more per- plexity SF-DiVeQ reaches. Note that for the experiments of Fig. 25, we still apply the learning rate scheduling by halving the learning rate of all parameters after 50% and 75% of training epochs.
34

---

Published as a conference paper at ICLR 2026
C.8 MISALIGNMENT OF CODEBOOK AND LATENT REPRESENTATIONS
As discussed in Sec. when training the VQ-VAE for the compression task using different learning rates and batch sizes, there would be cases in which Shannon's rate-distortion theory (Shannon, 1959) does hold. In other words, in these cases, when the codebook size is increased, the quality of reconstructions will not be improved with respect the objective metrics. In Fig. 26, we provide the quantitative results for the cases where the rate-distortion theory does not hold, and as a result, misalignment between codebook and latent representations occurs (see Fig. 4). The curves represent the results for VQ- VAE trained on CELEBA- HQ data set with various training setups.
0.6
25
0.5
8
LPIPS
20
SSIM
PSNR
+
0.1
0.5
RT
ST-GS
15
0.3
NSVQ
0.4
DiVeQ (ours)
SF-DIVeQ (ours)
5
10 11
11
Codchook hits
Codebook bits
Codehook bits
Figure 26: Misalignment of codebook and latent representations happens where the increase in codebook size does not improve the objective metrics. The figure shows the quantitative results VQ- VAE compression experiments on CELEBA-HQ data set for different VQ methods. Cases where misalignments happen are highlighted in red circles. Each curve shows the results only for one individual run. For the EMA approach, we consider the misalignment that happened in Fig. 17 for the 1 11-bit codebook.
Misalignments are highlighted with red circles in Fig. 26, which refer to the cases of [STE. 10-4, codebook bits=11], [RT Ir 5.5 10-4, codebook bits=6], [ST-GS Ir = 10-3, codebook bits=8], [NSVQ Ir = 10-3, codebook bits=11], and [DiVeQ] 10-3 codebook bits-1l]. Note that for the EMA approach, we con- sider the misalignment that happened in Fig. 17, when the batch size is 128 and (r = 5.5 10-4. According to all of our experiments, for our proposed SF-DiVeQ, the rate-distortion theory always hold true and as a result, no misalignment happens for SF-DiVeQ. For the sake of comparison, we consider the case of [SF-DiVeQ] Note that in difference to other plots in the paper, Fig. 26 shows the results for only one individual experiment, not the average on several different individual runs.
codebook bits=11] for plotting in Fig. 4.
The reported values in Fig. 4 show the distortion per bit (Dper J) for each of the quantization cases. Since the misalignments different VQ methods happen different codebook sizes, it is not fair to report the ordinary distortion for each of the quantization cases. Because smaller codebook sizes result higher amounts distortion. Therefore, for quantitative evaluation of quantization cases in Fig. measure the Dacrba metric that measures how much distortion remains per bit. D Dper-bu is computed as
K Pk log, Pk.
D HCC
Dper-bit
S.t.
D
Zn
2n
12 and H(C)
(17)
where D is the quantization distortion, (C) is entropy of the codebook index distribution, and pk refers to the empirical usage probability of the k--th codebook vector. Here, N and refer to the number of distribution samples and the number of codewords, respectively. Dpec-bu is suitable and fair metric to use to compare the case of misalignments, as it reflects the distortion (D) mixed with the effective codebook usage (H)..
35

---

Published as a conference paper at ICLR 2026
C.9 DIVEQ AND SF-DIVEQ IN RESIDUAL VECTOR QUANTIZATION
As mentioned in Sec. I of the paper. DiVeQ and Se- can also be used to train other variants of vector quantization like Residual VQ (RVQ, Chen al., 2010; Vali & Bäckström, 2023b). Different from ordinary vector quantization, RVQ uses multiple codebooks to quantize a distribution. Suppose an RVQ with three codebooks of {C², ,,,, c*). The first stage of RVQ quantizes the input x to & using the first codebook C
ب = c.- = arg min ||=
c) ll2
s.t.
c) is the j-th codeword of C²,
(18)
and computes the residual of the first stage r - 21. Similarly in the second stage, RVQ takes 11 as input and quantizes it to A using the second codebook C²
ct. arg min oj
||r, cill2 s.t. cj is the j-th codeword of C²
(19)
and computes the residual of the second stage as r2 = ri - r.. Then, for the third (or last) stage, RVQ takes 12 input and quantizes it to P2 using the third codebook Cs
22 = ct. = arg min || - e)'l2 s.t.
cg is the j-th codeword of C3.
(20)
Finally, RVQ computes the quantized input as
2 = c. + ct. + ct.
(21)
where & is the final hard quantized version of the input z.
In this section, we compare the performance of different VQ optimization techniques in the VQ- VAE compression task when quantizing the latent space by RVQ. We train the VQ-VAE models over 100 epochs with a batch size of 32 and an initial learning rate of Ir = 5.5 10 that is halved after 40 and 70 epochs. all experiments, we apply RVQ using three different codebooks over four different codebook bits B = {12, 18, 24, 30}. For instance, for quantizing with 30 bits, we apply RVQ with three quantization stages, each with a 10-bit codebook (or codebook size of 1024).
Fig. 27 and Fig. 28 show the quantitative comparison of different VQ optimization techniques when reconstructing the test set images of CELEBA-HQ and AFHQ, respectively. Note that in the plots, the reported value for each metric is the average of that metric over all test set reconstructions and over three different individual runs. According to the figures, our proposed DiVeQ and SF- Di consistently obtain higher SSIM, PSNR, and lower .PPPS values than the other methods over different codebook sizes (except for the SSIM of NSVQ in Fig. 28). Furthermore, when increasing proposed methods and STE is escalated. On
the codebook size, the performance gap between the other hand, similar to the result for ordinary VQ, the RT performs poorly for low RVQ codebook sizes, and its performance improves with the increase in codebook size.
0.
0,1
26
PSNR +
+ LPIPS
SSIM
34
0.3
STE
22
RT
0,6
NSVQ
DIVeQ (ours)
0.2
SF-DiVeQ (ours)
20
12
24
30
12
24
30
12
18
30
Codebook bits
Codebook bits
Codebook bits
Figure 27: Quantizing the latent space by Residual VQ: The proposed DiVeQ and SF-DiVeQ lead to consistent improvement in image reconstruction quality. Quantitative comparison of reconstructed images from the CELEBA-HH test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test set reconstructions and over three different individual runs.
36

---

Published as a conference paper at ICLR 2026
0.7
24
PSNR -
LPIPS
0.1
SSIM
0.0
22
03
STB
20
0.3
NSVQ
Ve
ST DiVei (ours)
0.4
18
12
24
30
12
18
24
30
12
18
30
Codebook bits
Codebook bits
Codebook bits
Figure 28: Quantizing the latent space by Residual VQ: The proposed DiVeQ and SF-DiVeQ lead to consistent improvement in image reconstruction quality. Quantitative comparison of reconstructed images from the AFHQ test set for different VQ optimization methods over different codebook sizes. Each curve is the average of the metric over all test reconstructions and over three different individual runs.
1. TRAINING LOGS
In this section, study the training logs when training the VQ-VAE for the image compression task on the AFHQ data set. Fig. 29 shows the codebook usage, perplexity (Eq. (15)), and reconstruction and perceptual losses over the course of training for different VQ optimization methods with an 1-bit codebook (or codebook size K = 2048). To make the plots look smoother, the curves are smoothed via Blackman windowing with window size of 15 (except for the codebook usage plot). The codebook usage represents the percentage of codebook used in each training epoch. Apart from the ST-GS method, all other VQ optimization techniques reach the full codebook usage (but happens bit late for the EMA approach). Note that the codebook usage is reported when codebook replacement (App. B. is actively applied during training for all VQ methods. Even though code- book replacement is applied, it cannot compensate for full codebook usage for the ST-GS method.
Perplexity refers to the number of unique codewords selected for quantizing a typical training batch, and it ranges from to codebook size K. Eq. (15) explains how perplexity is computed. According Fig. 29, the highest perplexities belong STE, EMA, and our proposed DiVeQ and SF-DiVeQ, while NSVQ reaches bit lower perplexity compared to these methods. However, for ST-GS and RT approaches, the perplexity is very high at the start of training, and it suddenly drops for ST-GS, whereas it smoothly decays for RT during training. Reconstruction loss is the MSE loss between the input image o and its reconstructed output I,, (see Fig. and Eq. (1)), and perceptual loss the L.IPS between and a,, According the losses, proposed Di DiV and SF-DiVeQ reach lower loss values than STE, EMA, RT, and ST-GS while converging faster than STE, EMA, and ST-GS. In contrast to all methods, NSVQ reaches the lowest loss values with big margin, especially for the perceptual loss. To make the difference between all VQ methods more clear, we plot the loss values on a logarithmic scale.
When training VQ in neural networks, computing the codebook usage once for the whole course of training is not enough make sure everything fine with the codebook optimization. Because the latent representation P= is dynamically changing, and as a result, one used codeword for the initial training iterations might become useless for the rest of the training by falling out of the latent distribution. Codebook replacement helps to keep the codebook usage at its maximum during training by actively reinitializing unused codewords. However, even a 100% codebook usage during the whole training does not guarantee properly learned codebook. Because it is possible that all codewords are located inside the latent representation Pz, but they are not necessarily well-fitted to P.. For instance, in Fig. 29, RT and NSVQ reach 100% codebook usage, but yield poor qualitative (Fig. 5) and quantitative (Fig. 6) results, and they are both prone to misalignment of codebook and latent representations (see Fig. 4).
37

---

Published as a conference paper at ICLR 2026
Codebook usage %
100
EMA
ST-GS
NSVQ
DIVeQ (ours)
SF-DIVeQ (ours)
10
20
30
40
50
60
70
80
90
100
Training epoch
2,000 4,000
Perplexity
0.5
1
1.5
2.5
3
3.5
-10
log (Reconstruction loss)
in6
3
-10*
log (Peroeptual loss)
10
Training iteration
-10*
Figure 29: Training logs do not necessarily reflect whether VQ codebook optimization is being carried out appropriately. Comparison of the codebook usage, perplexity, reconstruction, and perceptual losses for different VQ optimization methods, when training the VQ-VAE compression model with an 11 1-bit codebook on the AFHQ train set.
Similar to codebook usage and perplexity factors, reconstruction and perceptual losses cannot be suitable metric reflect whether VQ optimization method performs better than the other. For instance, in Fig. 29, NSVQ reaches a much higher perplexity than ST-GS and obtains the lowest loss values among the others, but leads to the worst qualitative (Fig. 5) and quantitative (Fig. 6) results over different data sets compared to all other approaches. Therefore, these metrics are useful only make sure the training process works fine.
38

---

Published as a conference paper at ICLR 2026
C.11 MORE SAMPLES GENERATED BY VQGAN
Apart from the quantitative evaluation of VQGAN generations over different VQ optimization meth- ods using the FID metric (Tables 2 and 7), it is important compare the generations qualitatively as well. Figs. 30 to 34 provide completely random generations from VQGAN for all five data sets with 9-bit codebook (or codebook size of K 512).
D DISCLOSURE OF THE USE OF LARGE LANGUAGE MODELS
In this paper. LLMs were used only for minor grammatical edits, word polishing, rephrasing. They did not contribute to research ideation, experiments, or core writing. All suggestions from LLMs were manually verified and edited by the authors prior to final inclusion.
EMA STE
M MEL CF
RT
STACS Cs
SF-DIVEQ
Figure 30: Generation task. Qualitative comparison of random generations for the CELEBA-HQ data set the VQGAN generation task for different VQ optimization methods, with 9-bit code- book (or codebook size of K 512).
39

---

Published as a conference paper at ICLR 2026
EMA
RT
ST-OS
OSSN
SF DIVEO CeR
Figure 31: Generation task. Qualitative comparison of random generations for the LSUN Church data set the VQGAN generation task for different VQ optimization methods, with 9-bit code- book (or codebook size of K 512).
40

---

Published as a conference paper at ICLR 2026
RT
NSVQ
DIVEO
SE-DIVEO
Figure 32: Generation task. Qualitative comparison of random generations for the FFHQ data set in the VQGAN generation task for different VQ optimization methods, with 9-bit codebook (or codebook size of K 512)
41

---

Published as a conference paper at ICLR 2026
ST-GS
DIVRO
SF- DENDD
Figure 33: Generation task. Qualitative comparison of random generations for the LSUN Bedroom data set the VQGAN generation task for different VQ optimization methods, with 9-bit code- book (or codebook size of K 512).
42

---

Published as a conference paper at ICLR 2026
RTT
ST-GS
NSVO
VEQ
Figure 34: Generation task. Qualitative comparison of random generations for the AFHQ data set in the VQGAN generation task for different VQ optimization methods, with 9-bit codebook (or codebook size of K 512)
43


Reviews:
[{"review_id": "Mlm2OSbY0g", "reviewer": "Reviewer_7ZbY", "summary": "This work aims to overcome two problems that previous vector quantization works have had.\nFirst, vector quantization uses a codebook that maps latents to the closest codeword, and the mapping operation is not mathematically differentiable.\nThat is, it suffers from a gradient collapse problem, which means the gradients will not pass through after the codebook in the backward pass.\nSecond, while training, several codewords remain inactive, and thus, they are not trained well.\n\nTo resolve the above problems, the paper proposes the following.\nFor the gradient collapse problem, the paper uses the reparameterization trick to align the direction of quantization error with the nearest codeword.\nBy representing a quantized vector with a latent vector and a directional noise and showing that this quantized vector is differentiable to the latent vector and the codeword, the paper defines vector quantization as a differentiable problem.\n\nFor inactive codewords, codebook replacement is widely used to replace inactive codewords with new codewords that are near active codewords.\nThis paper points out that codebook replacement is a heuristic method, and suggests that mapping latent to an interpolation of subsequent codewords, instead of mapping to a single codebook.\n\nThe above proposals yield high-quality images while achieving good indicators in both image reconstruction and image generation.", "questions": "The paper proposes SF-DiVeQ to eliminate the need for codebook replacement. Then, while training VQ with DiVeQ, are inactive codewords replaced by codebook replacement techniques?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The gradient collapse problem of vector quantization is a problem that many researchers have been trying to solve by detour for a long time, and the paper proposes to solve this problem directly.\n- The paper summarizes well about prior works that try to solve the same problem and their problems.\n- The paper explains the problem situation and the proposed method well, through mathematical formulas.\n- The proposed method can help generate high-quality images in both image reconstruction and image generation tasks, while achieving better indicators such as FID and LPIPS than other prior works.", "weaknesses": "- The paper proposes SF-DiVeQ to eliminate the need for codebook replacement. However, the reviewer wonders if SF-DiVeQ can really supplant codebook replacement, in case both subsequent codewords are inactive.\n- Unlike other prior works, the FID score of the proposed method becomes worse when the VQ bitrate becomes larger. It may imply that the proposed methods are not general.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JxeeAXupuD", "reviewer": "Reviewer_3NP6", "summary": "Quantization is a mathematically non-differentiable process, which poses a major challenge for the end-to-end training of neural networks. To overcome this, this paper proposed a differentiable vector quantization (DiVeQ) using the reparameterization trick. The proposed DiVeQ ensures the hard assignment in the forward pass while allowing gradient flow to enable end-to-end training. Its variant, space-filling DiVeQ (SF-DiVeQ), alleviates codebook underutilization and achieves full utilization. The proposed methods excel the existing quantization method in both investigated image compression and generation tasks and can be easily used as a drop-in replacement for the standard vector quantization method.", "questions": "1. Could the authors provide a theoretical analysis regarding the convergence properties or gradient flow of the proposed methods? \n\n2. The proposed SF-DiVeQ requires a specific initialization strategy. It would provide more insights if the authors could show the necessity of the pre-training step and the impact of the pre-training step on stabilizing the latent representation and improving the performance. Could the authors quantify the performance degradation when SF-DiVeQ uses the randomly initialized codebook? Is there a big performance drop or is it only marginal? Could the author also provide the extra complexity required by the initialization strategy?\n\n3. For DiVeQ, it is straightforward to the nearest codeword during training. But in SF-DiVeQ, a dithered codebook sampled from two subsequent codewords is used to quantize the input. How do the authors select these two subsequent codewords during training?\n\n4. Two different variance values were used for the investigated tasks. Could the authors explain how and why they chose different values for each task? Does this mean that a variance tuning step is required when applying to other domains, such as speech coding?\n\n5. Could the author give more discussion regarding the strength of the proposed method as a well-tuned EMA/STE with the codebook replacement algorithm can also achieve comparable performance? What do we benefit from the proposed method? Or does the codebook replacement algorithm contribute more to the performance?\n\n6. As mentioned in the weakness part, please change the term used to represent the codebook size. Codebook size is not the same as the bitrate.\n\n7. For the ablation on variance values of the proposed method, the colors in Figure 20 are not easy to distinguish, especially for variances with small values. Could the author adjust it using different line types?\n\n8. A few typo errors are found. \n   For example, on page 9, 'but they are not homogeneously scattered In ST-GS,' probably should be 'but they are not homogeneously scattered. In ST-GS'. \n   On page 16, 'we mask half the the true token indices forcing the'. Should be 'we mask half of the true token indices forcing the'.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed DiVeQ addresses the gradient collapse problem by modeling quantization as the addition of a simulated quantization error to the input. A key improvement over previous noise-substitution vector quantization (NSVQ) is that the direction of the error vector is aligned with the selected codeword to guarantee a good reconstruction.\n2. This paper also provides insights into existing quantization methods and establishes a basic understanding of the limitations. \n3. The proposed SF-DiVeQ addresses quantization error and codebook underutilization by extending quantization from discrete codewords in the codebook to continuous line segments connecting neighboring codewords. SF-DiVeQ achieves full utilization and improved representational efficiency. \n4. A codebook replacement algorithm was also given to enable stable codebook utilization and it can be applied to all investigated existing quantization methods.\n5. The proposed methods do not require a specifically designed auxiliary loss or hyperparameter tuning compared to existing methods.\n6. The proposed methods demonstrate consistent improvement over existing methods across multiple datasets on image compression and generation tasks. Comprehensive ablations on codebook size, batch size, learning rate, and the hyperparameter of the proposed methods are conducted.\n7. The proposed methods can act as a drop-in replacement with minimal changes and the authors have investigated the performance within a residual vector quantization design.", "weaknesses": "1. This paper analyzed the limitations of NSVQ and its poor performance due to the high randomness. For the proposed DiVeQ, the paper highlighted the fundamental theory and the definition of the simulated error and how it improves quantization accuracy by adjusting the hyperparameter. But a more detailed theoretical analysis of convergence properties and gradient flow would be appreciated to provide more insights into the proposed method.\n2. The proposed method approximates the quantization process by generating a directional noise vector and one component is sampled from a zero-mean Gaussian distribution. Adjusting the variance can affect the quantization accuracy (see Figure 3). Although the paper states that the variance parameter is not highly sensitive and a small value should be used, two different variance setups were used in the investigated compression and generation tasks. One can assume that some tuning or empirical selection is needed, especially in different domains or tasks. This is the same compared to existing methods, where hyperparameter tuning is needed.\n3. The proposed SF-DiVeQ requires a specific and empirical initialization strategy. This introduces extra complexity and might increase implementation overhead if the codebook size is large. The results of SF-DiVeQ without this specific codebook initialization strategy are not discussed.\n4. Similar to item #3, the recommended training setup for training SF-DiVeQ is to skip quantization for several initial epochs to stabilize the latent representation and then initialize the codebook vectors using the average of recent latent vectors. This process requires careful tracking of recent latent vectors and defining a specific number of training iterations to capture them, making it less of a drop-in replacement compared to existing methods that allow simple random initialization.\n5. The proposed methods are only evaluated on image compression and generation tasks. Further investigation into speech/audio coding would be appreciated to confirm a broader applicability of the proposed methods.\n6. The codebook replacement algorithm is used for all investigated methods. The proposed methods demonstrated good performance compared to existing exponential moving averages (EMA) and straight-through estimator (STE) methods by a smaller margin or performed comparably to EMA. This suggests that a well-tuned EMA can also deliver high-fidelity reconstruction. So, maybe, the codebook replacement algorithm plays a more important role in improving the overall performance while the contribution of the proposed DiVeQ is secondary.\n7. Throughout the whole paper, the term 'VQ bitrate' is used to represent the codebook size. But bitrate refers to the number of bits used to represent the given bitstream per unit of time. For the image compression task, bitrate is measured in bits per pixel. In the paper, it is a misuse of the term 'bitrate' and it should be 'n-bit codebook' instead of 'VQ bitrate of n'.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ptG58rK9Xe", "reviewer": "Reviewer_XAfN", "summary": "The paper aims to improve vector quantization (VQ) in VQ-VAE compression and VQGAN generation. The paper first analyzes the baseline methods and compares them to identify gaps. Then, the paper proposes DiVeQ, which applies reparameterization tricks to the baseline NSVQ. The paper then proposes an improved version of DiVeQ by replacing the codebook replacement method with codebook interpolation, as suggested by another baseline, SFVQ. The paper conducts experiments on VQ-GAN image generation tasks at 256x256 resolution, and the experimental results show that the proposed method achieves better SSIM, PSNR, FID, and LPIPS scores.", "questions": "Please refer to the Weaknesses section. In addition, these are minor questions and points.\n* Why are the proposed variants of DiVeQ and SF-DiVeQ in the supplementary materials, while they show better results?\n* The baselines could be outdated.\n* t-SNE plots are known to present misleading information, especially when compared to other t-SNE plots.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "* The paper provides an excellent background and related work section and presents a detailed comparison among them.\n* The figures are helpful for understanding the idea.\n* The experimental results show that DiVeQ achieves better results compared to the baselines.", "weaknesses": "* The contribution of the paper is not strong.\n\nThis reviewer believes that the proposed methods are minor extensions of the baselines, and the motivation and experimental results are not insightful enough. The paper proposes two methods, DiVeQ and SF-DiVeQ. DiVeQ is proposed by adding a direction vector to $v \\sim N(0,\\sigma^2I)$, and SF-DiVeQ is proposed by adding space-filling methods from SFVQ. In addition, although the paper’s title primarily emphasizes the reparameterization trick, the proposed methods are built on an existing baseline that already uses it. Also, since the motivation and experimental results do not present insightful analysis, this reviewer believes that the contribution of the paper is weak.\n* The experiments section should be strengthened.\n\nThe paper primarily conducts experiments on VQ-GAN, generating 256x256-resolution images. From the experiments, the paper compares image-related metrics such as FID and PSNR. However, this reviewer believes that the experimental results section could be improved by using different models and tasks and by providing more detailed analysis beyond the numerical results. For example, many papers conducted experiments on various models and tasks such as recent stable diffusion models [C1], visual-language-action tasks [C2], and Gaussian splatting [C3]. As these papers are cited in the DiVeQ paper, the reviewer encourages broadening the experiments, as the proposed methods are drop-in replacements for standard VQ layers, as mentioned in the paper.\n\n[C1] Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang, Yinpeng Chen, Lu Yuan, and Gang Hua. Designing a better asymmetric VQGAN for stable diffusion. arXiv preprint arXiv:2306.04632, 2023.\n\n[C2] Seonghyeon Ye, Joel Jang, Byeongguk Jeon, Sejune Joo, Jianwei Yang, Baolin Peng, Ajay Mandlekar, Reuben Tan, Yu-Wei Chao, Yuchen Lin, Lars Liden, Kimin Lee, Jianfeng Gao, Luke Zettlemoyer, Dieter Fox, and Minjoon Seo. Latent action pretraining from videos. In International Conference on Learning Representations (ICLR), 2025.\n\n[C3] Haishan Wang, Mohammad Hassan Vali, and Arno Solin. Compressing 3D Gaussian splatting by noise-substituted vector quantization. In Scandinavian Conference on Image Analysis, pp. 338–352. Springer, 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YI1EulZrXF", "reviewer": "Reviewer_Q4CB", "summary": "This paper addresses the non-differentiability of vector quantization (VQ) by proposing Differentiable Vector Quantization (DiVeQ) and a variant, Space-Filling DiVeQ (SF-DiVeQ). The methods enable gradient flow while preserving hard assignments, with a directional reparameterization of the quantization error. Both methods train end-to-end without auxiliary losses or hyperparameter tuning. Key results on VQ-VAE compression and VQGAN generation tasks demonstrate that DiVeQ and SF-DiVeQ consistently outperform existing approaches in reconstruction and sample quality across multiple datasets.", "questions": "1. In Section 3.2, SF-DiVeQ is described as quantizing inputs to \"random locations on the line connecting two subsequent codewords.\" Could the authors clarify the ordering of codewords along the space-filling curve?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. DiVeQ introduces a novel directional reparameterization of the quantization error. The proposed techniques are elegantly derived from the reparameterization trick, and allow for true end-to-end training without auxiliary losses or complex schedules.\n2. The authors perform a comprehensive evaluation across both image compression and image generation. The experiments are conducted on a wide range of standard, high-resolution datasets, and the proposed methods consistently outperform a diverse set of baselines.\n3. The paper's claims are further confirmed by ablation studies, which analyze the impact of key factors like the noise variance, batch size, and learning rate.", "weaknesses": "1. In Page 4 and Figure 2, the NSVQ incurs a higher distortion than the true nearest-codeword assignment with a fixed probability of approximately 0.67. This claim can be extended beyond a 2D plane, where the probability of incurring higher quantization errors rapidly approaches one as the dimension $D$ increases. This can be mentioned to strengthen the argument.\n2. In Section 4, the authors claim SF-DiVeQ avoids codebook misalignment (Fig. 4). However, the t-SNE plots show only a single run per method without providing statistical evidence (e.g., multiple runs with variance across seeds) to confirm that SF-DiVeQ consistently avoids misalignment. The authors are suggested to quantify the improvement across runs.\n3. The arguments that DiVeQ requires no hyperparameter tuning seem to be overstated. DiVeQ involves the hyperparameter $\\sigma^2$, and different values are used for compression ($10^{-3}$) and generation ($10^{-2}$).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)