Pith. sign in

REVIEW 4 major objections 5 minor 27 references

Beyond Diagonal Covariance: Flexible Posterior VAEs via Free-Form Injective Flows

T0 review · 4 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read This paper claims that a regularized free-form injective flow can serve as a variational autoencoder with a full-covariance Gaussian posterior, and that this posterior costs about as much as a diagonal one while improving model likelihood.

desk verdict Neat idea, but the central derivation breaks: the claimed full-covariance posterior is not the Laplace posterior even in the linear case, and the optimized loss is a stop-gradient surrogate, not the stated ELBO. read the letter →

arxiv 2506.01522 v1 pith:L7XIVR5E submitted 2025-06-02 cs.LG stat.ML

Peter Sorrenson , Lukas Lührs , Hans Olischläger , Ullrich Köthe This is my paper
classification cs.LGstat.ML
keywords variationalautoencoderfullcovariancefree-forminjectiveflowLaplaceapproximationpull-backmetricinvolutivityposteriorcollapseimagelikelihood
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper claims that the standard diagonal-Gaussian posterior used in variational autoencoders is not just a practical convenience but a genuine geometric limitation: for data manifolds whose induced metric has non-involutive eigen-line fields, no reparameterization of the decoder can make the pull-back metric diagonal. To get past this, it introduces FIVE, a VAE whose posterior is a free-form injective flow regularized so that $q(z|x) = \mathcal{N}(f(x), \sigma^2 f'(x)f'(x)^\top)$. This posterior is a full Gaussian covariance, but it is represented implicitly through the encoder Jacobian, so sampling and the KL term cost about the same as a diagonal VAE. If correct, this makes rich posterior covariances practical at scale and explains why diagonal VAEs distort curved manifolds. On MNIST and CIFAR-10, the paper reports that FIVE matches or surpasses the test log-likelihood of standard VAEs, full-covariance VAEs, and standalone flows.

What carries the argument

The load-bearing object is the variational posterior $q(z|x) = \mathcal{N}(f(x), \sigma^2 f'(x)f'(x)^\top)$, a full-covariance Gaussian whose covariance is written as the outer product of the encoder Jacobian with itself, so it never needs to be materialized. Its justification is a Laplace approximation: with decoder $g$ and noise $\sigma$, the posterior Hessian obeys $H(x)^{-1} \approx \sigma^2 f'(x)f'(x)^\top$. The supporting machinery is the pull-back metric $G = g'(z)^\top g'(z)$ from differential geometry, used to prove the diagonal-covariance obstruction through the Frobenius involutivity condition, plus a stochastic trace estimator and a stop-gradient trace identity that compute gradients of the log-determinant term without evaluating determinants. Sampling from the posterior is just a reparameterized draw $z = f(x) + \sigma f'(x)v$ with $v \sim \mathcal{N}(0,I)$, and the KL divergence to the standard normal prior is evaluated using $\operatorname{tr}(\sigma^2 f'(x)f'(x)^\top)$ and the approximate log-determinant gradient.

What would settle it

Take a decoder whose manifold has curvature radius comparable to or smaller than the observation noise $\sigma$ (for example $z = x^2 + y^2$ with $\sigma$ not very small) and compute the exact posterior $p(z|x) \propto \exp(-u(x,z))$ by dense integration in the low-dimensional latent space; if the KL divergence between $p(z|x)$ and $\mathcal{N}(f(x), \sigma^2 f'(x)f'(x)^\top)$ is large, or if FIVE's learned decoder no longer matches the true manifold's parameterization, the local-linearity premise fails.

Watch

Extended reading notes

Core claim

As the paper states it, a VAE that forces $q(z|x)$ to be diagonal Gaussian will at optimum force the decoder's pull-back metric $g'(z)^\top g'(z)$ to be approximately diagonal (Theorem 1), because the true posterior $p(z|x)$ approximates $\mathcal{N}(f(x), H(x)^{-1})$ with $H = I + \sigma^{-2} g'(f(x))^\top g'(f(x))$. Orthogonal coordinate systems with diagonal metric exist only when the eigen-line fields of this metric are involutive (Theorem 2), and in dimension three or higher curved manifolds can fail this test. FIVE sidesteps the constraint by defining the variational posterior directly from the encoder Jacobian, $q(z|x) = \mathcal{N}(f(x), \sigma^2 f'(x)f'(x)^\top)$, which is the Laplace-approximation form of $H(x)^{-1}$. The paper proves that in the linear case the two opposing Jacobian substitutions used in the derivation are consistent and the model recovers the data covariance exactly (Theorem 3). The empirical claim is that this full-covariance posterior, with no explicit covariance matrix ever formed, matches or exceeds the likelihoods of diagonal VAEs, learned full-covariance VAEs, and standalone free-form injective flows on MNIST and CIFAR-10.

Load-bearing premise

The argument rests on the Laplace/local-linearity assumption that the decoder $g$ is nearly linear at the scale of the noise $\sigma$, so the true posterior $p(z|x)$ is well approximated by a Gaussian whose inverse covariance is $\sigma^2 f'(x)f'(x)^\top$; the paper's own Limitations section notes this may break down when the data manifold curves sharply on that scale.

Editorial extensions

If this is right

  • A VAE can carry a full-covariance Gaussian posterior with sampling cost equal to one extra Jacobian-vector product per latent draw, so the diagonal-Gaussian restriction is no longer forced by computation.
  • Posterior collapse that appears in large-latent diagonal and explicitly full-covariance VAEs is reported to be much weaker for FIVE, which inherits the collapse resistance of free-form injective flows.
  • The involutivity criterion gives a principled way to identify data manifolds for which diagonal posteriors are provably insufficient: check whether the eigen-line fields of the decoder metric can be integrated.
  • Combining the FIF training objective with a KL-style regularizer yields a valid ELBO whose approximate KL term needs no determinant computation.
  • On MNIST and CIFAR-10, with both small and larger networks, FIVE reports test log-likelihoods that match or exceed standard VAE, FC-VAE, and FIF baselines.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If the Laplace assumption holds for real image manifolds, the same Jacobian-covariance construction could be applied to other implicit or explicit encoder-decoder architectures, not only FIF-trained ones, whenever the encoder Jacobian is available.
  • The geometric obstruction suggests a testable diagnostic: for a trained decoder, compute the commutators $[e_i, e_j]$ of the metric's eigen-frame; datasets whose manifolds violate involutivity should show a systematic gap between diagonal-VAE and FIVE likelihood.
  • Theorem 3's linear recovery result hints that FIVE may behave like a probabilistic autoencoder that finds a principal-subspace structure; a synthetic experiment with data drawn from a known low-rank Gaussian could directly check whether the learned encoder and decoder match the top-$d$ eigen-directions.
  • At larger $\sigma$ or sharper curvature, $q(z|x)$ may remain Gaussian while $p(z|x)$ becomes non-Gaussian; a two-dimensional toy with exactly computed posteriors could measure where the approximation degrades.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 5 minor

Summary. FIVE is a variational autoencoder whose approximate posterior is defined implicitly by the encoder Jacobian: q(z|x)=N(f(x), σ² f'(x)f'(x)^T) (Eq. 32). The paper argues that this yields a full-covariance Gaussian posterior at computational cost near that of a diagonal VAE, provides a differential-geometric argument (Theorem 2) that diagonal-posterior VAEs cannot represent arbitrary curved data manifolds, derives a regularized FIF training objective, and reports held-out test log-likelihoods on MNIST and CIFAR-10 in which FIVE matches or exceeds VAE, full-covariance VAE, and FIF baselines. The theoretical connection to the Laplace posterior is the main basis for the claim.

Significance. The geometric obstruction result (Theorem 2) is a clean and useful observation, and the empirical comparison is careful: test likelihoods are estimated by importance sampling with 100 samples, standard errors are reported over three runs, and two architectures per dataset are considered. If the posterior interpretation were valid, FIVE would be a practically attractive way to obtain full covariance without O(d²) parameters or determinant computations. However, the central derivation equating q(z|x) with the Laplace posterior is not valid in the finite-σ regime used in the experiments, and the theorem offered in support overclaims what is proved. The empirical likelihood results can stand on their own, but the paper's theoretical significance as a principled full-covariance VAE is currently not established.

major comments (4)
  1. [Section 4.3 / Appendix A.3, Eq. (31) and Lemma 1] The substitution f'(x) ≈ (g'(f(x)))^+ underlying Eq. (31) is already contradicted by the paper's own linear analysis. In Lemma 1 the global optimum of the FIVE loss is w* = √λ and v* = √λ/(σ²+λ). Then v* w* = λ/(σ²+λ) ≠ 1 for finite σ, so v* is not the pseudoinverse of w*. The resulting posterior q(z|x) = N(vx, σ²v²) has variance σ²λ/(σ²+λ)², while the Laplace posterior variance from Eq. (12) or (48) is σ²/(σ²+λ). These differ by the factor λ/(σ²+λ). Thus Eq. (31) holds only in the limits σ→0 or λ→∞, and the statement that q(z|x) is the Laplace posterior is unsupported in the finite-σ regime of the experiments.
  2. [Section 4.3, Theorem 3] The theorem statement overclaims exact recovery of the data distribution. The proof establishes W*W*^T = Σx and hence g#p(z) = N(0, Σx), but the marginal of the generative model is p(x) = ∫ p(x|z)p(z) dz = N(0, Σx + σ²I), which equals N(0, Σx) only when σ→0. The proof itself acknowledges this ('which is also equal to p(x) in the zero σ limit'), so Theorem 3 as stated for the FIVE loss with σ>0 is false. It should be restated as a σ→0 limit or as a statement about the decoder's pushforward, not about the model marginal.
  3. [Appendix A.3, Eq. (68); Appendix A.4, Eq. (93)-(94)] The objective actually optimized is not the ELBO of Section 4.2. Eq. (68) replaces the log-determinant term of the KL divergence with a stop-gradient trace surrogate, and the main text does not state that the training loss is this surrogate rather than the ELBO (34). In addition, the derivation in A.4 moves from Eq. (93) to Eq. (94) as though (g'(f(x)))^+^T g'(f(x))^T = I; for an injective decoder with d < n the product equals (g'(f(x)) g'(f(x))^+)^T, the transpose of a rank-d column-space projection, not the identity. The gradient formula therefore requires an additional assumption or correction. These issues matter because the paper's central claim is that FIVE optimizes a valid variational objective with a full-covariance posterior.
  4. [Section 5, Table 1 and Abstract] The empirical conclusion is weaker than the abstract's claim that full covariance 'substantially improves model likelihood.' On CIFAR-10 with the larger network, FIVE (2684.44±9.33) is not statistically distinguishable from the standard VAE (2680.36±13.86), and on MNIST with the larger network, FIF (1116.31±27.29) has numerically higher mean than FIVE (1115.17±7.67). The gains are consistent on small networks but not uniformly significant. This does not undermine the empirical contribution, but the summary claims should be calibrated to the table.
minor comments (5)
  1. [Section 1] There is a duplicated phrase: 'enrich the latent structure enrich the latent structure'.
  2. [Appendix A] Theorem numbering is duplicated between the main text and appendix (Theorem 4 restates Theorem 1, Theorem 5 restates Theorem 2, Theorem 6 restates Theorem 3); renumber the appendix theorems.
  3. [Eq. (35) and Eq. (27)] The displayed KL expression in Eq. (35) has an unbalanced parenthesis, and the stop-gradient operator SG should be defined before its first use in Eq. (27).
  4. [Limitations section] The Limitations section lists dataset scope but does not mention that the training objective is a stop-gradient surrogate; this should be disclosed as a limitation.
  5. [Section 4 / Section 5.1] The claim that computational cost is 'comparable to diagonal VAEs' should be quantified; the Hutchinson trace estimate and Jacobian-vector products add overhead per iteration.

Circularity Check

0 steps flagged · score 2.0 of 10

No significant circularity: the FIVE posterior is defined rather than fitted, likelihoods are evaluated on held-out data, and the same-author FIF citation is not load-bearing.

full rationale

The paper's central derivation is self-contained in the sense that q(z|x) = N(f(x), sigma^2 f'(x)f'(x)^T) is introduced as a definition in Eq. (32), not as a quantity fitted to the quantity it claims to predict. The Laplace approximation in Section 3 is an explicit assumption, and the 'in-out' substitution in Eq. (31) is an unproven approximation whose failure at the paper's own linear optimum (Lemma 1) is a correctness or validity concern, not a circularity: the approximation is not obtained by fitting a parameter to a target and then measuring that target. The empirical claims in Table 1 are tested on held-out test sets via importance sampling with 100 samples, which is an external benchmark and does not reduce to the training objective. The only same-author citation, to Sorrenson et al. (2024) for the free-form injective flow and its Hutchinson-style trace estimator, supplies a published, independently usable component and is not used to justify the posterior's covariance form or to forbid alternatives. Theorems 1 and 2 rest on standard differential geometry arguments rather than on a self-citation chain. Overall, no load-bearing step is circular by construction, so the circularity score is low.

Assumptions & free parameters 1 free parameters · 4 assumptions · 0 invented entities

The ledger is small: the model introduces no new physical entities and only one standard learnable parameter. The load-bearing assumptions are the Laplace/local-linearity condition and the in-out substitution, both of which are explicitly acknowledged as approximations in the paper.

free parameters (1)
  • learnable noise level sigma^2 = learned during training, value not reported
    The output noise variance is made trainable following Rybkin et al. It appears in the posterior covariance and affects the Laplace approximation, but it is a standard VAE parameter rather than a hand-tuned constant specific to FIVE.
assumptions (4)
  • domain assumption The Laplace approximation: u(x,z) is quadratic in z around f(x), so second and higher derivatives of g are negligible over the noise scale sigma
    Used in Eqs (20)-(23) and again in Eqs (30)-(31); the paper's own Limitations section says this breaks down on sharply curved manifolds.
  • ad hoc to paper The 'in-out' substitution f'(x) approx (g'(f(x)))^+ and its reverse in the log-det gradient are valid beyond the linear case
    Eqs (38)-(39) make two opposing substitutions; Theorem 3 only verifies consistency for linear maps in the sigma to 0 and d=n limit, so the nonlinear validity is an unproven premise.
  • standard math Frobenius integrability theorem and the equivalence of existence of orthogonal coordinates with involutive eigen-line fields
    Used to prove Theorem 2 in appendix A.2 and to motivate limitations of diagonal posteriors.
  • domain assumption The variational posterior q(z|x)=N(f(x), sigma^2 f'(x)f'(x)^T) is used for importance sampling even when it may not be close to the true posterior
    Used in Eqs (32) and (41)-(43); if q is a poor approximation, importance sampling estimates become unreliable, as acknowledged for FIF.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Beyond Diagonal Covariance: Flexible Posterior VAEs via Free-Form Injective Flows." pith.science (2026). https://pith.science/paper/L7XIVR5E

@misc{pith2026250601522,
  author       = {Pith},
  title        = {Pith review of: Beyond Diagonal Covariance: Flexible Posterior VAEs via Free-Form Injective Flows},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/L7XIVR5E}},
  note         = {Machine review of arXiv:2506.01522}
}
read the original abstract

Variational Autoencoders (VAEs) are powerful generative models widely used for learning interpretable latent spaces, quantifying uncertainty, and compressing data for downstream generative tasks. VAEs typically rely on diagonal Gaussian posteriors due to computational constraints. Using arguments grounded in differential geometry, we demonstrate inherent limitations in the representational capacity of diagonal covariance VAEs, as illustrated by explicit low-dimensional examples. In response, we show that a regularized variant of the recently introduced Free-form Injective Flow (FIF) can be interpreted as a VAE featuring a highly flexible, implicitly defined posterior. Crucially, this regularization yields a posterior equivalent to a full Gaussian covariance distribution, yet maintains computational costs comparable to standard diagonal covariance VAEs. Experiments on image datasets validate our approach, demonstrating that incorporating full covariance substantially improves model likelihood.

Figures

Figures reproduced from arXiv: 2506.01522 by the authors.

Figure 1
Figure 1. Construction of a local orthogonal coordinate system in 2 dimensions from eigenvectors [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. Left: Illustration of non-involutive eigen-lines. The planes spanned by [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗
Figure 3
Figure 3. Parameterizations of the manifold. Left: ground truth explicit parameterization [PITH_FULL_IMAGE:figures/full_fig_p017_3.png] view at source ↗
Figures from the paper (2 more)
Figure 4
Figure 4. Figure 4: Encodings of the manifold. Left: ground truth latent-space coordinate grid. The remaining [PITH_FULL_IMAGE:figures/full_fig_p017_4.png]
Figure 5
Figure 5. Figure 5: Different learned variational posteriors [PITH_FULL_IMAGE:figures/full_fig_p018_5.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

27 extracted references · 18 canonical work pages

  1. [1]

    Flows for simultaneous manifold learning and density estimation

    Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estimation. Advances in Neural Information Processing Systems, 33: 0 442--453, 2020

    2020

  2. [2]

    Rectangular flows for manifold learning

    Anthony L Caterini, Gabriel Loaiza-Ganem, Geoff Pleiss, and John P Cunningham. Rectangular flows for manifold learning. Advances in Neural Information Processing Systems, 34: 0 30228--30241, 2021

    work page 2021

  3. [3]

    Guild ai

    Guild AI Contributors. Guild ai. https://github.com/guildai/guildai

    work page

  4. [4]

    Inference suboptimality in variational autoencoders

    Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. In International conference on machine learning, pages 1078--1086. PMLR, 2018

    work page 2018

  5. [5]

    Free-form flows: Make any architecture a normalizing flow

    Felix Draxler, Peter Sorrenson, Lea Zimmermann, Armand Rousselot, and Ullrich K \"o the. Free-form flows: Make any architecture a normalizing flow. In International Conference on Artificial Intelligence and Statistics , 2024

    work page 2024

  6. [6]

    Harris, K

    Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, et al. Array programming with NumPy , 2020

    work page 2020

  7. [7]

    J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 9 0 (3): 0 90--95, 2007. doi:10.1109/MCSE.2007.55

    doi:10.1109/mcse.2007.55 2007

  8. [8]

    A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines

    Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics-Simulation and Computation, 18 0 (3): 0 1059--1076, 1989

    work page 1989

Show all 27 references
  1. [9]

    Auto-encoding variational Bayes

    Diederik P Kingma and Max Welling. Auto-encoding variational Bayes . In International Conference on Learning Representations , 2014

    2014

  2. [10]

    An introduction to variational autoencoders

    Diederik P Kingma, Max Welling, et al. An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12 0 (4): 0 307--392, 2019

    2019

  3. [11]

    Improved variational inference with inverse autoregressive flow

    Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in Neural Information Processing Systems, 29, 2016

    2016

  4. [12]

    Prince, and Marcus A

    Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43 0 (11): 0 3964--3979, 2021. ISSN 0162-8828, 2160-9292, 1939-3539. doi:10.1109/TPAMI.202...

    2021

  5. [13]

    Learning multiple layers of features from tiny images

    Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, Department of Computer Science, University of Toronto, 2009

    2009

  6. [14]

    Gradient-based learning applied to document recognition

    Yann Lecun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86 0 (11): 0 2278--2324, 1998

    1998

  7. [15]

    John M. Lee. Introduction to Smooth Manifolds, volume 218 of Graduate Texts in Mathematics. Springer, 2nd edition, 2013

    2013

  8. [16]

    Biva: A very deep hierarchy of latent variables for generative modeling

    Lars Maal e, Marco Fraccaro, Valentin Li \'e vin, and Ole Winther. Biva: A very deep hierarchy of latent variables for generative modeling. Advances in neural information processing systems, 32, 2019

    2019

  9. [17]

    Normalizing flows for probabilistic modeling and inference

    George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22 0 (57): 0 1--64, 2021

    2021

  10. [18]

    Variational Laplace autoencoders

    Yookoon Park, Chris Kim, and Gunhee Kim. Variational Laplace autoencoders. In International conference on machine learning, pages 5032--5041. PMLR, 2019

    2019

  11. [19]

    Pytorch: An imperative style, high-performance deep learning library

    Adam Paszke, Sam Gross, Francisco Massa, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019

    2019

  12. [20]

    Variational inference with normalizing flows

    Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530--1538. PMLR, 2015

    2015

  13. [21]

    Stochastic backpropagation and approximate inference in deep generative models

    Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278--1286. PMLR, 2014

    2014

  14. [22]

    Simple and effective vae training with calibrated decoders

    Oleh Rybkin, Kostas Daniilidis, and Sergey Levine. Simple and effective vae training with calibrated decoders. In International conference on machine learning, pages 9179--9189. PMLR, 2021

    2021

  15. [23]

    Exact solutions to the nonlinear dynamics of learning in deep linear neural networks

    A Saxe, J McClelland, and S Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In Proceedings of the International Conference on Learning Represenatations 2014. International Conference on Learning Represenatations 2014, 2014

    2014

  16. [24]

    Ladder Variational Autoencoders

    Casper Kaae S nderby, Tapani Raiko, Lars Maal e , S ren Kaae S nderby , and Ole Winther. Ladder Variational Autoencoders . In Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016

    2016

  17. [25]

    Lifting architectural constraints of injective flows

    Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, Lea Zimmermann, and Ullrich K \"o the. Lifting architectural constraints of injective flows. In International Conference on Learning Representations , 2024

    2024

  18. [26]

    VAE with a VampPrior

    Jakub Tomczak and Max Welling. VAE with a VampPrior . In International Conference on Artificial Intelligence and Statistics , pages 1214--1223. PMLR, 2018

    2018

  19. [27]

    Improving variational auto-encoders using Householder flow

    Jakub M Tomczak and Max Welling. Improving variational auto-encoders using Householder flow. arXiv preprint arXiv:1611.09630, 2016

    2016 arXiv

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.