Multi-Rate VAE: Train Once, Get the Full Rate-Distortion Curve
Authors: Juhan Bae
Michael R. Zhang
Michael Ruan
Eric Wang
So Hasegawa
Jimmy Ba
Roger Baker Grosse
Venue: ICLR 2023 notable top 5% (2023)
Keywords: Variational Autoencoders
VAEs
Hypernetworks
Response Functions
Hyperparameter Tuning

Abstract: Variational autoencoders (VAEs) are powerful tools for learning latent representations of data used in a wide range of applications. In practice, VAEs usually require multiple training rounds to choose the amount of information the latent variable should retain. This trade-off between the reconstruction error (distortion) and the KL divergence (rate) is typically parameterized by a hyperparameter $\beta$. In this paper, we introduce Multi-Rate VAE (MR-VAE), a computationally efficient framework for learning optimal parameters corresponding to various $\beta$ in a single training run. The key idea is to explicitly formulate a response function using hypernetworks that maps $\beta$ to the optimal parameters. MR-VAEs construct a compact response hypernetwork where the pre-activations are conditionally gated based on $\beta$. We justify the proposed architecture by analyzing linear VAEs and showing that it can represent response functions exactly for linear VAEs. With the learned hypernetwork, MR-VAEs can construct the rate-distortion curve without additional training and can be deployed with significantly less hyperparameter tuning. Empirically, our approach is competitive and often exceeds the performance of multiple $\beta$-VAEs training with minimal computation and memory overheads.

Published as a conference paper at ICLR 2023
MULTI-RATE VAE: TRAIN ONCE, GET THE FULL
CURVE
Juhan Bae! 12, Michael R. Zhang' :2, Michael Ruan', Eric Wang', So Hasegawa',
Jimmy Bal.2, 2 Roger Grosse 1,2,4
University of Toronto, Vector Institute, Fujitsu Limited, 4 *Anthropic
ABSTRACT
Variational autoencoders (VAEs) are powerful tools for learning latent represen- tations of data used in wide range of applications. In practice, VAFs usually require multiple training rounds to choose the amount of information the latent variable should retain. This trade-off' between the reconstruction error (distortion) and the KL. divergence (rate) is typically parameterized by a hyperparameter B. In this paper, we introduce Multi-Rate VAE (AE), a computationally efficient framework for learning optimal parameters corresponding to various 3 in single training run. The key idea is to explicitly formulate response function using hypernetworks that maps 3 to the optimal parameters. MR- VAEs construct a compact response hypernetwork where the pre-activations are conditionally gated based on B. We justify the proposed architecture by analyzing linear VAEs and showing that it can represent response functions exactly for linear VAEs. With the learned hypernetwork, MR-VAEs can construct the rate- aate-distottion curve without additional training and can be deployed with significantly less hyperparameter tun- ing. Empirically, our approach is competitive and often exceeds the performance of multiple 3-VAEs training with minimal computation and memory overheads.
INTRODUCTION
Deep latent variable models sample latent factors from a prior distribution and convert them to realistic data points using neural networks. However, computing the model parameters via maximum likelihood estimation is challenging due to the need to marginalize the latent factors, which is
intractable. Variational Autoencoders (VAEs) (Kingma & Welling 2013; Rezende et al. formulate a tractable lower bound for the log-likelihood and enable optimization of deep variable models by reparameterization of the Evidence Lower Bound (ELBO) (Jordan et al. VAEs have been applied in many different contexts, including text generation (Bowman et al. data augmentation generation (Norouzi et al. (2220), anomaly detection (An & Cho 2015 Park et al. 2022), future frame prediction (Castrejon et al. 2019), image segmentation (Kohl et al. 2018). and music generation (Roberts et al. 2018).
2014)
latent
1999).
2015
In practice. VAEs are typically trained with the 3- 3-VAE objective (Higgins et al. 2016) which balances the reconstruction error (distortion) and the KL divergence term (rate):
L 0) Exx(x) [E
[x)ool log Pe [x] + B (pp(() [DxL (z)()))) p(z)))
Distortion
Rate
where Pa (x z)) models the process that generates the data x given the latent variable (the "decoder") and 9s (z) x) is the variational distribution (the "encoder"), parameterized by 0 E RT and Φ E RP. respectively. Here, p(z) is prior on the latent variables, (x() is the data distribution, and /3 > 0 is the weight on the KI. term that trades off between rate and distortion.
On the one hand, models with low distortions can reconstruct data points with high quality but may generate unrealistic data points due to large discrepancies between variational distributions and priors (Alemi et al. 2018). On the other hand, models with low rates have variational distributions close to the prior but may not have encoded enough useful information to reconstruct the data. Hence,
"Correspondence jbaelcs................oo....... toronto.ed.

---

Published as a conference paper at ICLR 2023
AaleeDDitottion Curve
Training
RD
Methods
Parameters
Time (8)
Curve
150
VAR:
B-VAS
Standard VAE
18.81 x 10
4,832
X
Distortion
G-VAIS IKL Aroaaing
3-VAE
18.81 x 10
4,836
3-VAEs (+8 runs)
15.04 x 10
38,492
30-
3-VAEs (+ KL Annealing)
15.04 x 10a
38,640
MR-VAES (ours)
18.84 x 10*
5,040
100
(00
CA1
100
150
Figure 1: 3-VAEs require multiple runs of training with different KI. weights 3 to visualize parts the rate-distortion curve (Pareto frontier). Our proposed Multi-Rate VAEs (MR-VAFs) can learn the full continuous rate-distoris distortion curve in single run with small memory and computational overhead.
the KL weight 3 plays an important role in training VAEs and requires careful tuning for various applications (Kohl et al. 2018 Castrejon et al., 2019) Pong et al. 2019). The KL weighting term
also has close connection to disentanglement quality (Higgins et al.] 2016; (Brrgsss et al. 2018
Nakagawa et al.] 2021), generalization ability (Kumar & Poole, 2020 Bozkurt et al. 2021 data
compression (Zhou al. 2018 Huang et al. (020), and posterior collapse (Lucas et al. 2019 Dai
et al. 2020
Wang & Ziyin
2022).
By training multiple VAEs with different values of 33 we obtain different points rate- distortion curve (Pareto frontier) from information theory (Alemi et al., 2018). Unfortunately. as rate-distortion curves depend both the dataset and architecture, practitioners generally need to tune 3 for each individual task. In this work, we introduce a modified VAE framework that does require hyperparameter tuning on and can learn multiple VAEs with different rates in single training run. Hence, we call our approach Multi-Rat Rate VAE (MR-VAL VAE).
We first formulate response functions (3) and 0 (3) Giibonn al. 1992) which map the KL weight 3 to the optimal encoder and decoder parameters trained with such 3. Next, we explicitly construct response functions 2 and (3)) using hypernetworks (Ha et al. 2016), where ys R denotes hypernetwork parameters. Unlike the original VAE framework, which requires retraining the network to find optimal parameters for some particular B, response hypernetworks can directly learn this mapping and do not require further retraining.
While there is a lot of freedom in designing the response hypernetwork, we propose a hypernetwork parameterization that is memory and cost-efficient yet flexible enough to represent the optimal parameters over a wide range of KL weights. Specifically, in each layer of a VAE, our MR-VAD architecture applies an affine transformation to log 8 and uses it to scale the pre-activation. We justify the proposed architecture by analyzing linear VAEs and showing that the MR VAE architecture can represent the response functions on this simplified model. We further propose modified objective analogous to Self-Tuning Networks (MacKay et al., 2019: Bae & Grosse 2020) to optimize response hypernetworks instead of the standard encoder and decoder parameters.
Empirically, we trained MR- VAES to learn rate-distortion curves for image and text reconstruction tasks over a wide range of architectures. Across all tasks and architectures, MR-V MR-VAEs found competitive or even improved rate-distortion curves compared to the baseline method of retraining the network multiple times with different KL weights. We show a comparison between V-VAE (with and without KL annealing (Bowman et al. 2013)) and MR- MR-VAE with ResNet- encoders and decoders (He et al.] 2016) in Figure I MR- VAES can learn multiple optimal parameters corresponding to various KI. weights in a single training run. Moreover, MR-VAF3 do not require KI. weight schedules and can be deployed without significant hyperparameter tuning.
Our framework is general and can be extended to various existing VAE models. We demonstrate this flexibility by applying MR-VAES to 3-TCVAES (Chen et al. 2018), where we trade-off the reconstruction error and total correlation instead of the reconstruction error and rate. We show that TA-VAEs can be used to evaluate the disentanglement quality over a wide range of 33 values without having to train 3-TCAES multiple times.
2 BACKGROUND
2.1 (8-) VARIATIONAL. AUTOENCODERS
Variational Autoencoders jointly optimize encoder parameters & and decoder parameters 0 to min- imize the 3-VAE objective defined in Eqn. While the standard EI BLO sets the KI. weight to 1,
2

---

Published as a conference paper at ICLR 2023
VAEs are commonly trained with varying KL weight 3. As the KL weight 3 approaches 0, the VAE objective resembles the deterministic autoencoder, putting more emphasis on minimizing the reconstruction error. When the KL weight is large, the objective prioritizes minimizing the KL divergence term, and the variational distribution may collapse to the prior.
It is possible interpret VAES from an information-ttheoretii framework ÇAlemi et al. 2018), where the 3-VAE objective is a special case of the information-bottleneck (IB) method (Tishby et al. 2000 Alemi al. 2016). In this perspective, the decoder po (x|z) serves lower bound for the mutual information 1,(x; xx z) and the rate upper bounds the mutual information. The relationship can be summarized as follows:
H D S (xxx2) S R
where D is the distortion, H the data entropy that measures the complexity of the dataset, and R is the rate. Note that we adopt terminologies "rate" and "distortion" from rate-distortion theory (Thomas Joy 2006; Alemi et al., [2018] Bozkurt et al., 2021). which aims to minimize the rate under some constraint on distortion. The 3- B-E objective in Eqn. | can be seen as a Lagrangian relaxation of the
rate-distortion objective:
Co D
In this view, training multiple VAEs with different values of B corresponds to different points on the rate-distortion curve and distinguishes models that do not utilize latent variables and models that make large use of latent variables (Phuong al. 2018). We refer readers Alemi al. (2018) for more detailed discussion of the information- thheoretic framework.
2.2
LINEAR VAES
The linear VAE a simple model where both the encoder and decoder are constrained to be afine transformations. More formally, we consider the following problem from Lucas et al. (2019):
(x M/Dz 1,,"I)
l (2) x)) NV((x - -4)x 1),C),
where C is a diagonal covariance matrix that is shared across all data points, E is the encoder weight, and D the decoder weight.
While the analytic solutions for deep latent models typically do not exist, the linear VAB has analytic solutions for optimal parameters and allows
reason about various phenomena in training
VAEs. For example, Dai et al. (2018) show the relationship of linear VAE, probabilistic PCA
(pPCA) (Tipping & Bishop 1999,, and robust PCA (Candès et al. [011] Chandrasekaran et al. 2011 f) and analyze the local minima smoothing effects of VAEs. Similarly, Lucas et al. (2019) and Wang &
Ziyin 2022) use linear VAFs to investigate the cause of the posterior collapse (Razavi et al. 2019).
2.3 RESPONSE FUNCTIONS
Response (rational reaction) functions (Gibbons et al. 1992 Lorraine & Duvenaud 2018 in neural networks map a set of hyperparameters to the optimal parameters trained with such hyperparameters. In the case of 3-VAEs, response functions map the KL weight 3 to optimal encoder and decoder parameters that minimize the 3- V-V objective:
Φ² (3) = arg min L3 0)
DERP
0 (B) arg min ( 0).
#ERm
Approximation of response functions has various applications in machine learning, including hyper- parameter optimization (Lorraine & Duvenaud 2018), game theory (Fiez et al., 2019 Wang et al.. 2019) and influence estimation (Bae et al. 2022).
3 METHODS
In this section, we introduce Multi-Rate VAEs (MR VAEs), an approach for directly modeling VAE response functions with hypernetworks. We first formulate memory and cost-efficient response
3

---

Published as a conference paper at ICLR 2023
Layer
Log
Affine Transform
Figure 2: MR- VAE architecture scales the pre-activations for each layer of base VAE. This scaling term is generated with an affine transformation on the log KL. weight followed by an activation function.
hypernetworks (3) and O+ (3) and justify the proposed parameterization by showing that they recover exact response functions for linear VAEs. Then, we propose a modified training objective for MR-VAEL and formally describe algorithm for learning the entire rate-distortion curve in single training run.
3.1 RESPONSE HYPERNETWORKS FOR VAES
We propose explicitly modeling response functions with hypernetworks to construct the rate-distortion curve, opposed training multiple VAEs with dif- ferent KL weights B. The formulation of response hypernetworks, shown Figure 3 is advantageous as we can infer the optimal parameters for some / by simply changing the inputs to the hypernetwork. As response functions map scalar to high-dimensions parameter space Rm-P, there are many possible archi- tectural designs for such response hypemetworks.
B
10
VAEs
Response Function
B= 0.00
Figure Instead of training several VAEs for each desired KL weight B. MR-VAE. learn the response functions with hypernetwork in a sin- gle training
Our Multi-Rate VAEs (MR- VAEs) construct a com- pact approximation of response functions by scaling each layer's weights and bias by function of the KL weight. More formally, consider the i-th layer of VAE, whose weight and bias are expressed as W()) m+++ xm and (6) Rm respectively. We directly model response functions with hypernetworks v E R by parameterizing each weight and bias as follows
w (8) (6)
w hyper log( +
(e)
hhpeer
mw
base
b Φ (B) o (i) hyper log( b Phyper
(i)
buse:
where O and Crow indicate elementwise multiplication and row- row-wise rescaling. Here, w Whyper b b Rm are vectors that are used in applying an affine transformation to the log KL weight. We further define the elementwise activation function (+)) (.)
if i-th layer belongs encoder 0.
(x)=
I+e (ReLUU( -cxp(x))) if i-th layer belongs to decoder Φb
This specific choice of activation functions recovers the exact response function for linear VAEs. which we elaborate on in Section 3.2] We further note that the above activation functions scale the base weight and bias between and 1. In Appendix [ show that the choice of activation function does not significantly change the performance of MR-VARR VAEs and we can use sigmoid gatings for both encoders and decoders in practice.
Observe that the hypernetwork parameterization in Eqn. 6 is both memory and cost-efficient: needs 2m;+1 additional parameters represent the weight and bias and requires 2 additional elementwise
We show the hyperntwork parameterization for convolution layers in Appendix
4

---

Published as a conference paper at ICLR 2023
multiplications per layer in the forward pass. Given activations a R from the previous layer, the resulting pre-activations s(1) can alternatively be expressed as:
s()) w (3)a
(8)
W hyper log(
base a (1
+
b
(9)
hyper
base
The architecture for MR- VAEs is illustrated in Figure This response hypernetwork is straightfor- ward to implement in popular deep learning frameworks (e.g.. PyTorch (Paszke et al. 2019) and Jax (Bradbury et al. 2018) by replacing existing modules with pre-activations gated modules. We
provide sample PyTorch code in Appendix B.2
3.2 EXACT RESPONSE FUNCTIONS FOR LINEAR VAES
Here, we justify the proposed hypernetwokk parameterization introduced in Section 3.1 by analyzing the exact response functions for linear VAEs (Lucas et al. 2019). We show that MR-Ves VARS can precisely express the response functions for this simplified model. Consider the following problem
(analogous to the setup introduced in Section
2.2
(x(x (D (1) + μ. 1))
(10)
4 (z (E E n). c())))
Compared to the linear VAE setup in Eqn. we represent the encoder weight, diagonal covariance, and decoder weight the product of two matrices for compatibility with our framework, where E(1), C()), and (L) are all square matrices. This decomposition can be understood as constructing two-layer linear networks for both the encoder and decoder.
We apply the MR-VAE formulation (shown in Eqn. 6 to all parameters and construct response hypernetwork for each encoder weight, diagonal covariance, and decoder weight. The following theorem shows that the MR-VAE architecture is capable of representing the response functions for linear VAEs.
Theorem 1. Consider the linear VAE model introduced in Eqn. 10 Let the encoder weight, diagonal (
covariance matrix, and decoder weight be parameterized Eap (8) E ( (1) (3), (B)
(2)) (B) (B)), respectively. Then, there exist hypernetwork C (3), and ( D vp parameters xp* E Ih such that E (3), C.. (3), and D (3) are the exact response functions for linear VAEs.
The proof is shown Appendix A.2 Note that we use a separate activation function (described in Eqn. 7 for encoder and decoder parameters exactly represent the response functions. Motivated by the analysis presented this section, IRRV use the same hypernetwor parameterization for deep VAEs.
3.3 TRAINING OBJECTIVE FOR RESPONSE HYPERNETWORE
To learn the optimal parameters for various ranges of the KL weight 3, we propose to use an objective analogous to the Self-Tuning Networks (STN) objective (Lorraine & Duvenaud, 2018 MacKay
et al. 2019 Bae & Grosse 2020). The hypernetwork parameters vp are optimized to minimize the Tollowing objective:
4/5 =arg min (( = arg min E (oo((()(o(()) (Fexp(n) (& (cxp(n)). 0 (cxp(n)))) (11)
vER*
where U[log(u), log(b) is uniform distribution with range log(a) and log(b). The proposed objective encourages response hypernetworks to learn optimal parameters for all KL weights the range between and 6. Unlike the STN objective, sample the hypermetoork inputs from uniform distribution with a fixed range (instead of Gaussian distribution with a learnable covariance matrix) interested in learning the global response function (instead of the local response function).
3.4 TRAINING ALGORITHM
The complete algorithm for MR- VAE summarized in Alg. D When training MR- VAEs, we approximate the objective in Eqn. by sampling 7) in each gradient descent iteration.
5

---

Published as a conference paper at ICLR 2023
Algorithm 1 Multi-Rate Variational Autoencoders (MR-VAEs)
Require: & (Hypemetwork parameters), (learning rate), (a, b) (sample range)
while not converged do
B Dirain (llgg (), log(b))
Sample mini-batch
b Sample inputs to the hypemeeworr
((b) = (xpp(n) @ ( (exp(()));
Compute the MR-VAE objective
v < 7/ ((b) end while
Update hypernetwork parameters
Normalizations. As is widely considered beneficial for optimization to normalize the inputs and activations (LeCun et al., 2012] Montavon & Mäller 2012 loffe & Szegedy) 2013 [ae & Grosse 2020), we standardize the inputs to the hypernetworks (after log transform) to have a mean of 0 and standard deviation of 1. Since the inputs are sampled from fixed distribution 2/ (log(a), log(b)), can apply fixed transformation based on the sample range [a, bj specified before training.
Memory and Computation Costs. As detailed in Section 3.1 MR-VAE architectures introduce additional vectors for each layer of the base neural network. During the forward pass, MR- VAEs require additional elementwise operations each layer. Across all our image reconstruction experiments, MR-VABE require at most 5% increase in parameters and wall-clock time compared to 3-VAEs.
Hyperparameters. While MR-VAE does not require hyperparameter tuning on the KI. weight, it needs two hyperparameters a and 6 that define the sample range for the KL weight. However, we show that MR- VAEs are robust to the choices of these hyperparameters in Section 5.3 and we use fixed value a = 0.01 and b = 10 for our image and text reconstruction experiments.
RELATED WORKS
Rate-distortion with VAFs. Alemi et al. (2018) and Huang al. (020) advocate for reporting points along the rate-distortion curve rather than just the ELBO objective to better characterize the representation quality of the latent code. For instance, a powerful decoder can ignore the latent code and still obtain a high marginal likelihood as observed by |Bomman et al. (2015) and Chen et al. (201)). Alemi al. (2018) show that this problem can be alleviated by choosing < 1, but this approach still requires training VAE for each desired information theoretic trade-off. YYngg et al. (2020) study the use of modulating activations autoencoders for image compression, but unlike MR-VAFs, they do not learn generative model and require learning a separate network for each target rate. Concurrently, Collins et al. (2022) directly provide 3 as inputs to VAEs to Iearn the rate-distortion trade-off for particle physics discoveries.
Calibrated Decoders. Instead of using a fixed decoder variance, calibrated decoders (Lucas et al. 2019: Rybkin et al.) (200)) update the decoder variance during training and do not require tuning the KL weight B. In the case where Gaussian decoders are used, the decoder variance o² is equivalent to setting 3 = 2o2. Calibrated decoders may be desirable for applications such density modeling as it is trained with the standard ELBO objective (3 = 1). However, many applications of VAEs involve directly tuning 3. In addition, calibrated decoders cannot construct rate-distortion curve and may not be appealing when the aim is to get better disentanglement. MR VAEs can be directly applied in both of these settings.
Multiplicative Interactions. Multiplicative interactions have long history in neural networks. Earlier works include LSTM cells (Hochreiter & Schmidhuber] 1997). which use gating mechanisms to model long-term memories, and Mixture of Experts (Jacobs et al. 1991), which selects learners with a softmax gating. The FiLM network (Perez et al. 2018) and conditional batch normalization (Perez ett al.) 2017 Brock al.] 2018) further extend the gating mechanism by scaling and shifting the pre-activations, Dayakumar et al. (2020) analyze the expressive power of multiplicative interactions within network. Our work multiplicative interactions to learn the response function for 3-VAFs.

---

Published as a conference paper at ICLR 2023
Response Hypernetworks. Our use of hypernetworks for learning best response functions is related to Self-Tuming Networks (STNs) (MacKay et al. 2019; Bae & Grosse 2020) and Pareto Front Hypernetworks (PFHs) (Navon et al. 2020). Lorraine & Duvenaud (20naaadd(088)) (2018) first constructed response hypernetworks for bilevel optimization (Colson et al., 2007) and optimized both parameters and hyperparameters in single training run. To support high dimensional hypernetwork inputs, STNs structured response hypernetworks similar to MR-V,ES, whereas PFHs use chunking that iteratively generates network parameters. In contrast to these works, the architecture of MR- RARS is specifically designed for learning 3-- B-VAE's response functions and requires significantly less computation and memory.
EXPERIMENTS
Our experiments investigate the following questions: (1) Can MR-VAES learn the optimal response functions for linear VAËs? (2) Does our method scale to modern-size 3-VABs? (3) How sensitive are VRES to hyperparameters? (4) Can MR-VAEs be applied to other VAE models such as 3-TCVAEs (Chen al. 2018)?
We trained MR-VAEs to learn VAES with multiple rates in a single training run on image and text reconstruction tasks using several network architectures, including convolutional-basse VAEs, autoregressive VAEs, and hierarchical VAEs. In addition, we applied MR- VAEs to the 3-TCVAE objective investigate if our method can be extended to other models and help find VAEs with improved disentangled representations. Experiment details and additional experiments, including ablation studies, are provided Appendix (aaendixCanndd Appendix D respectively.
5.1 How DO URRVAES PERFORM ON LINEAR VAES?
We first considered two-layer linear VAE models (introduced in Section 3.2) to validate our theoreti- cal findings. We trained VAEV on the MNIST dataset (Deng) 2012) by sampling the KL weighting term ,3 from 0.01 to 10. Then, we trained 10 sepa- rate two-layer linear B- VAEs, each with different ,
MNIST (Linear VAE)
3-VAE
800
MR VAE
Distortion
40-
values. The rate-distortion curve for MR- VAEs and 3-VAEs are shown in Figure 4 We observed that the rate-distortion curve generated with MR-VAES closely aligns with the rate-distortion curve generated with individual retraining. Furthermore, by explicitly formulating response functions, MR-VAES does not require retraining VAE for each |8.
Rate
Figure comparison of MR- VAEs 8- VAEs on linear VAFs. MR- VAFs learn the optimal rate- distortion curve in single training run.
5.2 CAN VR-V SCALE TO MODERN-SIZE ARCHITECTURES?
We trained convolution and Res architectures on binary static MNIST (Larcchllle & Murray 201T), Omniglot (Lake al. 2015), CIFAR-I0 (Krizhevsky et 2009), SVHN (Netzer et al. [001 and CelebA64 (Liu et al. 2015) datasets, following the experimental set-up from Chadebec et al. (2022). We also trained N VAES on binary static MNIST, Omniglot, and CelebA datasets using the same experimental set-up from Vahdat & Kautz 2020)) Lastly, we trained autoregressive LSTM VAEs on the Yahoo dataset with the set-up from He et al. (2019).
We sampled the KL weighting term from 0.01 to 10 for both MR-VAES and B-VAEs. Note that the training for 3-VAES was repeated 10 times (4 times for NVAEs) with log uniformly spaced / to estimate rate-distortion curve. We compared MR-VAES with individually trained 3-VAEs and show test rate-distortion curves in Figure 5 Note that traditional VAEs typically require scheduling the KL weight to avoid posterior collapse (Bowman et al. 2015 Fu et al. 2019: Lucas et al.) 2019) and we train 3- VAES with both a constant and KL. annealing schedule. Across all tasks, MR- VAES achieved competitive rate- dittoristotioon curves with the independently trained VAEs and even improved performance on the MNIST, Omniglot, and CelebA datasets. Note that we focus on the visualization of the rate-distortion curve advocated for by Alemi al. (2016); Huang et al. (2020).
7

---

Published as a conference paper at ICLR 2023
MNEST
Outatiglor
CIPAR-ID
VAEE
MR VAE Tnelltt
Distotton
Oinoltes
Distetion
Disiotio
70
30
Rate
Rate
Kote
SVHN
CelchA
Yaluou
200
ANAC
Distortion
41
Distontoo
200
Disuntion
500
800
100
50
0-
100
20
10
DM0
Role
Rule
Rite
MIJST (NVAL)
Ontiigtoo (NVAI)
Caleba (NVVA))
125
VAR MENAA
100
Desutti (ed))
Divortion
Divanison
70-
50
8
50
.00
125
1.50
55
50
70
.00
135
130
Ratc
Rate
Raio
Figure 5: A comparison of MR-VAFs and 8-VAEs (with and without KI. annealing). 3-VAEs trained multiple times on several KL weighting to construct rale-distortion curve. Across all tasks, MR-VAEs achieve a competitive rate-distortion curve with B-VAEs on a single run of training.
Omniglot
SVEN
CelehA
Aufive Uniia (A1))
A-VAE
Audie USiis (AL))
Audive Ebita (AZ)
ore o
10
SVIN
ColebA
VAT NAR
11777
30
10 107
Figure 6: A comparison of independently trained 3-VAEs (with KI. annealing) and MR-VAFS proxy metrics such (top) active units (AU) and (bottom) Fréchet inception distance (FID). Note that lower FID better.
As MR- VAEs directly learn the optimal parameters corresponding to some in a single training run, can use them analyze the dependency between the KL weight and various proxy metrics with a single training run. To demonstrate this capability of MR-VAES we compute the Fréchet Inception Distance (FID) (Heusel al. 2017) for natural images and active units (AU) (Burda et al.) (0166) in Figure 6 Notably, by using a shared structure for various KI. weights, we observed that MR- MR-VAES are more resistant dead latent variables with high KL weights and generate more realistic images
with low KL weights.
We further show samples from both 8-VAES and MR-VAEs on SVHN and CelebA datasets which used -eaNet--bsse encoders and decoders in Figure 7 One interesting property is that MR- VAEs have consistency among the generated images for the same sampled latent since base VAE weights are shared for different B. A practitioner can select the desired model for their task or trade-off between compression and generation quality by feeding different Bs into the hypernetwork at inference time.
5.3 How SENSITIVE ARE MR- -AES TO HYPERPARAMETERS.
While the MR- VAE eliminates the need to train the network multiple times with different KL weights and KI. schedule, it introduces two hyperparameters and that determine the range to sample
8

---

Published as a conference paper at ICLR 2023
/-Es
VIR-VAES
8 10
8 60
8
0 0.1
B
78 AL 2 270
78 052
0.01
Figure 7: Generative samples for 3-VAES and MR-VAEs CelcbA and SVHN datasets. We use the same latent variables sampled from the prior for all columns. For MR-VAEs, there consisteney among the generated images each column since intermediate weights are shared for different 3. While higher values result in worse reconstruction loss for both models, MR- VAEs produce images that are of higher quality.
160
0.00
150
10
(-TEVAL
SER-WA WAB
Disortio
125
0.1
Disoction
1:2
Messi
25
75
1001 1.00
Figure 8: (left and middle) An ablation studying the effect of the sample range for MR- VAEs. For each figure, we fix one value (b = 10 or a = 0.01) and change the other. (right) A comparison of independently trained 3-TCVAEs and MR-VAEs the Mutual Information Gap (MIG) metric (the higher value is better)
the KI. weight. Here, we show that VR-SAEs are insensitive and robust to the choices of the hyperparameters and can be fixed through various applications. We trained ResNet encoders and decoders on the Omniglot dataset with the same configuration rate-distortion curves Figure 8 the left, we fixed = 10 and changed the sample range a in the set {0.001, 0.01, 0.1, }} and in the middle, set (10, 1, 0.1). While the length of the rate-distortion curves differs with different sample ranges, observed that the area under the curve is similar across all configurations.
Section 5.2 and show the test
fixed 0.01 and modified the sample range b in the
5.4 CAN MR- VAES BE EXTENDED TO OTHER MODELS?
To demonstrate the applicability of MR-VERVARR VAES other VAE models, trained MR-V VABs on the 8- TCVAE objective. The weight in 8- 2-TCVAE balances the reconstruction error and the total correlation instead of the reconstruction error and the KL divergence term.
We trained VAES composed of MLP encoders and decoders on the dSprites dataset, following the set-up from Chen et al. (2018). Since we are interested the disentanglement ability of the model, we sampled the weighting / between 0.1 and 10. We compare 3-TCVAES and MR-VA RA-VAEs (trained with the 3-TCVAE objective) by examining their performance on the Mutual Information Gap (MIG) in Figure 8 We observed that MR-VAE: are more robust to the choice of the weighting term and achieve competitive final results as the baseline 3-TCVABs without having to train multiple models.
6 CONCLUSION
In this work, we introduce Multi-Rate VAE (MR-VAE), an approach for learning multiple VAEs with different KL weights in a single training run. Our method eliminates the need for extensive hyperparameter tuning on the KL weight while only introducing a small memory and computational overhead. The key idea is to directly learn map from KL weights to optimal encoder and decoder parameters using a hypernetwork. On various tasks and architectures, we show that MR- VAEs generate competitive multiple VAEs with various values of 3. Moreover, our framework is general and straightforward to extend to various existing VAE architectures.
even better rate-distortion curves compared to the baseline of retraining

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We would like to thank James Lucas, Rob Brekelmans, Xuchan Bao, Nikita Dhawan, Cem Anil, and Silviu Pitis for their valuable feedback on this paper. We would also like to thank Sicong Huang and many other colleagues for their helpful discussions throughout this research. Resources used in this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute www.wectorinstitute.. partners).
REFERENCES
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken elbo. In International Conference on Machine Learning, pp. 159-168. PMLR, 2018.
Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint 1rXiv:1612.0000 2016.
Jinwon An and Sungzoon Cho. Variational autoencoder based anomaly detection using reconstruction probability. Special Lecture on IE, 2(1):1-18, 2015.
Juhan Bae and Roger B Grosse. Delta-stn: Efficient bilevel optimization for neural networks using structured response jacobians. Advances in Neural Information Processing Systems, 33: 21725-21737, 2020.
Juhan Bae, Nathan Hoyen Ng, Alston Lo, Marzyeh Ghassemi, and Roger Baker Grosse. If influence functions are the answer, then what the question? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL httpp:://peereeiiieeeee
net fortmmormmm i iddiidion hzbguh9zk
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from continuous space. arXiv preprint arXiv:1511.06344 2015.
Alican Bozkurt, Babak Esmaeili, Jean- nan-BBptiste Tristan, Dana Brooks, Jennifer Dy, and Jan-Willem van de Meent. Retelregulatizzttion and generalization in variational autoencoders. In International Conference on Artificial Intelligence and Statistics, pp. 3880-3888. PMLR, 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: http://ggtttt.tt/ wtt. www.///ooo/e//pooocooooo
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1096, 2018.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. 2016.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in bela-vae. arXiv preprint arXiv:1804.03599 2018
Emmanuel J Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM (JACM), 58(3):1-37, 2011.
Lluis Castrejon, Nicolas Ballas, and Aaron Courville, Improved conditional vrnns for video prediction. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 7608 -6117 2019.
Clément Chadebec. Louis J. Vincent, and Stéphanic Allassonničre. Pythae: Unifying generative autocncoders in python benchmarking case. arXiv preprint arXiv:2206.08833 2022. URL
httts:///axxiv org/ abss 2206 000090000000
Venkat Chandrasekaran, Sujay Sanghavi, Pablo A Parrilo, and Alan S Willsky. Rank-sparsity incoherence for matrix decomposition. SIAM Journal on Optimization, 21(2):572-596 2011.
10

---

Published as a conference paper at ICLR 2023
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disen- tanglement in variational autoencoders. Advances in Neural Information Processing Systems, 2018.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint 1rXiv:661..211 2016.
Jack H Collins, Yifeng Huang, Simon Knapen, Benjamin Nachman, and Daniel Whiteson. Machine- learning compression for particle physics discoveries. arXiv preprint arXiv:2210.1148)) 2022.
Benoît Colson, Patrice Marcotte, and Gilles Savard. An overview bilevel optimization. Annals of operations research, 153(1):235-256, 2007.
Bin Dai, Yu Wang. John Aston, Gang Hua, and David Wipf. Connections with robust pea and the role of emergent sparsity in variational autoencoder models. The Journal of Machine Learning Research, 19(1):1573-1144 2018.
Bin Dai, Ziyu Wang. and David Wipf. The usual suspects? reassessing blame for posterior collapse. International Conference on Machine Learning, pp. 2313-2322. PMLR, 2020.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6): (9(6):441-144 2012.
Tanner Fiez, Benjamin Chasnov, and Lillian Ratliff. Convergence of learning dynamics in stackel- berg games. arXiv preprint arXiv:1906.012//1 2019
Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. Cyclical annealing schedule: A simple approach to mitigating kl vanishing. arXiv preprint 10X10:1903.10145, 2019.
Robert Gibbons et al. A primer in game theory. 1992.
David Ha, Andrew Dai, and Quoe V Le. Hypernetworks. arXiv preprint 1609v:1609.09/1 09106, 2016.
Junxian He, Daniel Spokoyny, Graham Neubig, and Taylor Berg-Kirkpatrics Lagging inference networks and posterior collapse in variational autoencoders. arXiv preprint urXiv:1901 1901004400,, 2019.
Kaiming Hc. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
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. 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1780, 1997.
Zhiting Hu, Haoran Shi, Bowen Tan, Wentao Wang, Zichao Yang, Tiancheng Zhao, Junxian He, Lianhui Qin, Wang, et al. Texar: A modularized, versatile, and extensible toolkit for text generation. In ACL 2019, System Demonstrations, 2019.
Sicong Huang, Alireza Makhzani, Yanshuai Cao, and Roger Grosse. Evaluating lossy compression rates of deep generative models. In International Conference on Machine Learning, pp. 4444-4454,
PMLR, 2020.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456 PMLR, 2015.
11

---

Published as a conference paper at ICLR 2023
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79 87, 1991.
Siddhant M Jayakumar, Wojciech M Czarnecki, Jacob Menick, Jonathan Schwarz, Jack Rae, Simon Osindero, Yee Whye Teh, Tim Harley, and Razvan Pascanu. Multiplicative interactions and where to find them. 2020.
Michael I Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):1832233, 1999.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprini arXiv:1412.6980, 2014.
Diederik P Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes. arXiv preprint
Simon Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph Ledsam,
Klaus Maier-Hein, SM Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilistic u-net for segmentation of ambiguous images. Advances neural information processing systems, 31, 2018.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Abhishek Kumar and Ben Poole. On implicit regularization in 3-vaes. In International Conference on Machine Learning, pp. 5480-5490. PMLR, 2020.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, (1006266):132-133333 2015.
Hugo Larochelle and lain Murray, The neural autoregressive distribution estimator. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 29-37. JMLR Workshop and Conference Proceedings, 2011.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
Jonathan Lorraine and David Duvenaud. Stochastic hyperparameter optimization through hypernet- works. arXiv preprint arXiv: 18011000199999 2018.
James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! linear vae perspective on posterior collapse. Advances in Neural Information Processing Systems, 32, 2019.
Matthew MacKay, Paul Vicol, Jon Lorraine, David Duvenaud, and Roger Grosse. Self-tuning networks: Bilevel optimization of hyperparameters using structured best-response functions. arXiv preprint arXiv: 030880 2019.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset. http:///ttt...coo//eeemmm///ssss
2017.
Grégoire Montavon and Klaus-Robert Müller. Deep boltzmann machines and the centering trick. In Neural networks: tricks of the trade, pp. 621-637. Springer, 2012.
Akira Nakagawa, Keizo Kato, and Taiji Suzuki. Quantitative understanding of vae non-linearly scaled isometric embedding. In International Conference on Machine Learning, pp. 7916-7926. PMLR, 2021.
Aviv Navon, Aviv Shamsian, Gal Chechik, and Ethan Fetaya. Learning the pareto front with hypemetworks. arXiv preprint trXiv:2010.041000 2020.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 201 1.
12

---

Published as a conference paper at ICLR 2023
Sajad Norouzi, David J Fleet, and Mohammad Norouzi. Exemplar vae: Linking generative models, nearest neighbor retrieval, and data augmentation. Advances in Neural Information Processing Systems, 13:8753-8764, 2020.
Sconho Park, George Adosoglou, and Panos M Pardalos. Interpreting rate-distortion of variational autoencoder and using model uncertainty for anomaly detection. Annals of Mathematics and Artificial Intelligence, 90(7):335-752, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, Segh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Ethan Perez, Harm De Vries, Florian Strub, Vincent Dumoulin, and Aaron Courville. Learning visual reasoning without strong priors. arXiv preprint 1oXiv:777.00017, 2017.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Mary Phuong, Max Welling, Nate Kushman, Ryota Tomioka, and Sebastian Nowozin. The mutual autoencoder: Controlling information in latent code representations. 2018.
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-fit: State-covering supervvid reinforcement learning. arXiv preprint mxXiv:/9033 93390,698, 2019.
Razavi, Aäron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delsa-vaes. arXiv preprint arXiv:1901.03410 2019.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wienstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286 PMLR, 2014.
Adam Roberts, Jesse Engel, Colin Raffel, Curtis Hawthorne, and Douglas Eck. hierarchical latent vector model for learning long-term structure in music. In International conference on machine learning, pp. 4364-4373, PMLR, 2018
Oleh Rybkin, Kostas Daniilidis, and Sergey Levine. Simple and effective vae training with calibrated decoders, 2020.
Robert Sicks, Ralf Korn, and Stefanie Schwaar. A generalised linear model framework for 8- variational autoencoders based on exponential dispersion families. J. Mach. Learn. Res., 22:233-1, 2021.
MTCAJ Thomas and Thomas Joy. Elements of information theory. Wiley-Interscience, 2006.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 61(3):611-62 1-622, 1999.
Naftali Tishby, Fernando Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004405, 2000.
Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Neural Information Processing Systems (NeurIPS), 2020.
Yuanhao Wang, Guodong Zhang, and Jimmy Ba. On solving minimax optimization locally: A follow-the-ridge approach. arXiv preprint 17..: 2019.
Zihao Wang and Liu Ziyin. Posterior collapse of linear latent variable model. arXiv preprint wwXiv:2205.04009 2022.
Fei Yang, Luis Herranz, Joost Van De Weijer, José A Iglesias Guitián, Antonio M López, and Mikhail G Mozerov. Variable rate deep image compression with modulated autoencoder. IEEE Signal Processing Letters, 27:331-335, 2020.
13

---

Published as a conference paper at ICLR 2023
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
Lei Zhou, Chunlei Cai, Yue Gao, Sanbao Su, and Junmin Wu. Variational autoencoder for low bit-rate image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 2617-2624 2018.
14


Reviews:
[{"review_id": "x0YZym1XQUy", "reviewer": "Reviewer_47TW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes MR-VAE, an algorithm to train a VAE that corresponds to an interval of $\\beta$ values in a single training run.  The proposed method trains a hypernetwork that modifies the weights of a regular VAE at inference time for a given value of $\\beta$, thus eliminating the need to retrain a VAE for each value of $\\beta$.  The choice of hyperparameters for the hypernetwork is based on the analysis on linear VAEs, and the empirical results demonstrate that the method generalizes well to modern VAE architectures.", "summary_of_the_review": "This paper provides an answer to a somewhat obvious but practically important question of whether it's possible to train a VAE that can operate on a continuum of rate-distortion points.  The presentation is clear, with clear experimental results to back up the usability (i.e. R-D performance, computational/memory cost, stability to hyperparameters).  I believe this is a valuable contribution to the field.", "strength_and_weaknesses": "**Strengths**\n* Very well-written paper with a clear motivation.\n* Strong empirical results on diverse datasets.  In particular, the paper addresses the important questions that naturally arise: (1) does MR-VAE introduce performance degradation?  (2) does MR-VAE generalize to larger, complicated architectures?  (3) how sensitive is MR-VAE to the choice of its (only) additional hyperparameter $(a, b)$?  The experiments convincingly and positively answer all these questions.\n\n**Weaknesses**\n* I did not find any particular weakness.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aHsAekpXnqS", "reviewer": "Reviewer_fEL1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces Multi-Rate VAE, which is a hyper-network that outputs the parameters of a $\\beta$-VAE network, when given a value of $\\beta$. The rate-distortion curve created using the the networks computed from this hyper-network matches the one that is obtained by repeatedly training $\\beta$-VAEs for different values of $\\beta$. This hyper-network also has similar number of parameters as that of the $\\beta$-VAE network it outputs, and has similar training time. The construction of the hyper-network is theoretically principled, because the paper proves that the hyper-network can output the optimal $\\beta$-VAE when the VAE is linear.", "summary_of_the_review": "I think the idea of hyper-networks outputting $\\beta$-VAE for different values of $\\beta$ without the need for retraining from scratch for each $\\beta$, is of practical importance. Furthermore, the proposed framework is both theoretically principled and empirically verified to work.", "strength_and_weaknesses": "Strengths:\n1. The proposed hyper-network construction seems to be theoretically principled. This is because for linear VAEs, the paper provides a theorem that proves that indeed the optimized hyper-network can output parameters for the optimal linear $\\beta$-VAE.\n2. The hyper-network is of practical importance since it bypasses the need to retrain $\\beta$-VAEs for different values of $\\beta$, while fine tuning for optimal $\\beta$.\n3. The hyper-network has similar size and training time as the $\\beta$-VAE it outputs.\n4. The results have been empirically verified on MNIST, OMNIGLOT, CIFAR10 and CelebA datasets, where it is shown that the rate-distortion curve created using the the networks computed from this hyper-network matches the one that is obtained by repeatedly training $\\beta$-VAEs for different values of $\\beta$.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "BEUf6gLt_s", "reviewer": "Reviewer_MZus", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper considers $\\beta$-VAEs and treats the weight $\\beta$ of the KL-divergence term as a hyper-parameter. The authors develop a hypernetwork model which learns the optimal VAE parameters (i.e. weights of the encoder and decoder) simultaneously for all $\\beta$ values within a pre-specified range. The approach is experimentally tested on different datasets like MNIST, SVHN and others as well as for different VAE models like standard VAEs and hierarchical VAEs.", "summary_of_the_review": "The paper proposes an interesting hypernetwork approach for estimating the optimal weights of $\\beta$-VAEs simultaneously for all $\\beta$ values within a pre-specified range for given encoder/decoder architectures. The technical part (especially in  the supplementary material) is sometimes hard to follow and there are a few questions and issues remaining open. My current recommendation is a lower bound. The final recommendation will depend on the clarification of these questions in the discussion phase.\n\nAfter rebuttal discussion: The clarifications provided by the authors answered almost all my questions and are convincing. I am raising my recommendation to 8.", "strength_and_weaknesses": "Paper strengths:\nThe authors start from simple log-linear Gaussian VAEs and prove theoretically that a simple affine transform of the encoder/decoder weights parametrised by $\\beta$ provides the dependence of the optimal wights on the hyper-parameter $\\beta$. They conjecture that the hyper-network for complex, non-linear Gaussian VAEs is obtained by applying such affine transforms of the base weights in each layer of the encoder/decoder pair. The experiments compare achievable data likelihoods for VAEs (with  fixed architecture), when learned with fixed betas (aka response curve), with the one obtained from learning the hypernetwork. It is striking to see how good the corresponding rate curves are matched by the hypernetwork for different VAE variants.\n\nPaper weaknesses:\nThere are a few conceptual choices and issues which remain unclear for me:\nI was not able to fully follow the proofs given in the supplementary material for the log-linear Gaussian VAEs:\n- Step from (18) to (19,20): Taking the expectation over the data and maximising it w.r.t. to $\\mu$ will give a $\\mu_{MLE}$ that depends on the other model parameters. It remains unclear how to derive the gradients of the objective w.r.t. the latter.\n- Theorem 1: It remains unclear to me, why it is necessary to represent the parameter matrices of the model as products of \nmatrices? Why are then the affine transforms applied directly to the VAE weights for the non-linear models in (8,9)?\n\nIt remains unclear to me, why applying the affine transform in all layers of the encoder/decoder of non-linear VAEs is sufficient for obtaining the required hypernetwork. Is this may be related to the known rescaling properties of ReLU networks? Would this work also for other non-linearities?. And more general, can this approach be used also for non-Gaussian VAEs with discrete latent variables?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%