Discrete Variational Autoencoding via Policy Search
Authors: Michael Drolet
Firas Al-Hafez
Aditya Bhatt
Jan Peters
Oleg Arenz
Venue: ICLR 2026 Poster (2026)
Keywords: Discrete Variational Autoencoding
Reinforcement Learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We present a method for learning stochastic discrete latent encodings using a novel policy search formulation.

Abstract: Discrete latent bottlenecks in variational autoencoders (VAEs) offer high bit efficiency and can be modeled with autoregressive discrete distributions, enabling parameter-efficient multimodal search with transformers. However, discrete random variables do not allow for exact differentiable parameterization; therefore, discrete VAEs typically rely on approximations, such as Gumbel-Softmax reparameterization or straight-through gradient estimates, or employ high-variance gradient-free methods such as REINFORCE that have had limited success on high-dimensional tasks such as image reconstruction. Inspired by popular techniques in policy search, we propose a training framework for discrete VAEs that leverages the natural gradient of a non-parametric encoder to update the parametric encoder without requiring reparameterization. Our method, combined with automatic step size adaptation and a transformer-based encoder, scales to challenging datasets such as ImageNet and outperforms both approximate reparameterization methods and quantization-based discrete autoencoders in reconstructing high-dimensional data from compact latent spaces.

Published as a conference paper at ICLR 2026
DISCRETE VARIATIONAL AUTOENCODING VIA POLICY SEARCH
Michael Drolet'; Firas Al-Hafez', Aditya Bhatt Jan Peters 12,334 & Oleg Arenz
1 Technical University of Darmstadt German Research Center for AI (DFKI)
3 Hessian.Al Robotics Institute Germany (RIG) Darmstadt, Hesse, Germany
ABSTRACT
Discrete latent bottlenecks in variational autoencoders (VAEs) offer high bit ef- ficiency and can be modeled with autoregressive discrete distributions, enabling parameter-efficiet multimodal search with transformers. However, discrete ran- dom variables do not allow for exact differentiable parameterization; therefore, discrete VA typically rely on approximations, such as Gumbel-Softmax repa- rameterization or straight-through gradient estimates, employ high-variance gradient-free methods such as REINFORCE that have had limited success on nigh-dimensional tasks such as image reconstruction. Inspired by popular tech- niques in policy search, we propose training framework for discrete VAEs that leverages the natural gradient of a non-parametric encoder to update the paramet- ric encoder without requiring reparameterization. Our method, combined with encoder, scales to chal-
automatic step size adaptation and a transformer- lenging datasets such as ImageNet and outperforms both approximate reparame- terization methods and quantization- -ased discrete autoencoders in reconstructing high-dimensional data from compact latent spaces.
INTRODUCTION
Discrete representations play critical role in numerous fields, including telecommunications, biol- ogy, and robotics. While real-world data is often continuous, introducing a discrete latent bottleneck can offer unique advantages, such increased bit-efficiency, multimodal modeling, and the ability to leverage tools from combinatorial optimization. These properties make discrete representations particularly useful for tasks that benefit from structured or compact latent spaces. However, encod- ing continuous information into discrete representations and decoding it back remains a challenging endeavor. In particular, gradient-based optimization requires a differentiable transformation to map continuous inputs to discrete latents, but this is impeded by the non-differemiiabiiityy of operations like rounding or taking the argmax, which are essential for obtaining hard assignments. Because of this difficulty, we often left with approximate techniques, such as Gumbel Softmax reparameter- ization and straight-through esimation Jang et al.) (2016).
However, approximate reparameterization with Gumbel-Softmax Dang al. (2016) is sensitive to its temperature hyperparameter, forcing a trade-off between increased gradient variance at low tem- peratures and large approximation error at high temperatures. Furthermore, for larger bottlenecks, backpropagating the soft-assignment through the autoregressive sampling process has high memory footprint and can suffer from vanishing gradients. Variance-reduucioon schemes such as the Gumbel- Rao Monte Carlo estimator (GR-M RCCK) Paulus et al.) (2020) help but do not fully resolve the scaling problem. On the other hand, rector-quantizationn methods, such as the Vector-Quantize Variational Autoencoder (VQ-VAE) Van Den Oord et al. (2017) and Finite Scalar Quantization (FSQ) Mentzer (et al. (002)) avoid the problems of discrete reparametcrization, but still rely on approximations due to the use of straight-through estimation. Furthermore, their latent distribution is intractable, which prevents them from maximizing the latent entropy in the Evidence Lower Bound (ELBO) ssore-function-base based gradient estimators, such RE- BAR (Tucker et al., 2017) or Muprop (Gu et al., 2015), are promising alternative as they can obtain unbiased gradients optimize the ELBO in the discrete setting. However, due their high gradient variance, the success of these methods has so far been limited.
and demands special training objectives.
Miihhall Drolet (m chael .drlleree et @t@tuuuuueee -tuudarree rmstadt. .d) is the corresponding author. Code and supplementary material available at drolet. 10/ daps

---

Published as a conference paper at ICLR 2026
To overcome these limitations, propose Discrete Autoencoding via Policy Search (DAPS), training framework that optimizes the ELBO for discrete, autoregressive encoders without any dis- crete reparameterization straight-through tricks. DAPS casts encoder learning KL-reuularized policy-search problem: we form a closed-form, nonparametric target distribution q" and update a parametric encoder go by weighted maximum likelihood, avoiding backpropagation through the sampling path. A single scalar trust-region parameter is adapted automatically using an effective sample size (ESS) objective, yielding stable step sizes across tasks and loss scales. We additionally include robotics experiment to demonstrate that DAPS can solve complex, high-dimensional con- trol tasks beyond the pixel-based domain. In particular, we show that the latent space learned by DAPS can be effectively leveraged to train whole-body expressive motions in robotics simulator, conditioned on a behavior context. Empirically, DAPS trains on high-dimensional datasets (e.g., ImageNet) with superior reconstruction quality, while offering (i) explicit entropy/bit-rate control via the entropy regularization parameter 3, (ii) stochastic discrete latents amenable to downstream search, and (ii) stable training behavior.
2 PROBLEM SETTING AND RELATED WORK
We seek to maximize the Evidence Lower Bound (ELBO) of the data likelihood p(x), introduced in the VAE (Kingma& & Welling 2013 Rezende et al. 2014), but using a discrete latent variable. As in the continuous case, the ELBO given by the likelihood of the reconstruction minus the KL
divergence between the approximate latent posterior and latent prior:
CELBO 0..
Ex log (x|2)) - BKx.. (2(()) p(2)).
(1)
where the generative model (x/z) and the recognition model go (z|x) are neural networks parame- terized by ø and 0, respectively. The expectations are computed by first sampling x from the dataset X and then sampling z from the recognition model (o (z) Higgins et al. (2017) introduced the coefficient 3 as a hyperparameter, generalizing the original VAE which can be recovered for / = 1.
In the continuous setting, estimate the gradient of the expectation with respect to for gradient-based optimization. How- ever, we consider discrete latent embeddings, which we organize as sequence of B blocks, z = (z....., BB) with V possible classes per block (also known as the vocab/codebook size). Furthermore, while our method can be applied general architectures, we explicitly focus on au- toregressive encoders
choose a Gaussian recognition model and reparameterize to
= II gl 22,............... 1,
due to the higher expressiveness compared to approximating the joint distribution as product of marginals. However, by using discrete latents, exact differentiable reparameterization is not possible, and approximate backpropagation through the autoregressive sampling (backpropagation through time) is challenging because of high gradient sensitivity. We will now discuss related work, outlining existing solutions for discrete latent variables and comparing their applicability our problem setting.
Approximate Discrete Reparameterization. Methods based approximate reparameterization typically apply the Gumbel-Softmax trick, an approximation of the exact but non-differentiable Gumbel-Max reparameterization, that was originally proposed by [aag al. (2016) and applied to the binary MNIST dataset. Given x ~ p(x), let e represent the unnormalized class probabilities such that go (z x) softmax (go (f)x)). Reparameterizing the latent variable can be performed by first where u ~ Uniform(O, 1) and is arbitrarily close to zero. Whereas the Gumbel-Max trick would exploit that 2hard = arg max, E + 9; follows the desired distribution, Phand <(2/x), the Gumbel-Softmax computes Z softmax ((e g)/T). The temperature т is hyperparameter that needs to be carefully chosen to trade off the accuracy of the approximation and the variance of the gradient. In order to pass hard assigned labels to the VAE decoder, Jang et al. (016) combine the Gumbel-Softman reparameterizztion with a straight-through estimator, by using the exact but non-differentiable Gumbel-Man reparameterization during the for- ward pass, and the approximate but differentiable Humbel-Softman reparameterization during the backward pass, which can be straightforwardly implemented Z + sg( (haard Zwa).. ). where
sampling Gumbel noise: = - log( - log(u+e))
2

---

Published as a conference paper at ICLR 2026
hhrd refers to the one-hot encoding of 2hard and sg prevents the gradient from flowing through its argument. Although the Sumbel-Softmaal trick was successfully used by the discrete VAE (dVAE) in DALL-E (Ramesh et al., 202I) to generate high-quality images, its sensitivity to the temper- ature parameter is a major hurdle when applying it in practice. Furthermore, backpropagating through the soft-assignments increases the memory footprint, and for autoregressive models the backpropagation- theeggattoo-toouuh--iie through-time can suffer from vanishing or exploding gradients. We will compare to GumbeJ-Softmax in our experiments, where we also consider GR-MCK (Paulus al. 2020), modification that uses Rao- Boo-Bacckellizattoon for variance reduction.
Vector Quantization. VAEs based on vector quantization are arguably the most popular models for learning autoencoders with a discrete latent bottleneck. The YQ-VAE encoder first outputs continuous latent vectors and then maps them to their nearest (in terms of Euclidean distance) vector parameterized embedding table Exp. This quantization discretizes the bottleneck but introduces a non-differentiable argument minimization in the computational graph. To update the encoder with respect to the reconstruction loss, VQ-VAE: employ straight-throw estimation, in a similar fashion to the Gumbel-Softmax estimator. Namely, it uses the non non-differentiable argument minimization during the forward pass, but skips the quantization step during the backward pass, which can be implemented using stop-gradien, operator 2g == + sg(2z z), where and Zy are the encoder outputs before and after quantization, respectively. However, we cannot compute the ELBO in Eq. T since we cannot obtain a tractable discrete distribution over the latent variable. Instead, the VQ-VAE optimizes different loss, given by
C = Ex log (x((g) + [|sg(z) evp (z)l + B11z z sg(ev(z))13
(2)
where E (2) is the nearest embedding vector and z, == z + sg(e (z) 2) uses straight-through estimation. The expectation with respect to z go (2)) equates to a single sample, i.e., the output of the deterministic recognition model. A more recent method, FSQ Mentzer et al. (202)). replaces the learned codebook with rounding mechanism. This design choice results in increased parameter efficiency compared to VQ- VAE (due to the small number of levels typically used), while also avoid- ing the tricks needed to enforce high codebook usage. Similar the VAF, foeus generative methods, we explicitly
modeling from compact latent spaces, but in contrast to quantization- regularize for entropy, as staying close to the prior distribution requires more direct control of the encoder entropy. Finally, it is worth noting that the use of heterogeneous FSQ levels can make FSQ more challenging to implement in downstream tasks compared other discrete VAEs.
Gradient-Free Optimization. Another branch of methods circumvents the challenges of back- propagating through the discrete sampling process by relying on zero-order methods. Episode- based policy search (Deisenroth et al., 2013) a subfield of reinforcement learning which aims to optimize policy To (a)s,, which is probability distribution over actions given the state to maximize the expected return R. which is the cumulative reward received over an episode. REIN- FORCE [Wiliiams] 1992) is method that can be used to obtain an unbiased estimate of the gradient of the expected return with respect to the policy parameters 9, without requiring backpropagation through an expectation over the policy. Due to the close relation between reinforcement learning and variational inference (Neumann] 2011 Arenz et al., [0188] Levine, 2018), this approach can be straightforwardly applied optimize the ELBO [ where the recognition model go(zix x) takes the role of the policy, and the log-likelihood
Pa Px the role of the return (Tucker et al., 2017]] Mnih & Gregor 2014, Gu al. 2013 Grathwohl et al. 2017). As REINFORCE sulters from high variance, these methods typically employ variance reduction techniques, such as introducing con- trol variates (Mnih Gregor 2014) or combinations of gradient estimators (Tucker et al. 2017
Grathwohl et al. (017).
However, despite these efforts, the gradient variance of these zero-order methods is still large, re- sulting in limited success compared to approximate reparameterization or vector quantization. Yet, there have been significant advances the field of policy search since the introduction of REIN-
FORCE (Williams
1992) that have been overlooked in the field of discrete VAEs. In particular.
trust-regions based on the Kullback-Leible divergence to the previous policy (Peters et al. 2010 Schulman 2015) and zero-order natural gradient estimates (Kakade) 2001| Peters & Schaal 2008
Wierstra et al. 20144) are standard techniques in reinforcement learninn that also seem promising for training discrete VAEs. REPS (Peters al. 2010) combines both ideas by iteratively solving the
3

---

Published as a conference paper at ICLR 2026
optimization problem of maximizing the expected return of non- non-parameteric policy subject to trust-region constraint to the previous policy. REPS has been adapted to many problems, including hierarchical control (Daniel et al. 2016), model-based RL (Abdolmaleki et 2015), deep RL (Ab-) dolmaleki et al. 2018), and variational inference (Arenz et al. 2018). Methods that are similar in nature ours include On-Policy Maximum Posteriori Optimization (V-MPO) (Song et al.) 2019) and Supervised Policy Update (SPU) (Vuong et al.) 2018). As in REPS, these methods construct a nonparametric target policy by solving a constrained optimization problem and then update the pa- rameterized policy by minimizing its KL divergence from the target policy, resulting in a weighted maximum likelihood objective. Similar to V-MPO and SPU, we use neural network based policy. We also incorporate similar techniques from Lower Bound Policy Search (LBPS) (Watson & Peters 2023
which optimizes the effective sample size (ESS) to adapt the size the trust region.
DISCRETE AUTOENCODING VIA POLICY SEARCH (DAPS)
We will now present our method for optimizing the ELBO (Eq. 1 with respect to expressive discrete encoders, such as transformers, that can capture the correlations between different latent dimensions through auto-regressive sampling. We avoid computationally expensive and high-variance back- propagation through time by using insights from reinforcement learning that enable us to optimize the encoder using weighted maximum likelihood. In the following, will present the encoder update and the decoder update separately, although they are updated together during training.
Introducing the return R(z, x) log p (x|z) 3 log p(z) we can that the ELBO (Eq. can be treated as an entropy regularized return maximization problem,
(44 =/ p(x) Σ go (o(z))) dx + (H(((x)
(3)
which is closely related to maximum entropy reinforcement learning (Levine 2018). However, stan- dard reinforcement learning typically considers sequential decision making, where agent observes the current state x, chooses an action according to its policy (z(x)), which causes transition the next state and yields a step-based reward, and aims learn the policy that maximizes the accumulated reward. As we are not considering such sequential decision making in our setting, our work is more closely related episode- based policy search (Deisenroth et al.) 2013). a subfield of reinforcement learning that directly aims to maximize the policy with respect to the return, without considering the step-based interactions.
Building on this connection, we can apply techniques from reinforcement learning and policy search to update the encoder of our discrete autoencoder. Namely, we estimate the state value V(xx). which corresponds to the expected return and entropy for state x, and use it for variance reduction. Fur- thermore, we split the encoder updates into two steps, where we first perform a closed-form update on non-parametrie encoder qz) and then realize this update the parametric encoder qe (z x) using maximum likelihood. While these techniques were introduced in reinforcement learning, and sometimes RL terminology to highlight the connections, the following description of our method is self-contained and does not assume prior knowledge in reinforcement learning.
Baselines for Variance Reduction. When using Eq. 3 optimize the ELBO with respect to the encoder, subtracting a state-dependent baseline b(x) from the return R(z, x) does not affect the optimal solution, and can reduce the variance of Monte-Carlo estimates without introducing bias (Greensmith et al., 2004 Deisenroth et al. 2013). A typical baseline in reinforcement learning is to use the Value function Vxx). which corresponds to the expected entropy-regularize return. However, we want to estimate a suitable baseline using a small number & of latent samples, and, therefore, estimating the expected return would be highly sensitive to bad outliers. Instead, we use the soft-maximum of the returns, which is an optimistic estimate of the Value function, and accurate under the assumption that the samples were drawn from the optimal encoder (Ziebart, 2010). Hence, replace the return with the estimated advantage
A(z, x) R(2, x) log
exp R(z x), x² ()(x))).
Optimizing the Nonparametric Encoder. We will start by defining a constrained optimization problem for the objective given by Eq. 3 The procedure is based on REPS, where the objective is
4

---

Published as a conference paper at ICLR 2026
Algorithm DAPS: Weight Computation and Loss Functions
Procedure GETWEIGHTS 0, Φ, x)
Procedure UPDATEFN Φ, x)
I: for k 1 to K do
1: {w" ) fz1) GETWEIGHTS(0..O .,
z² t sample from (o [x)
2: // Decoder update
3: R +- log (x +
3: [(Φ) t K Ex log (x
4: xz,,)) t R log Ex exp(R*)
4: // Encoder update
5: log q (z|x) x
A A x) + 77 log go (z/x))
5: [(0) Σx w" log (z x)
1 B
// update for !
6: w + (xx)
Ex 3"('""
xx)
90(et
7: L(n) ( -0) ) ESSueget
7: return stopGrad( (u*)) (2"1))
return ((0), ((4). C(n)
to maximize the expected return (or advantage) while satisfying a constraint on the KL divergence between the policy and the most recent parametric policy go (z)xx,,
max [ p(xx) Σ (z) x) dr (H(((2)x))
(4)
DKL(9(Z)x) l1 (o(2(x)) Sn
[ =1.
We can use Lagrangian multipliers to convert the constrained optimization problem into an uncon- strained one. As shown in Appendix A.I the optimal solution is given by
AA(z, x) + + log qo (z(xx)
(z)x) exp
(5)
77 + 3
where 1 is Lagrangian multiplier: The term 1) controls the trust region, i.e., how large of a step we take toward the optimal policy, while 33 controls the policy entropy. While Eq. providds us with a solution for the optimal policy, the normalizer, since summing over all actions is intractable when using a high-dimensional latent bottleneck. However, will now show that update the parametric policy using weighted maximum likelihood.
only evaluate it given particles, and we cannot compute
use Eq. 5 to evaluate importance weights to
Updating the Parametric Encoder. Now that we have the form of the optimal nonparametric policy, we can move our parameterized policy toward it using maximum likelihood, which enables to train it without requiring backpropagation through the autoregressive sampling. However, because we cannot sample from our nonparametric distribution q", we have to resort to importance sampling, where we use the current parametric encoder as the proposal distribution. Hence, we maximize weighted maximum likelihood objective,
((0) / p(x) DKL (q" (zix) [z((x)) dx ΣΣ wik log go const (6)
where wi q² (x)s are the importance weights, and the constant term is the entropy of 9" which does not affect optimization. However, due to the unknown normalizer of q", we can only evaluate the importance weights up to constant factor, and therefore resort to nelf-normalized importance weighting using weights w; = we/ E, Wy. nolmanormalized importance weights benefit from a lower variance but introduce a bias to the approximation, which diminishes asymptotically for large sample sizes. The new objective can then be summarized maximizing the log- ligelikeli of our parameterized policy, weighted by q² The generative model is fixed during the policy update, and the process for updating it will now be described in more detail.
Decoder Update. The generative model and the recognition model (i.e., the policy) are updated using coordinate descent, which separates the optimization of the decoder from the optimization of the encoder. The generative model p& (x (z) takes in a latent code z and outputs the parameters of distribution over When optimizing the ELBO with respect Φ. the objective reduces to maximizing the expected log-likelihood under the current recognition model:
[(Φ) = - p(x)) Σ
go (z|x)) xx) log P (x|2) dx Σ Exrgp(aix.) P 2))
x
5

---

Published as a conference paper at ICLR 2026
The approximation follows from Monte Carlo integration using randomly sampled data points and latent samples. Note that this objective corresponds to the standard loss of the VAE in Eq. How- ever, because we do not reparameterize the latent variable, the update is performed independently from the policy update, in contrast to VAEs, where both the generative model and recognition model updated in one backward pass. Hence, our generative model update does not depend 0, and the KL divergence to the latent prior can be dropped.
Step Size Adaptation using Effective Sample Size. The multiplier 1) in Eq. controlss the step size of our update based on the relative entropy between our nonparametric target o² and qo. To adapt 1 automatically, we follow prior work using the effective sample size (ESS) (Maia Polo & Vicente 2023 Metelli et al., 2220; Watson & Peters (202)). which provides a tractable proxy for the order-2 Rényi divergence. We update 1) so that the minibatch ESS matches a desired target level urgu
N
K k-
witk
Wek (= m.:n)
ESS,,
(8)
N
9e ( x.1)
;=1 P
Further details and the connection to Rényi divergences are provided Appendix A.2 In practice,
7 is treated as a trainable parameter and updated with SGD on (ESS,, - pngg)) 2 We found
ESS E [K/4, 3K/4] to yield stable convergence. During training, we observe that 1) decays smoothly over time, resulting in smoothly decreasing step sizes during optimization, while account- ing for the task-specific scale of the reconstruction loss.
4 EXPERIMENTS
Datasets. We evaluate on four domains of increasing scale and complexity. We will use the term block size to denote the number of indices in the latent code (i.e., the maximum length of the trans- former decoder), and the term vocab size to denote the number of possible values each index can take (i.e., the number of embeddings in the VQVV codebook). The datasets are as follows. MNIST: 28 x 28 (binary) with 64-bit bottleneck (vocab size 256; block size 8). CIFAR-10: 32 x 32 with a 576-bit bottleneck (vocab size 512; block size 64). ImageNet-256: 256 x 256 with 10,240-bit bottleneck (vocab size 1,024; block size 1,024). LAFAN: motion dataset (~4.6h expressive hu- man motion, 5 subjects) with a 640-bit bottleneck (vocab size 1,024; block size 64). For LAFAN, motions are retargeted to the Unitree HI robot and encoded as 32 x 121 tensors of relative poses. Across datasets, we enforce the same bottleneck size for all methods to ensure fair comparison.
Architectures. We adopt MLP generative model for MNIST (inspired by (Kingma & Welling 2013), and a ResNet decoder for CIFAR/ImageNet (with additionally a learnable variance clipped to (0.01, 1.0) (Van Den Oord et al. 2017). The recognition model is vision transformer (Alexey) 2020). enabling either autoregressive sampling (non-VQ methods) or direct encoding of codebook vectors (VQ methods) as depicted in Appendix For completeness, we also tested the ResNet encoder from (Van Den Oord et al. (017)) in the VQ setting, but observed no advantage. For LAFAN, both recognition and generative models are transformer-basee methods rely on decoder-only transformer while ELBO- -o--daee methods use an encoder-decoder architecture, we increase the number of layers and hidden units in the VQ decoders match their total parameter count as closely as possible to ours.
(See Appendix Because VQ-based
Training. All methods are trained with identical batch sizes and numbers of optimizer steps per dataset. We use Adam with weight decay and a base learning rate of 3 x 10 applying cosine decay when appropriate. For MNIST and CIFAR, we train for 250k steps with batch size of 256 on an RTX-3090. For LAFAN, we also train for 250k steps with a batch size of 256 on an A100- 40GB. For ImageNet, we run for 300k steps with batch size of 64 on A100-80GB GPUs (see Appendix[H] H)). Each experiment is repeated with 10 random seeds on MNIST, CIFAR, and LAFAN, and with 5 seeds on ImageNet due to the higher cost. In practice, we find that the autoregressive Gumbel Softmax is unstable at the base learning rate, requiring per-dataset tuning. Because of this instability, we exclude it from ImageNet experiments. In contrast, autoregressive GR-MCK remains stable in the ImageNet setting, though it requires more GPU memory.
We benchmark DAPS against a range of discrete latent variable models: (1) Gumbel- Softmax, (2) GR-MCK, (3) FSQ, and (4) VQ-VAE, For DAPS, Gumbel Softmax, and GR-MCK, we
Baselines.

---

Published as a conference paper at ICLR 2026
Vector Codebook
ResNet Decoder
Transformer Encoder
FCN
X- x(x(()
Po(x/e)
Pass
Pass
Pass 3
Pass
Figure 1: Overview of DAPS. Left: Images are split into patches and embedded by a feed-forward network before entering the encoder. Middle: The decoder generates latent sequences autoregres- sively: each step conditions on the encoder output and previously sampled latent codes (via causal masking, shown gray). Right: The generative model decodes latent embeddings into an image.
also implement von-autorrereesivve versions (denoted by the "NA" suffix), where the decoder-only transformer outputs all logits in single forward pass. Despite being less expressive, Gumbel-NA outperforms Gumbel on most datasets, which we attribute the difficulty of gradient estimation in autoregressive sampling. For completeness, we also evaluate the Gaussian VAE (equipped with the same latent capacity) and a PPO-style autoencoding baseline; however, we omit them from larger- scale experiments due their substantially weaker performance. For each method, we perform Bayesian sweep over hyperparameters (see Appendix [ D). centered around the default values used in the original works.
Ablations. We study the two main hyperparameters of DAPS and the ESS target-to evaluate the method's sensitivity their values. Specifically, we consider grid with 3 on CIFAR-IO. Our results show that {0.1, 1.0, 5.0, 10.0} and ESS target E {K/4, Appendix (). Intuitively, this schedule promotes annealing 3 yields the strongest performance (see exploration early in training and gradually shifts the focus toward high-quality reconstructions later on. In contrast, DAPS shows little sensitivity the choice of ESS target. This parameter primarily influences the adaptation of during training to satisfy the entropy objective. Overall, these findings highlight that /3 the more critical hyperparameter for performance, while the ESS target has comparatively minor effect, underscoring the robustness of DAPS to its setting. Finally study the utilization of the latent vocabulary (see Appendix G The marginal distribution q2z:[x) is dis- played in Fig. 2 indicating strong codebook utilization for DAPS compared to FSQ and VQ-VA.
DAPS (1+00.0))
DAPS
F3Q
VOLWEE
200
10e
Figure 2: Latent Code Utilization on MNIST Validation Dataset,
5 RESULTS
We now turn to an empirical evaluation of DAPS. We present key metrics, including the PSNR (peak signal-to-noise ratio), the 3-ELBO (the 8-weighted latent prior KL divergence plus the log likelihood of the reconstruction), and the FID score (computed using all 50k validation samples for ImageNet and all 10k validation samples for CIFAR-10).

---

Published as a conference paper at ICLR 2026
Method
MNIST (64 bits)
CIFAR-10 (576 bits)
ImageNet (1.28 KB)
LAFAN (640 bits)
3-ELBO PSNR
8-ELBO PSNR
FID
B-ELBO PSNR
FID
6-ELBO PSNR
FSQ
18.42
24.19
163.00
24.24
54.54
36.19
VQVAE
18.45
24.19
164.30
23.83
65.01
31.04
DAPS
-46.54
18.23
1185.51
25.21
157.27
87.0k
24.66
48.65
-949.78
36.81
DAPS-NA
-46.96
18.36
977.39
25.02
156.33
78.8k
24.40
57.43
1050.25
32.90
GRMCK
-62.25
16.78
217.45
22.69
179.88
60.7k
23.01
73.21
-1008.12
34.11
GRMCK-NA
-68.52
16.22
61.96
22.71
172.65
21.4k
21.85
99.15
-1530.51
26.93
Gumbel
-68.07
16.30
704.92
23.74
169.87
-998.49
34.51
Gumbel-NA
-47.09
18.21
785.35
24.27
162.04
85.2k
24.49
51.66 -1400.05
27.89
PPO
-104.66
14.56
416.31
23.30
176.29
Table Summary mean validation metrics (for best value obtained per seed) across seeds.
MNIST
CIFAR
Imageellee
LAFAN
-00
100k
-1k
Vall logpix00)
60
30%
-100
-1k
-1.8k
100k
200k
10Ck
Step
DAPS
Gumbel
VQg VAF
GR-MEK
Autrinng restive
Non-Autoongrreiee
Figure 3: Validation reconstruction log- log-likelihoods, log p(x]z), throughout training (3 0.01).
To facilitate better comparison between ELBO-based methods (which incorporate additional entropy regularization) and VQ-based methods, we present our results in two separate tables. In Table T the final value of 3 is treated as hyperparameter rather than part of the objective. As result, ELBO- based methods use a final | of 0.01 for MNIST, CIFAR, and LAFAN in Table Meanwhile, a final of 1.0 is used for ImageNet, which find works well due the large magnitude of the recon- struction log-likelihood relative to the latent prior KL divergence. We also seek to compare methods with the original ELBO objective mind. Consequently, Table presents a comparison of ELBO- based methods for 3 = 1. In both settings, DAPS exhibits strong performance, demonstrating the flexibility of our method in accommodating various objectives.
MNIST. DAPS achieves comparable performance baseline methods in terms of B-ELBO and PSNR on MNIST. However, unlike VQ-based methods, DAPS also produces visually high-quality digits when decoding samples from the latent prior (see Appendix, Fig. 21). We also train label- conditioned model on the learned latent space, the results for which are shown in Appendix Fig. 20 CIFAR-10. Our method outperforms all baselines CIFAR-10 10 across the relevant metrics. No- tably, non-auvoregressive DAPS (DAPS-NA) achieves a slightly lower FID than DAPS, though un- derperforms DAPS other experiments. Compared non- non-VQ- VQ-based approaches, DAPS reaches a high ELBO quickly and stably, without the hyperparameter sensitivity observed in methods such as autoregressive Gumbel Softmax (e.g.. temperature tuning).
Method (8 1) MNIST (64 bits)
CIFAR-10 10 (576 bits)
B-ELBO PSNR
3-ELBO PSNR
FID
DAPS
-80.64
17.13
857.77
25.27
158.61
DAPS-NA
-86.00
16.54
585.02
24.97
157.16
GRMCK
-101.86
14.83
-211.33
21.98
184.65
GRMCK-NA
-114.61
14.21
-780.04
20.90
183.46
Gumbel
-93.72
15.68
405.18
23.51
172.68
Gumbel-NA
-80.80
17.21
394.19
24.20
163.98
Table Mean validation metrics for 1.0 on MNIST and CIFAR-10.
8

---

Published as a conference paper at ICLR 2026
Figure 4: ImageNet 256 validation reconstructions. Top bottom: DAPS, FSQ, Groundtruth.
ImageNet. We evaluate DAPS the ImageNet-2566 dataset, demonstrating its ability recon- struet high-dimensional data using low-bit bottleneck. DAPS significantly outperforms traditional methods terms of reconstruction quality while maintaining high bit efficiency, Given randomly- sampled images from the validation set, DAPS produces high-quality reconstructions, as shown in Fig. 4 and Fig. 23 Finally, DAPS achieves the best 8-ELBO, PSNR, and FID scores comparatively.
Figure 5: latent sequence for hopping is decoded by DAPS and fed through inverse kinematics.
We also evaluate DAPS on a LAFAN dataset for the Unitree HI robot provided in Lo- coMuJoCo (Al-Hafez et al. 2023), which contains to-ontto-robot retargeted motion trajectories with inherent temporal dependencies and multimodal characteristics. Quantitatively. DAPS excels reconstructing this data, outperforming baseline methods in terms of 3-ELBO and PSNR. DAPS generates coherent and high-quality motion trajectories that align well with the true behaviors in LAFAN (see the supplementary videos). This makes DAPS amenable to downstream robotics tasks such hierarchical control, where high-level policy guide low-level policy by specifying desired trajectory, given a context. A reconstructed motion is shown in Figure 5
LAFAN.
Figure 6: A latent space sequence for dancing decoded by reinforcement learning policy and executed in physics simulator.
Downstream Task:
Goal-Conditioned Robot Control.
We construct a dataset of motion segments paired with contextual signals and targets. high-level policy is trained to process history of relative body poses while conditioning on two context signals: (i) a desired future COM velocity and (ii) a text description of the motion. The policy is trained to autoregressively generate discrete latents in the DAPS latent space in a distillation setup similar to PoseGPT (Lucas et al. 2022). These latents effectively replace raw motion references, providing a command space defined by text and COM trajectories that guides a low-level motion imitation policy (see Appendix(K). Supplementary videos illustrate the resulting robot behaviors and highlight the utility of DAPS in downstream control tasks. An example motion shown in Figure 6

---

Published as a conference paper at ICLR 2026
6 DISCUSSION
Overall, DAPS provides statc-of-the-art performance in reconstructing high-dimensional datasets, both for images and sequential motion data, outperforming existing methods across key metries. The scalability of DAPS, especially with respect to high-dimensional and sequential data, makes strong candidate for a variety of generative tasks. Our approach discrete variational autoen- coding via policy search offers a viable alternative to traditional methods like Gumbel-Softmax and VQ-VAE, particularly when dealing with nigh-dimensional datasets. Using standard set of VAE hyperparameters, the main parameters for DAPS are the desired effective sample size and B. This makes the training process comparable to that of VQ-VAE, which requires the specification of a commitment coefficient. Both methods introduce hyperparameters that are not highly sensitive, making them relatively easy to tune. While autoregressive discrete sampling allows for expressive multimodal latent generation, it introduces a computational bottleneck, particularly the sequence length increases. This limits the scalability of both DAPS and Gumbel-Softmax when compared to traditional VAEs and VQ-VAES, However, DAPS avoids backpropagating through autoregres- sive sampling, allowing it to scale to sequence lengths of 1,024 and likely beyond. This provides an advantage in terms of computational efficiency and scalability within the class of autoregressive methods, particularly on large datasets such as ImageNet 256. Our formulation also allows for fine- grained control over the tradeoff between compression and generalization through the of 33 in the ELBO. By adjusting B, we can effectively manage the tradeoff between generating high-quality re- constructions and generalization. This flexibility is key strength of DAPS over other methods, such as FSQ and VQ-VAE, which do not allow for easy control of entropy. Furthermore, our approach demonstrates stable training across seeds, which critical for real-world applications.
CONCLUSION
We introduced DAPS, policy-search-basee framework for training discrete variational autoen- coders without relying on reparameterization techniques. Our method integrates reinforcement learning concepts to optimize the variational encoder using weighted maximum likelihood, pro- viding both scalability and efficiency. Experiments on high-dimensional datasets such as ImageNet and LAFAN demonstrate DAPS' superior performance in reconstruction quality and its ability to generalize well across different tasks. By combining insights from policy search with variational autoencoding. DAPS opens the door to the efficient use of discrete latent bottlenecks for large-scale generative modeling. Its flexibility in balancing compression and generalization, along with its abil- ity to handle complex data like images and motion trajectories, positions DAPS as a powerful tool for both generative modeling and downstream control tasks. This work also paves the way for further exploration of policy-driven techniques in the optimization generative models.
ACKNOWLEDGMENTS
This work was supported by the German Science Foundation (DFG) under project RTG2761. We gratefully acknowledge support from the hessian.A AI Service Center (funded by the Federal Ministry of Education and Research, BMBF, grant no. 011S22091), the hessian Innovation Lab (funded by the Hessian Ministry for Digital Strategy and Innovation, grant no. S-DIW04/0013/0M the Lichtenberg high-performance computer of TU Darmstadt. This work was also supported by a hardware donation from NVIDIA through the Academic Grant Program.
and
REFERENCES
Abbas Abdolmaleki, Rudolf Lioutikov, Jan Peters, Nuno Lau, Luis Pualo Reis, and Gerhard Neu- mann. Model-based relative entropy stochastic search. Advances in Neural Information Process- ing Systems, 28, 2015.
Abbas Abdolmaleki, Jost Tobias Springenherg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- tin Riedmiller. Maximum posteriori policy optimisation. arXiv preprint arXiv:1806.0669: 2018.
Firas Al-Hafez, Guoping Zhao, Jan Peters, and Davide Tateo. LocoMuloCo: A comprehensive imitation learning benchmark for locomotion. In 6th Robot Learning Workshop, NeurlPS, 2023.
10

---

Published as a conference paper at ICLR 2026
Dosovitskiy Alexey. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv: 2010.1 2020.
Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient-free variational inference using policy search. In International conference on machine learning, pp. 234-243. PMLR, 2018.
Christian Daniel, Gerhard Neumann, Oliver Kroemer, and Jan Peters. Hicrarchical relative entropy policy search. Journal of Machine Learning Research, 17(93)11500, 2016.
Mare Peter Deisemroth, Gerhard Neumann, and Jan Peters. A survey on policy search for robotics. Foundations and Trends0 in Robotics, 2(1-2):1 2013.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. arXiv preprini arXiv:1711.01123. 2017.
Evan Greensmith, Peter L Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. Journal of Machine Learning Research, 5(Nov))1471-1550,
2004.
Shixiang Gu, Sergey Levine, Ilya Sutskever, and Andriy Mnih. Muprop: Unbiased backpropagation for stochastic neural networks. arXiv preprint arXiv:1511.05176, 2015.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatioo with gumbel-softmax, arXiv
preprint arXiv:1611.0111 2016.
Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint
2018
Thomas Lucas, Fabien Baradel, Philippe Weinzaepfel, and Grégory Rogez. Posegpt: Quantization- based 3d human motion generation and forecasting. European Conference on Computer Vision, pp. 417-435. Springer, 2022.
Felipe Maia Polo and Renato Vicente. Effective sample size, dimensionality, and generalization in covariate shift adaptation. Neural Computing and Applications, 35(25):18187-18199 2023.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantiza- tion: Vg-vae made simple. arXiv preprint arXiv:2309. 15505, 2023.
Alberto Maria Metelli, Matteo Papini, Nico Montali, and Marcello Restelli. Importance sampling techniques for policy optimization. Journal of Machine Learning Research, 21(141):1-75, 2020.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In
International Conference on Machine Learning, pp. 1791-1799, PMLR, 2014.
Gerhard Neumann. Variational inference for policy search changing situations. 2011.
Max Paulus, Chris J Maddison, and Andreas Krause. Rao-blackwellizingg the straight-through gumbel-softmax gradient estimator. arXiv preprint wwXiv:2010004838, 2020.
Jan Peters and Stefan Schaal. Reinforcement learning of motor skills with policy gradients. Neural networks, 2144)6622-697, 2008.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference Artificial Intelligence, volume 24, pp. 1607-1612, 2010.

---

Published as a conference paper at ICLR 2026
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Iyya Sutskever. Zero-shot sost-to-image generation. International conference on machine learning, pp. 8821-8831. Pmlr, 2021.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
John Schulman. Trust region policy optimization. arXiv preprint arXiv:1502.05477, 2015.
Francis Song, Abbas Abdolmalcki, Jost Tobias Springenberg, Aidan Clark, Hubert Soyer, Jack W Rae, Seb Noury, Arun Ahuja, Siqi Liu, Dhruva Tirumala, al. V-mpo: On-policy maximum a posteriori policy optimization for discrete and continuous control. arXiv preprint arXiv: 12238, 2019.
George Tucker, Andriy Mnih, Chris Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Tim Van Erven and Peter Harremos. Rényi divergence and kullback-leibler divergence. 1EEE Transactions on Information Theory, 6777):377,-3820, 2014.
Quan Vuong, Yiming Zhang, and Keith W Ross. Supervised policy update for deep reinforcement learning. arXiv preprint 1rX1v 2018.
Joe Watson and Jan Peters. Inferring smooth control: Monte carlo posterior policy iteration with gaussian processes. In Conference Robot Learning, pp. 67-79. PMLR, 2023.
Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhuber. Natural evolution strategies. The Journal of Machine Learning Research, 15(1):949-98)) 2014.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229-256, 1992.
Brian Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, University of Washington, 2010.
12

---

Published as a conference paper at ICLR 2026
APPENDIX
A DAPS DERIVATIONS
A.1 SOLVING THE CONSTRAINED OPTIMIZATION PROBLEM
We wish solve the following constrained optimization problem:
max [ [p(x) E ((2, dr (H(((z/x))
(9)
s.t. Dxl(9(z(x) go(z)x)) cn:
[z(zx)= =1
Using Lagrangian multipliers, we can convert this into an unconstrained optimization problem:
I n. = p(x) Σ ((x)x)x((x, 2) + (H(((z)x))
+ (( [ p(x) Σ ((22x) log g(x)x)p(x)
g0(z)x)p(x)) dx
(10)
x
+ p(x)
(xx))) Σ [((|x)) dx
Diffcrentiating J with respect to q(z)x) xx) and setting to zero yields:
Σ
og p(x))(()x) log go()))
89
-mlog((21x) x(x) + A(x, 2) - B logg(z/x)) dx
_ p(x) (1)g 17 (z)g(z)x) - A(x) + A(x, 2) - 3 log ((/2x) 3
(11)
Re-arranging terms gives us the optimal nonparametric target distribution, g (z(x/x):
(z/x) exp A(x, z) + 17 1o qo (z)x)
+ 8 + A(x)
exp
17 3
174 3
(12)
A(x, 2)
3 log qo (z|x)
1 B + X(x)
++x))
(2|x) exp
exp
17 + 3
+ 3
we can see that the optimal policy is posterior distribu- tion with prior that is the parametric policy qo (z)x). This gives us the following dual function:
whereby re-writing the last line in Eq. 12
((n,A(x)) = / ((x) + 176,
+ [ p(x) Σ q- (2|x))
(x/ex))(n log g0(2))) - (n + 3) (0)g x(x) + A(x, 2)) dx
(xx))(x)dxx+xxx)x + men + (1 (1+ 3) p(x) Σ 4 (z|x)dx
(13)
Since the optimal X(x) normalizes q² (z|x), have:
13

---

Published as a conference paper at ICLR 2026
1 B x
2) + +1 log go(z)xx
Σ exp
exp
17 9+
+
(14)
A(x, 2) + 1) log g0 21/x) -
xxx)) = (1 + 3) log Σ cxp
n7 B
Plugging this into the dual, we get:
A(xx, z) + 17 log qo (z)xx ) 3 3ddx + men
(x)) p(x))(( (1) + 1og Σ exp
11 B
+ (1) ( p(xx Σ q" (z|x))xx
(15)
( 3) / p(x) log Σ exp
A(xx 2) + 11 log qo (2)xx)
] |xx + 176,
+ 3
We can now differentiate with respect to n to optimize our stepsize.
6g(n, [1 + B) / p(x) log E exp A(xx2) 2) + +1 log (2(()))) )] dx + men)
(16)
δn
ōn
7 + 8
Letting
En exp
((xx2)+ log (x
and log v. have that oe w- and:
n+
3 log (0(
A(xx 2))
exp A(xx, 2) + r log go (z(x)))
Σ
(17)
(1) 3) 3)2
7 B
Applying the chain rule, using
5u by
bu 6v 8z 8n.
get:
A(x,z)+n no (2)
8g(n, (x)))
Blog 90 (z(x) -((x,2) 2)
exp
p(x)
Σ
n+B
8n
x
17 + 3
Ex exp
A(xxx' ) log go (x' xx)
n+0
+ +og Σ exp
[((x,2) 2) + + log go (2/xx)
]] dx + cy
17 7 8
(18)
= - p(x) Σ log (0(2|x) (z(2))
q- (z|x) + log Zyp dx En
DKL (q" (z|x) g0(2)x))]
Ey - Ex-p(x)
It can be seen that computing this value requires computing the partition function, Zge, and this requires cnumerating over all possible sequences of z. Similarly, the outer summation cannot be easily computed, rendering the optimization of n using the KL divergence difficult in our setting. For this reason, we make use of the Effective Sample Size to optimize as described below.
A.2 ESS-BASED TRUST REGION
Setup. Given arbitrary datapoint let go (z|x) be the proposal distribution. Let q" (z|x) denote the unnormalized target from Eq. 5 with partition function
7" (x/x)
((x)
[z)x))
(x/x)
7(x)
Drawing Z1-k ~ (-(x))), define the normalized importance weights
9º (zx
(Zx x)
14

---

Published as a conference paper at ICLR 2026
Definition.
The effective sample size (ESS) associated with weights fwk] is
(E k=
ESS (x)
Σ k= w
Scale invariance.
If instead we had used the unnormalized target q", then
9 (zx
Wk
[(x) wk.
go (zx
The ESS is invariant to such rescaling:
Σx w6 )2
E k 2x
EEx wk) Σ w
Σx
Z Σx w
Thus one may compute ESSK k using either q" q"
Normalized ESS ratio. For convenience, we define the normalized ESS ratio
1 K
1 1
R Em=
(x(x)
(x)
K 2k=1
R Σ W
This form makes both numerator and denominator empirical means, so the strong law of large num- bers applies directly.
Definition (Rényi-2 divergence). the order-2 Rényi divergence is
For two distributions p, g with absolutely continuous w.r.t. g.
D2(p11g) log Σ p(z)2
((z)
Lemma (Population ESS and Rényi-2). Assume Egg [w²] < 0o, which holds in the case of finite discrete distributions with full support. Then, as K + 00,
(Eg. (w))
(x(x)
p(x)
exp( D2(a" (lq0)).
Ego [w2)
Eg [w2]
Proof. By the strong law,
1
1
Wk
[w),
w - Eayo (w2),
K
k=1
almost surely. Since q" is normalized,
(z(xxx 1.
E [w]
Σ 20(2/xx
qo (z(2(x))
Moreover, by the definition of D2.
77" (2|x)
9" (z|x)2
Eggo [w]] Σ (0(2
=exp(Da(("lloo))
90 (z(x
(o(z
Therefore p(x) 1/ /R exp(-D2(1 l|go)).
Corollary (KL trust region from ESS target). By monotonicity of Rényi divergences in their order (Va Erven & Harremos 2014).
DKL ((q4)) = Di(a" S D2(9" [[go)
Hence, if Px (x) is adapted to match target level Ptargeto then the update satisfies the Kl. trust- region bound
DKL(g" < - log Prarget.
Remarks. (i) PK is biased unite-sample estimator of the population p(x). but converges almost surely by the strong law. (ii) In practice treat 1) a trainable parameter and update it with stochastic gradient descent on (PK Ptarget) )2 thereby driving PK toward Ptarget-
15

---

Published as a conference paper at ICLR 2026
B TRAINING DYNAMICS AND STABILITY
MNIST
(schodiic)
Normalized ESS (ta3get-0.33 0.75)
Val KL to prior (4)
Val loppC-le) (1)
-40
0.70
0.50
-50
0.20
20
-100
0.00
1024
2002
100%
200k
120k
200k
1004
200k
Step
Step
Bhall
Step
Step
001
CIFAR-10
(schedule)
Normalized CSS (tarett003g)
Val KL prior (1)
Val logpeedet (1)
500
0.31
-2k
100%
700k
100+
200k
100k
7004
10Ck
700k
Step
Step
Final
Step
Step
ImageNet 256
(echedule)
Normalized ESS
(target=0 25)
Val KL price (1)
Voll Iogariele (1)
0.95
100%
W
46.
100%
0.22
100k
2004
300%
200%
2004
500k
103k
2004
300k
100k
200k
300k
St.p
Stop
Final
Stop
Stop
LAFAN
(schodule)
Nomallzzdd ESS (target-0.33)
Vol KL to prior (1)
Val logu a) (?)
-1A
0.54
300
-2%
-3%
0
1.0k
200k
100%
200%
1024
2024
0
1.0kk
200k
Stop
Stop
Final
Stop
Stop
Figure 7: A summary of key metrics for DAPS throughout training.
16

---

Published as a conference paper at ICLR 2026
C ABLATIONS: 3 AND ESS-TARGET (CIFAR-10)
We scan 3 E {0.1, 1,5, 10) and ESS- SSS-target E (K/4, K/2,3K/4) and report the metrics below.
(schedule)
Val logp(z|=)
Val Latent Prior KL (4)
10
400
1k
300
0
200
-1k
100
-2k
100k
200k
200k
100k
200k
Step
Step
Step
0.1
5 0.01
5 to
Figure 8: (1) An overview of different | values, averaged across the 3 different ESS targets. (2) The scheduled / curves are averaged over the 10 seeds from the final experiments, which use an ESS target of 0.33.
Val log p(1]=) (1)
Val Latent Prior KL (1)
150
400
100
300
200
50
100
100k
200k
100k
200k
100k
200k
Step
Step
Step
ESS Target
0.5
Figure 9: The joint performance various 33 and ESS combinations. It can seen that the choice of ESS target has the desired effect, namely, the attenuation of the 77 trajectory throughout training.
D ABLATION: NUMBER OF BASELINE SAMPLES (MNIST)
Advantage
Val ELBO (t)
Val L2 (1)
-100
-10
-120
-15
140
5
100k
200k
100k
200k
100k
200k
Step
Step
Step
Value Samples
16
32
Figure 10: Key metrics w.r.t the number of samples used in the advantage estimation.
17

---

Published as a conference paper at ICLR 2026
E SUMMARY OF KEY METRICS
MNIST
PSNR
SSIM
(AELBO
(-EL8O
15
-55
16
0.8
100
0.2
200k
Stop
VQ-VAE
F5Q
GK-MCK
Autorggrenta
Non-Autorepresiin
CIFAR-10
PSNR
Val FID
# ELBO
SSIM
1000
160
3000
50%
100k
150k
200k
250k
50%
100k
150k
2006
25th
50k
100k
150k.
200k
250k
50k
100k
150k
200k
250k.
Gradient Step
DAPS
Curbed
VO-VEE
F5Q
MR-MCK
FPO
Autorogrosiee
Non Non-Ausiergnnsiieg
ImageNet 256
Val PSNR
Val FID
# ELBO
Val SSIM
120
100000
24
000000
0.5
100k
200k
300k
100+
200k
500k
100k
200x
300k
1.0kk
200k
500+
Gradient
Autoong
LAFAN
JIELBO
300
1500
1500
20
2000
2000
Gradient Step
DAPS
Gumbel
VQa VAL
1SQ
Gil MUK
Autiong monuve
Non-Autee Peee-Autooggmette
Figure 11: A summary of key metrics averaged over all seeds (for 3 + SD.
18

---

Published as a conference paper at ICLR 2026
F BOTTLENECK SCALING EXPERIMENTS
-00
A-ELBO
PSNR
-80
16
Val
-100
-120
140
Tiny
Small
Medium
Large
Tiny
Small
Medium
Large
Size
Size
DAPS
Gumbel
VQ-VAE
Figure 12: Across bottleneck sizes (tiny: B-4, V=4 = 8 bits; small: B=4, V=16 16 bits;
medium: B-8, V- 32 bits; large: 8-8, V- = 64 bits), DAPS exhibits consistent scaling behavior on MNIST. In the left plot, validation PSNR increases predictably with latent capacity (bits B log, V). and DAPS remains competitive with VQ-VAE across all settings. Notably, VQ-VAE does not incorporate an entropy term, whereas DAPS (trained with 8 = 0.01 for all sizes) includes the comparison reflects slightly different optimization objectives.
mild entropy regularization, In the right plot, DAPS generally outperforms Gumbel across all bottleneck sizes in validation 3- ELBO, with the gap widening as capacity increases. Together, these results verify that DAPS scales as expected with latent capacity while maintaining sufficient reconstruction performance.
21
Val PSNR
Ø-ELBO
Tiny
Small
Medium
Large
Tiny
Small
Mediuu
Large
Size
Size
DAPS
Gumbel
VQ-VAE
Figure 13: Across bottleneck sizes (tiny: B-4, V-512 = 36 bits; small: B-16, V-512 144 bits; medium: B-64, V-64 = 384 bits; large: B=64, V= = 576 bits), DAPS exhibits strong scaling behavior CIFAR. In the left plot, validation PSNR improves with latent capacity, and DAPS remains competitive with, and slightly outperforms, VQ-VAE as the bottleneck increases. A similar trend to Fig. 12 appears in the right plot: DAPS generally outperforms Gumbel in vali- dation B-ELBO, while not exhibiting an advantage at the tiny bottleneck size. For all models and bottleneck configurations in Figs. 12 and 13 we performed hyperparameter sweep and report the best-performing setting for each. We use 5 seeds for the Tiny, Small, and Medium runs, and the 10 seeds from our main experiments for the Large runs.
19

---

Published as a conference paper at ICLR 2026
G
LATENT CODE UTILIZATION
Seed
Seed
Seed 2
DAPS (3=0 0.0
5000
4000
DAPS
Number Hits
-300
F50
2000
1000
VQ-VAE
Figure 14: Marginal (empirical) distribution of g(2:) the full CIFAR validation set. Each grid shows the full vocabulary V of the recognition model, reshaped for visualization, with cell values in- dicating the number of times each latent code is used. DAPS effectively utilizes all codes, whereas VQ-VAR struggles to do so. Although FSQ aims to address this limitation and indeed achieves stronger utilization, its vocabulary usage remains much sparser than DAPS, highlighting the favor- able entropy properties induced by our objective.
20

---

Published as a conference paper at ICLR 2026
Seed
Seed
Seed
DAPS (3-g 01)
600
500
DAPS
400
Number of Hits
300
PSQ
200
100
VQ-VME
Figure 15: Marginal (empirical) distribution ((z.) on the full MNIST validation set. Each grid shows the full vocabulary V of the recognition model, reshaped for visualization, with cell values indicating latent-code usage. DAPS again achieves broad and balanced utilization of the vocabulary, whereas FSQ exhibits noticeably sparser usage and VQ- VAE contains several low-hit regions. This further highlights the favorable entropy properties promoted by our objective.
21

---

Published as a conference paper at ICLR 2026
H COMPUTATIONAL OVERVIEW
We provide a brief comparison of the computational charactcristics of the methods considered, in- cluding memory usage, parameter counts (for both recognition and generative models), and training speed (measured in gradient steps per minute).
H.1
MODEL PARAMETER SUMMARY
Recognition Model
XUP
MNIST
CIAR
Imageket 250
LAHAN
1.20
1.00
0.00
Autorepentee
sttrrggeentoo
Generative Model
109
MNIST
CIFAR
Imagehic 256
LARNN
Model
Model
Midel
Model
Autorggraais
Nen-Autorggemaiii
Figure 16: comparison of model parameter counts.
H.2 AUTOREGRESSIVE MODELS: RESOURCE USAGE AND SPEED
Speed Comparison
MNST
CIRRR
ImageNet 256
LAFAN
850
Aok Stops/Mt
600
100
00000
400
300
500
Gumbel
LNJ
IKO GPU
Memory Comparison
MNET
20.0
Menray
0.4
500
00
Mode
Mide
Mode
DAPS
MRC
Figure 17: A comparison of resource usage and speed for autoregressive models.
22

---

Published as a conference paper at ICLR 2026
H.3 NON-AUTOREGGRESIVE MODELS: RESOURCE USAGE AND SPEED
Speed Comparison
Imagelle:
6909
10000
Avt. Sk.26 10nn
(000
230
500000
100
2001
Model
Memory Comparison
MNIST
CIRAR
Imagelttt 256
LAFAN
10
8--
Menney
21-1
0.4
10-
02-
00-
D-
Moddl
Model
Model
Medd
Figure 18: A comparison of resource usage and speed for non aon-autrregresive models.
H.4 COMPUTATIONAL CONSIDERATIONS
Model parameter comparison Figuure 16). We compare the parameter counts of the recognition and generative models across methods. All baselines use parameter budgets that are comparable to, or larger than, those of our primary model (DAPS-aute autrrerreeeesssi ensuring a fair evaluation. Minor deviations reflect architectural differences inherent to each method.
Autoregressive resource usage (Figure 17). We next compare memory usage and training speed for the autoregressive variants. On ImageNet 256, both Gumbel and GR-MCK exceed the memory capacity of a single GPU and must be trained across two GPUs; in contrast, DAPS fits the full batch on single A100, though we use 2 GPUs in practice. Although we do not evaluate Gumbel on Ima- geNet due to instability, we report its memory and speed characteristics for completeness. Memory usage for GR-MCK is dominated by the MC-K factor, explaining its substantially larger footprint relative to Gumbel. We note that only 8 value samples are used for LAFAN Fig. 17 and [8 which slightly degrades performance (PSNR=35.15 and (-ELBB0=985 for DAPS; PSNR=32.50, B-elbo=- 1067 for DAPS-ÑA). Overall, we observe that DAPS scales more favorably on ImageNet compared Gumbel, which hypothesize is related to the interplay between backpropagation through time (BPTT) gradients and the quadratic scaling properties of transformers. These results highlight that implementing DAPS is practical even larger bottleneck sizes.
Non-autoregressive resource usage (Figure 18). As expected, we observe that DAPS-NA is faster and more memory-efficient than its autoregressive counterpart. However, DAPS-NA is generally slower than baseline methods. Nonetheless, the computational cost remains practical in absolute terms (e.g., under 10 minutes on MNIST and under 2 hours on CIFAR for 250k gradient steps). However, we focus on the autoregressive variant throughout the paper, given that DAPS-NA gener- ally underperforms autoregressive DAPS. Finally, we again observe that the MC-K parameter drives GR-MCK's memory and speed profile. Overall, both variants of DAPS achieve hetter performance at the cost of higher computation. However, this does not prevent scaling to challenging data sets such ImageNet, even for the autoregressive variant.
23

---

Published as a conference paper at ICLR 2026
1
VISION MODEL ARCHITECTURE AND PARAMETER SUMMARY
Vision Recognition Model Architecture
ResNet Generative Model Architecture
Patch Feed Forward (Encoder Input)
Latent Emhedding (ResNet Input)
Dense
Kemei (Dasch H) Bas (H)
Table (V, H)
I. Latent Embedding (Decoder Inpot)
ResNet Decoder
Einbed
Tuble (V, H)
Come
Kemel (3,3, 3, 24, 1gg) Bias (Kg))
Residual Block
Transformer
Conv
Kornel (3, 3, Hg. H2) Bas (H2)
Encuder Block 1
Batch Norm
Scale (Hfz) Dius (1fg)
Layer Norm
Scale (.1) + Bias (21)
Cony
Kernel (1, 1, Hg. Hg) Dias (Hg)
MuitiHead Attention
Query (H, Ni, H (NN)) )
Key (H, No. H(Nr))
Batch Norm
Seale (Hg) Bins (Hg)
Valuee (H, NA. M/NN))
Cut NN. H/N,, H)
Layer Norm
Scale (8)) + Bias (H)
Residual Block
Donse
Kornel (H, HR + Bios (HR)
Repeat: Com
BN Couw - BN
Denve
Kornel H) (H)
Com'T
Kornel (4, Hg. Ha) Bia (Ma)
Rutch Norm
Scale (Hg) Baa
Encuder Block N
Repeat: I.N MHA INN
Com'T
Scale (4, Hg. Bias (6)
Encuder-
Dncuder-Decuder Decoder Block I
Generative Model. The vision generative model architecture is a standard ResNet decoder. This model takes as input latent sequence (of length 64 for CIFAR-10, and length 1, 024 for Ima- geNet) and converts it to sequence of em- beddings using the embedding table. This is then reshaped into square tensor of embeddings (8x 8x8 x 128 for CIFAR-10, and 32 32 x 128 for ImageNet) to processed by the ResNet. After which, convolutional transpose layers used to upsample the tensor to the final image dimension, including additional 3 channels for the RGB variance. See [.] and for the corresponding variable values.
Layer Norm
Scale (82) + Bias (21)
MultiHead Attention
Query (R, Ni. 8 M/NA)
Key (H, No, H.NN)
Value (J, NA. H (
(N,, H/N,,
Layer Norm
Scale (II) + Bius (21)
MultiHead Altention
Query (H, No. H/N
Key (H, No. H(N,,)
Value (J, No. H (.
(N,, H/N H)
Layer Norm
Scale (H) Bits (H)
Donse
Kornel (H, HR) + Biss RH
Dence
(,R H) (11)
Enewder-Deeoder Block N
Repear: I.N MHA IN X-MHA
Layer Norm
Rins (H)
Dense
Kemel (H, V Rins (V)
Recognition Model. The recognition model architecture is a standard vision transformer. Cross attention is used for autoregressive-based methods, while self-attention is used for VQ- based methods. We also tried using the ResNet encoder proposed in VQ-VAE, but found per- formance benefit. See L. and I.2 for the corre- sponding variable values.
24

---

Published as a conference paper at ICLR 2026
1.1 IMAGENET-256 VISION HYPERPARAMETERS
Conlig (ImageNet-256)
model:
block_size: 1024
vocab size: 1024
sccabssize:
embed_dim: 128
recognition model:
transformer:
num heads:
N.h
num_layers:
N
autoregressive:
non non_auttoregressive:
mip_ration
qkv_ 128
num pan__aatches:: 1024
32x32 patches
patch_ ffwd: (dense (128)
generative model: :
resnet: hidden dim:
activation: relu
norm: batch_norm
residual:
num blocks: 2
N
hidden dim: 64
activation: relu
norm: batch_ norm
Network Hyperparameter Summary. The configuration above lists the architecture-specifi hy- perparameters used for training ImageNet 256 (shared by all methods). Variable names, indicated by comments on the right-hand side, correspond the variables the Vision Transformer Diagram featured in[ The autoregressive sampling used by ELBO-based methods can be sped up through the use of KV caching, as past activations do not need to be recalculated on subsequent forward passes. Our implementation uses decoder cache window size of 8 blocks.
Method (final ()
Commit. coeff.
Temp.
Init. 0
Target ESS,
MC-K K (DAPS)
LR
DAPS (1.0)
50
0.25
3e-4
DAPS-NA (1.0)
50
0.25
3e-4
GRMCK (1.0)
1.0 + 0.75
10
10
1c-4
GRMCK-NA (1.0)
1.0 + 0.1
10
10
3e-4
Gumbel-NA (1.0)
1.0 + 0.1
50
3e-4
FSQ
0.25
3e-4
VQVAE
0.25
3e-4
Table 3: ImageNet256 model-specific hyperparameters (final | = 1.0).
25

---

Published as a conference paper at ICLR 2026
1.2 CIFAR-10 VISION HYPERPARAMETERS
Conlig (CHFAR-10)
model:
sioe_ssze:: 64
vocab size: 512
embed_dim: 128
recognition model:
transformer:
num heads:
N.h
num_layers:
N
autoregressive:
non non_autoregressive:
halp_aatoo:
autoregressivv: 4
non non_auu non_autoregresssvee 6
dim: 128
num image_ patches: image_patches: 64
8x8 patches
patch ffwd: [dense (128)
generative :
resnet:
hidden dim: 64
A_2
activation: relu
norm: batch_norm
residual:
num blocks:
hidden 64
activation: relu
norm: batch_norm
Network Hyperparameter Summary. The configuration above lists the architecture-specife hy-
perparameters used for training CIFAR-10 (shared by all methods). Variable names, indicated by
comments on the right-hand side, correspond to the variables in the Vision Transformer Diagram
featured in This design follows the architectures presented in Van Den Oord al. (2017) and
Alexey 2020).
Method (final (
Commit, coeff.
Temp.
Init.
Target ESS, MC-K K (DAPS)
LR
DAPS (0.01)
0.33
DAPS-NA (0.01)
0.33
GRMCK (0.01)
+ 0.75
3e
GRMCK-NA (0.01)
1.0 + 0.1
100
3e-4
Gumbel (0.01)
1.0
1e-4 4 > le-5
Gumbel-NA (0.01)
1.0 + 0.1
6
3e 3e-4
FSQ
0.25
3e-4
VQVAE
0.25
3e-4
DAPS (1.0)
6
0.33
3e-4
DAPS-NA (1.0)
6
0.33
8
3e-4
GRMCK (1.0)
1.0 + 0.75
3
100
3e-4
GRMCK-NA (1.0)
1.0 - 0.75
3
100
3e-4
Gumbel (1.0)
1.0 + 0.75
6
3e- 4 + 3e-5
Gumbel-NA (1.0)
1.0 0.1
6
3e
Table 4: CIFAR-10 model-specific hyperparameters for both 3 0.01 and 1.0.
26

---

Published as a conference paper at ICLR 2026
1.3
MNIST VISION HYPERPARAMETERS
Conlig (MNIST)
model:
slock_size: 8
vocab size: 256
sccabssize:
embed_dim: 64
recognition model:
transformer:
num heads:
N.h
num_layers:
autoregressive:
non :
halpraatio:
autoregressive: 4
non ron_autoxeereessive 6
vq fsq: /
gkv_dim: 54
image patches: 8
patch_ ffwd: [dense [
generative_modee
ffwd:
[dense (64), dense (256)
activation: relu
Network Hyperparameter Summary. The configuration above lists the architecture-specifi hy- perparameters used for training MNIST. Variable names for the recognition model, indicated by comments on the right-han hand side, correspond to variables in the Vision Transformer Diagram fea- tured in The generative model fully connected, feed-forward network with two hidden layers.
Method (final
Commit, coeff.
Temp.
Init.
Target ESS,
MC-K (DAPS) LR
DAPS (0.01)
0.5
0.33
3e-4
DAPS-NA (0.01)
0.5
0.5
3e-4
GRMCK (0.01)
1.0 -> 0.75
3e-4
GRMCK-NA (0.01)
1.0 + 0.1
0.5
100
3e-4
Gumbel (0.01)
1.0
0.5
30-4
Gumbel-NA (0.01)
-
3e-4
FSQ
0.25
3e-4
VQVAE
0.25
3e-4
DAPS (1.0)
0.75
3e-4
DAPS-NA (1.0)
0.5
3e-4
GRMCK (1.0)
10 - 0.75 0.75
100
3e-4
GRMCK-NA (1.0)
1.0 0.75
3e-4
Gumbel (1.0)
1.0 - 0.75
3e-4
Gumbel-NA (L.0)
10 100000 0.1
3e-4
Table 5: MNIST model-specific hyperparameters for both 3 = 0.01 and 3 = 1.0.
27

---

Published as a conference paper at ICLR 2026
J TRAJECTORY MODEL ARCHITECTURE AND PARAMETER SUMMARY
Trajectory Recognition Model (Transformer)
Trajectory Generative Model (Transformer)
1D Resample Convolution (Time)
Latent Embedding Tranntormer Input)
Con''
Kernel [K F.. H - Bizs (2)
Embed
Table (V. .2)
Transformer
Transformer
Encoder Block
|
Decoder Block
Layer Norm
Scale (.2) - Bias (H1
Layer Norm
Scale (21)
Bius (21
MultiHead Attention
Head Attention
Query (H, No. H/NG)
Key (H, No. H/NN,)
Query (F7, Nr. M/Na)
Key (H, Nh. H/N()
Value (H, Nr, (/N
Out (Nr, H/N,, H)
Value (H, Nr, H/N,)
Out (N,, H/NN,
Layer Norm
Scale (H) + Bits (H)
Layer Norm
Scale (H) - Biss (H)
Dense
Kornel (H, HRR a Biss (HR)
Kornel (H, HR) + Biss (HB)
Kornel (HR, H) (H)
Kernel (HR, H) Bias (H)
Encuder Block
N
Decoder Block N
Repeat: LN MHA LN
Reyear: LN - ->
Enewder-Deoddr Block
Layer Scale (H)
Layer Norm
Scale (I2) + Bius (21)
Dense
Kernel (Hr R) + Biss (H)
Multi Head Attention
Query (H, No. H/N J
Key (H, No. H/N,,)
ID Resample Convolution (Time)
Value (JI, N₁. H (
Out (Nr, H/NN, 11)
Come
Kamei (Kgg H, H) Bias (H)
Layer Norm
Scale (H) + Bits (H)
Output Head
MultiHead Attention
Query (H, No. H/N,)
Key (H, No. H/N,))
Dense
Kernel (H. F) (F)
Value (H, M/Nx)
(Nu, M/N,,
Generative Model. The trajectory generative model architecture is similar in symmetry the trajectory recognition model. After processing the embedded latent sequence, the sequence is resampled by a ID convolutional network to reach the desired trajectory segment length. The result is a continuous trajectory of length 32 with 127 features per timestep.
Layer Norm
Scale (H) + Bits (H)
Dense
Kernel (H, HT) + Bius (2)
Dence
(HR, H) Bias (I2)
Encuder- Decuder Block
Repeat: LN - MILA - LN + X MIIA - LN
MLP
Noom
Scale
Bins (H)
Donse
Kemel (HI, V)+ Dias (V)
Recognition Model. The trajectory recogni- tion model architecture shares similar struc- ture to the vision recognition model. Instead of using an image patcher (as in the ViT), we use a ID convolutional network to embed the trajectory sequence. Cross attention is used methods, while self-
for autoregressive- attention is used for VQ-based methods.
28

---

Published as a conference paper at ICLR 2026
J.1 LAFAN TRAJECTORY HYPERPARAMETERS
Conlig (Trajectory: LAFAN)
model:
block_size: 64
vocab size: 1024
V
embed_dim: 128
trajectory:
length: 32
Segment Tength
feature dim: 121
F
recognition model:
gesample_convll
cutput__channela. {mmddl embed_dim)
stride: 2
kernel size: 4
K 1
transformer:
heads:
N.h
num_layers:
autoregressive:
non non_autoreen non_autorrerressiv:: 6
mip_ratio:
R
qkv_dim: 128
generative model:
transformer:
<<: *zecognition ctclioon_ooeleeemm .rranseede..rrrnsfrreeerrr
resample_convld: convld:
output_length, #ltraajectory.Ienthh
syyput_channells $/model.cmbed_ce cmbed_dim)
stride: 2
kernel size:
K_2
output_ffwee ffwd: [denset (trajectory Feature_dim]]
Network Hyperparameter Summary. The configuration above lists the architecture-specff. hy- perparameters used for training the LAFAN dataset. Variable names, indicated by comments on the right-hand side, correspond the variables in the Trajectory Transformer Diagram featured in [ The recognition model and generative models are highly symmetric. These networks reconstruct randomly sampled trajectory segments length 32 with feature dimension 121. In order decou- ple the block size from the length of the trajectory segment, we use 1D convolutional network to resample the trajectory length the model's block size. This operation is then reversed the output of the generative model to achieve the desired segment length.
Method (final ()
Commit, cocff.
Temp.
Init. 3
Target ESS,
MC-K K (DAPS)
LR
DAPS (0.01)
2
0.33
3e-4
DAPS-NA (0.01)
1
0.33
16
3e-4
GRMCK (0.01)
1.0
2
10
1e-4 + 1e-5
GRMCK-NA (0.01)
1.0
1
10
3e-4
Gumbel (0.01)
1.0
2
1e-4 4 + 1e-5
Gumbel-NA (0.01)
1.0
0.75
3
3e-4
FSQ
0.25
3e
VQVAE
0.5
3e 4
Table 6: LAFAN model-specific hyperparameters (final 0.01).
29

---

Published as a conference paper at ICLR 2026
K
DOWNSTREAM ROBOTICS EXPERIMENT DETAILS
Dancing
Running
Walking
Figure 19: In our downstream experiment, we use different language commands to generate motion styles with DAPS, and use physical consistent decoder given by RL policy and robot simulator to decode the discrete latent representation.
We train low-level policy with DeepMimic imitate motions from the LAFANI dataset on Unitree HI robot simulated in MuJoCo. Unlike the conventional DeepMimic formulation, where motion reference information for the next timestep (site positions, joint angles, and velocities) ex- plicitly given as goal to the policy, we provide the policy with compact embedding of the desired future motion as the goal. We condition the high-level policy-which produces these embeddings on language information and motion goals. For the former, we use BERT embeddings of brief description of the motion file. For the latter, we provide the encoder with the last 32 timesteps and future 32 timesteps of the center of mass (COM) trajectory, relative to the robot's base frame. The resulting output a compact vector of 64 latent indices.
We train the low-level policy conditioned these embeddings with PPO using the LocoMujoco framework (Al-Hafez et al.) 2023). The policy transforms the embedded indices with a ID con- volutional layer, then concatenates the resulting features with standard proprioceptive information: joint angles, joint velocities, IMU information (projected gravity and torso angular velocities), and the previous timestep's action, ultimately feeding these into a multilayer perceptron of size [2048, 1024, 512] with elu activations, resulting in 19-dimensional vector of torques to be applied to each motor of the Unitree HI robot. The policy executed 100 Hz.
L DATASETS AND PREPROCESSING
Images. For ImageNet-256,, we randomly crop and flip the images. All images are standardized using their empirical mean and standard deviation. CIFAR-10 and MNIST are also standardized similarly.
Trajectories (LAFAN). For the LAFAN dataset, the motion sequences normalized to have zero mean and unit variance along each feature.
30

---

Published as a conference paper at ICLR 2026
DAPS
Gumbel
4
0
6
0
D
0
0
6
(
1
/
/
1
/
1
/
1
/
1
2
2
2
2
2
2
2
2
2 3 4 6 6
2
2
2
2
2
2
2
2
7
2
2
3
3
3
3
3
3
3
3
8
5
3
3
3
3
3
3
3
3
3 4
4
4
4
4
9
4
/
4
4
4
4
4
7
9
4
4
4
4
5
5
5
5
5
5
5
5
5
5
5
5
$
5
5
5
8
5
5
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
?
7
7
7 8 9
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
9
4
7
9
8
9
9
9
7
9
9
4
7
9
9
9
9
9
GR-MCK
PPO
6
0
6
0
0
6
Q
O
O 7
O
0
8
O
6
0
6
0
O
1
7
1
1
1
/
/
/
1
9
/
/
/
2
2
2
2
2
2
2
1
9
8
2
2
8
0
2
3
3
5
3
3
3
3
3
8
3
3
8
9
3
8
3
3
8
3 4
7 4 5 6
4
4
4
4
0
7
4
4
4 8
4
4
4
4
4
4
4
4
9
5
0
5
5
5
5
5
5
5
5
5
8
5
5
9
5
5 6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
/
7
7
7
2
7
7
7
3
7
7
9
7
7
7
7
7
7
7
7 9
3
5
8
8
8
8
3
9
8
8
θ
8
8
8
8
3
8
8
9
9
8
9
7
9
9
9
9
9
7
9
7
1
9
9
9
9
VQVAE
FSQ
6
6
0
0
0
6
0
O
0
O
0
0
0
0
0
0
0
0
0
0
1
/
2
1
1
1
1
/
/
1
1
1
1
7
/
1
3
2
2
3
2
2
0
2
2
5
2
2
2
2
3
2
2
2
2
7
3
3
3
3
3
3
3
3 4 5 6
3
3
3
3
3
3
3
3
3
3
3 4
4
4
4
4
4
4
4
4
4
7
4
4
4
4
4
4
4
4
5
5
3
5
5
5
15
5
5
5
5
5
3
9
5
5
5 6
5 6
6
6
6
6
6
3
6
6
6
G
6
6
6
6
6
6
7
7
7
7
7
1
7
7.
7
7
7
7
7
7
7
7
2
4
2
7 3
7 6
3
8
8
8
8
8
8
9
8
8
8
8
8
8
8
8
8
8
2
9
9
9
9
9
9 9
8
9
9
9
9
6
9
9
9
9
Figure 20:
Label-conditioned samples from downstream model using discrete flow matching.
31

---

Published as a conference paper at ICLR 2026
DAPS
Gumbel
3
7
3
3
5
6
7
3
3
9
8
7
4
6
4
2
7
0
3
3
4
0
8
8
8
2
8
7
4
5
3
4
7
5
6
3
0 2
7
6
9
9
8
0
3
8
2
3
3
7
4
4
5
4
1
3 7
3
3
3
1
2
9
8
8
8
5
3
3
4
x
5
7
3 7
4
6
[
4
0
3
2
3
2
9
8
8-
1
9
8 0 +
1
2
0
9
8 9 7 4 5
3
5
8
9 9
8
8
8
3
3
4
3
4
5
2
8
4 3 2 8
3
3
9
8
3
8
3
9
8
41
5
8
8
6
7
8
7
3
7
8
7
5
8
2
3
1
2
4
3
8
3
3
2
3
3
0
6
9
3
2
8
5
5
4
3
4
1
à
6
4
2
/
a
3
8
2
2
7
3
Q
7
3
S
3
GR-MCK
PPO
2
8
3
3
9
2
0
4
5
4
5
7
9
7
0
3
9
8
7
3
9
8
0
3
7
4
3
8
5
9
3
0
8
8
9
6
8
2
6 9
5
4
7
8
9
2
$
8
3
8
0
7
5
5
4
7
5
9
0
6
9
9
0
8
3
9
5
8
8
3
3
6
5
9
9
4
0
5
8
0
1
9
2
6
4
4
8
6
5
5 4
4
9
6
Q
9
3
8
4
4
6
2
5
4
7
7
5
5
8
6
6
3
3
8
5
0
0
3
9
7
5
2
?
2
1
2
8
6
4
7
8
9
6 0
8
3
7
0
8
8
2
7
4
8
3
2
2
4
2
6
2
7
4
7
0
7
6
VQVAE
FSQ
6
3
4
2
4
7
8
2
4
3 5 3 3
4.
7
5
30
8
3
7
22
7 3 3
3
2 &
2 σ 9
/
3
53
24
2
の 8
7
3.
5 8 6 2
3/
- 5 & 7 3
55
3
3
0 4
9
3 3
29
2
3
2
2
5 7
55
22
:-
3
7
3
8
K
35
52
X
7
5
3
%
3
4 -
'
$
1
22
4
3
2
7
2
2
4
32
7
?
$
3
Figure 21: Samples obtained by decoding random, uniform prior latent codes.
32

---

Published as a conference paper at ICLR 2026
DAPS
GR-MCK
VQ-VAE
FSQ
PPO
True
Figure 22: CIFAR validation reconstructions vs. ground truth.
33

---

Published as a conference paper at ICLR 2026
DAPS
GR-MCK
VQ-VAE
FSQ
True
Figure 23: ImageNet 256 validation reconstructions vs. ground truth.
34


Reviews:
[{"review_id": "CqvGnGOAIo", "reviewer": "Reviewer_tBAf", "summary": "The main challenge in training discrete VAEs lies in backpropagating through the discrete sampling process. In this work, the authors propose Discrete Autoencoding Policy Search (DAPS), a novel training framework for discrete VAEs inspired by policy search methods from reinforcement learning, particularly Relative Entropy Policy Search (REPS). The method exploits the gradient of a non-parametric encoder (which can be optimized in closed form but not directly sampled) to update the parametric encoder without relying on reparameterization, while also incorporating automatic step-size adaptation. This approach eliminates the need for differentiable soft relaxations or zero-order gradient estimators, which often lead to biased or high-variance gradients and unstable training.", "questions": "-\tWould this method still be beneficial when considering simpler models without an autoregressive structure, and a single vector as latent representation? One of the motivations of this approach is to improve memory efficiency with respect to approaches such as Gumbel-Softmax, does the same argument apply in simpler scenarios?\n-\tIn this context, is generation performed by sampling from the uniform prior assumed during training? Similar to VQ-VAE, does this method require training an autoregressive prior in a second stage? How is the generation process?\n-\tIn the appendix, the authors present label-conditioned samples. How is label-conditioned generation implemented? In the downstream robotics example, the high-level policy is also conditioned on language information, how is this text-conditioning realized in the DAPS generative process? \n-\tAccording to the authors, what are the main limitations of the method, and what are the primary directions for future research?\n\n\nI would be happy to increase my score if the authors adequately address the main weaknesses and questions.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Given the increasing interest in learning discrete representations in the literature, I believe exploring new methods that enable the effective training of discrete latent variable models is relevant. While connections with reinforcement learning have previously been explored to address the challenge of backpropagating through discrete sampling (most notably through methods like REINFORCE) to the best of my knowledge, this is the first work to apply particle-based policy search methods to train discrete VAEs. This approach effectively circumvents the reparameterization problem and enables more stable training procedures. The authors also show promising results on complex, high-dimensional datasets, additionally highlighting the potential of this method for broader applicability in downstream robotics tasks.", "weaknesses": "-\t**Method description.** Overall, I believe the description of the method could be improved. While some aspects may be intuitive for readers with a reinforcement learning background, they may appear opaque to others with expertise in deep generative models but not reinforcement learning (such as myself). For instance, why is the advantage function used instead of the reward function? In Equation (5), the loss for updating the encoder involves sampling from the parametric posterior, why does this not introduce the same backpropagation issue? I noticed that the training algorithm includes a stop-gradient operator, but I believe this should be explicitly discussed in the main text, since a central aspect of the method is preventing backpropagation through the discrete sampling process.\n\n-\t**Writing.** Additionally, I think the overall writing could be improved. In the introduction, several acronyms (e.g., ELBO) are not defined until later sections, and some variables referenced in the algorithm are never described in the text. Likewise, the metrics reported in the tables are not defined (e.g., log p, L2).\n\n-\t**Method.** While I find the framework interesting, it's unclear to me the ratio between performance gain and computational efficiency, as it is something that the authors do not address. They mention efficiency in terms of memory relative to Gumbel-Softmax-based methods, but not with respect to other baselines, nor in terms of training time. Also, as I stated before, I think key aspects of the method, such as the sampling procedure, remain unclear to me after reading the paper. I would appreciate if the authors can clarify it during the rebuttal.\n\n-\t**Evaluation.** I find the evaluation of the method limited. While the authors provide some reconstruction examples, showing that qualitatively their method performs comparably to state-of-the-art approaches, they do not include quantitative metrics to assess reconstruction quality, such as PSNR, comparing original and reconstructed samples. For generation, although FID scores are reported for some methods, the sample size used is not specified, which is important for comparison with state-of-the-art results. Additionally, it is unclear why these metrics could not be obtained for certain model configurations. Regarding the latent representations, the authors do not analyze their use. This is a notable limitation of VQ-VAEs that FSQ aims to address. It would be interesting to investigate how this alternative training framework, which allows for more stable training, impacts posterior inference. Furthermore, it would be valuable to study how the dimensionality of the latent space influences overall performance (could the improved training allow for more compact latent spaces?), especially since the latent representations appear to be high-dimensional. Also, I think the use of the ELBO as evaluation metric can be misleading, for example, a model with posterior collapse ELBO may still be high because reconstruction is good, but the latent code carries no information.", "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": "tkdJJ9dr7b", "reviewer": "Reviewer_WViT", "summary": "This paper introduces DAPS (Discrete Autoencoding via Policy Search), a novel method for training discrete variational autoencoders (VAEs) that leverages policy search techniques from reinforcement learning. The core idea is to frame the encoder update as a KL-regularized policy optimization problem, avoiding backpropagation through discrete sampling and enabling the use of autoregressive encoders. The method is evaluated on image (MNIST, CIFAR-10, ImageNet) and motion (LAFAN) datasets, demonstrating improved performance in reconstruction quality and downstream tasks like robot control.", "questions": "In Table 1, were the encoder architectures for all baseline methods (Gumbel, GR-MCK, FSQ, VQ-VAE) identical to the one used for DAPS? If not, how can we be sure that the performance gains are due to the proposed training algorithm and not a more powerful or better-suited architecture? Please clarify this in the experimental setup.\n\nPlease see also the above Weaknesses.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The integration of policy search principles (like REPS and trust-region adaptation via ESS) into discrete VAE training is a novel and potentially impactful contribution.\n\nThe paper provides extensive experiments, showing that DAPS can scale to challenging datasets like ImageNet and outperform strong baselines such as Gumbel-Softmax, GR-MCK, FSQ, and VQ-VAE on key metrics (e.g., FID, log-likelihood).\n\nThe method demonstrates stable training.", "weaknesses": "The paper's logic is often indirect and difficult to follow. The connection between the core objective (discrete autoencoding) and the RL framework is not smoothly bridged, making the \"big picture\" and the intuitive justification for the approach unclear. In fact, both sides can be directly bridged via reverse KL minimization.\n\nThe notation is frequently confusing and inconsistent. \n- A critical example is the description of the parametric encoder $q_{\\theta}(z|x)$. It is initially presented as a standard categorical distribution, but later described as an autoregressive model decomposed into a product of conditional (categorical) distributions. This conflict is deeply confusing. Furthermore, Figure 1 visually suggests a non-autoregressive structure (e.g., $z_1$, $z_2$, $z_3$ are not shown as inputs for $z_4$), creating a direct inconsistency between the text, the formula, and the diagram.\n- Eqs. (3)-(5) are presented with terminology and notations from reinforcement learning, despite Section 3 being mainly about discrete autoencoding. At least, the Advantage $A(z,x)$ should be more explicitly and clearly connected to the VAE objective.\n- The recognition loss in Line 9 of Algorithm 1 is different from Eq. (6), particularly considering the gradient with respect to the parameter $\\theta$.\n\nThe introduction of the constraint $D_{KL}[q(z|x) || q_{\\theta}(z|x)] \\le \\epsilon_{\\eta}$ in Eq. 4 actually changes the optimization from a reverse KL (in the standard ELBO) to a forward KL w.r.t. parameter $\\theta$. Reverse KL is mode-seeking, while forward KL is mode-covering. The paper provides no discussion or analysis of this fundamental shift, its implications for the latent representation, or its potential to cause conflicting behaviors during training.\n\nThe inclusion of a \"Goal-Conditioned Robot Control\" experiment feels disconnected from the primary narrative of discrete autoencoding. The paper fails to establish a clear, compelling link between this experiment and the core contribution discrete autoencoding.", "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": "RF2kNNr80B", "reviewer": "Reviewer_i7u6", "summary": "This paper introduces Discrete Autoencoding via Policy Search (DAPS), a novel training framework for discrete VAEs that avoids reparameterization or straight-through gradient estimation. DAPS formulates encoder learning as a KL-regularized policy search problem, and the parametric encoder is updated via weighted maximum likelihood using importance sampling. A trust-region parameter is adapted automatically using the effective sample size (ESS) to ensure stable training. Experiments on MNIST, CIFAR-10, ImageNet-256, and LAFAN motion data show that DAPS outperforms Gumbel-Softmax, GR-MCK, FSQ, and VQ-VAE in reconstruction quality and bit efficiency.", "questions": "1. Section 5 mentions that DAPS generates coherent and high-quality motion trajectories for LAFAN, but Figure 3 only shows a single frame without quantitative motion quality metrics. Could you provide quantitative metrics comparing DAPS against baselines on motion-specific evaluation criteria, and clarify how 'coherence' was assessed (e.g., via human evaluation or physical plausibility checks)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. DAPS is a principled training framework that reformulates discrete VAE encoder optimization as a KL-regularized policy search problem, enabling stable and scalable training without reparameterization. The methods leverage modern RL techniques that have been underexplored in discrete VAEs.\n2. DAPS demonstrates state-of-the-art reconstruction quality on ImageNet-256, and outperforms GR-MCK and Gumbel-Softmax in both log-likelihood and FID across MNIST, CIFAR-10, and ImageNet. It scales efficiently to sequence lengths of 1,024.\n3. The step size of DAPS adapts automatically via ESS, eliminating the need for manual tuning of the trust-region parameter.", "weaknesses": "1. The baselines involved in the experiments lack the most recent work in 2024/2025; therefore, it is not clear how DAPS compares to recent SOTA methods.\n2. Section 3 states that DAPS avoids computationally expensive and high-variance backpropagation through time by using weighted maximum likelihood. However, the paper doesn't empirically quantify the computational savings compared to Gumbel-Softmax or GR-MCK.\n3. The abstract claims a 20% improvement on FID  for ImageNet 256. It would be beneficial to have a hypothesis test to determine whether this improvement is significant.", "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)