{"id":"c27811d5-701f-4446-9ddd-fd846161bf2e","arxiv_id":"2412.16899","paper_version":2,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"LMMVAE extends variational autoencoders with linear-mixed-model-style random effects, improving dimensionality reduction for correlated data.","lead":"This paper introduces LMMVAE, a variational autoencoder that adds random-effect latent variables to capture correlations between similar observations, such as repeated measures or nearby locations. It reports lower reconstruction errors and better downstream classification than standard VAEs on simulated and real tabular data.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Spatial loss is not a valid ELBO under the correlated prior; also Eq. (9) is the per-observation loss while matched to mini-batch averaging, and the matrix-normal KL in Appendix B mixes q- and p- dimensions.","rationale":"The reader's verdict is CONDITIONAL, and my concern matches their weakest_assumption: the spatial loss in Eq. (9) is used without deriving the KL for the correlated prior B ~ MN(0, K, D). This is load-bearing because the paper's headline claim explicitly covers spatial data and the spatial experiments are a major part of the evidence. The categorical and longitudinal variants appear internally consistent: for the categorical case the loss is derived (Appendix B), and the longitudinal case extends the categorical sum over K polynomial terms under a diagonal or block-diagonal prior, though it too relies on an independence assumption between polynomial terms that is only tested empirically. The spatial variant is qualitatively different because K introduces cross-observation correlation, which is exactly the paper's claimed contribution, yet the ELBO is never written down for it. The left-multiplication by Psi*^{1/2} enforces correlation in samples but does not correct the KL term; the paper even says 'With the current LMMVAE architecture one could think of B's posterior for a single categorical feature as MN(0, Iq, D), and so by multiplying it on the left by Psi*^{1/2}... The loss function, therefore, remains the same as in (9)', which conflates a reparameterization of the variational distribution with a change of prior. The reparameterization changes the variational distribution whose KL enters the ELBO; one cannot keep the independent-prior KL unless the rotated distribution happens to have the same KL, which requires proof. I am not flagging the mini-batch averaging of levels as the primary concern because the paper acknowledges it, and it is a practical device that still computes an unbiased-ish estimate; the spatial ELBO gap is more fundamental. The NLL metric comparison is also concerning because LMMVAE's likelihood includes the random-effect term, so lower NLL for LMMVAE may partly reflect a different generative model, not tighter inference; but the reconstruction error comparisons do not depend on this, so the central empirical claim survives in categorical and longitudinal settings. I therefore keep CONDITIONAL rather than REJECT: the categorical and longitudinal evidence is strong, but the spatial evidence needs to be re-derived or explicitly labeled heuristic before the paper's full claim can be accepted. The concrete test is a closed-form derivation of the spatial KL divergence, which is unambiguous and can be done by a referee in a few lines.","tokens_in":28064,"tokens_out":2542,"duration_ms":29960,"concrete_test":"Independent re-derivation of the ELBO for the spatial model: write the generative model p(x|u,b) N(f(u)+Zb, sigma_e^2 I) with prior p(b) = N(vec(B); 0, K ⊗ D) and variational q(b|x) = N(mu_b(x), diag(tau_b^2(x))). Compute DKL[q||p] in closed form for the per-observation and per-mini-batch objectives. If the resulting loss differs from Eq. (9) by a nonzero term involving K (e.g., sum of |mu_b,k|^2 times K^{-1} diagonal entries and log|K| terms that do not cancel), then the spatial training is optimizing a heuristic objective. If instead the paper's implied variational family (after rotating by Psi*^{1/2}) makes the KL equal to Eq. (9), the derivation must be written out explicitly. This check settles whether the spatial loss, and hence the spatial experiments, are variational or heuristic.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"The central claim is that LMMVAE improves reconstruction and NLL on unseen correlated data across categorical, longitudinal, and spatial scenarios. The spatial claim is the least secure. Section 3.3.3 changes the RE prior from B ~ MN(0, Iq, D) to B ~ MN(0, K, D), with rows correlated through the RBF kernel K, yet states \"The loss function, therefore, remains the same as in (9)\". Equation (9) was derived in Appendix B for an independent prior, D = diag(sigma_b^2), using a standard diagonal Gaussian surrogate q(b|x). With B ~ MN(0, K, D), the KL divergence DKL[q(b|x)||N(0,K)] contains an additional term involving (mu_b)^T K^{-1} (mu_b) and the log-determinant of K and the trace of K^{-1} times the variational covariance, unless the variational family is itself restricted to a degenerate or specially structured form. The paper does not derive this KL. Simply left-multiplying samples by Psi*^{1/2} makes the samples correlated, but the loss in (9) still uses the independent-prior KL, so the optimized objective is not the ELBO of the spatial generative model. Moreover, the spatial predictor for unseen locations replaces the posterior covariance with Psi* from training locations, which is a heuristic kriging-type approximation rather than a variational posterior for new locations, as the \"Unknown mode\" results show. The paper's abstract and Section 3.1 also claim generality of Eq. (6), but the loss derivations cover only the diagonal-D categorical case; longitudinal and spatial modifications are asserted by analogy, not derived. This does not invalidate the categorical and longitudinal empirical results (which appear strong), but it does invalidate the spatial variant as a principled variational method and weakens the claim that Eq. (6) is general. Since spatial experiments (Tables 13-16, 23-25) are a substantial part of the paper's evidence, the central claim should be qualified.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes LMMVAE, a variational autoencoder that augments the standard fixed latent model with a linear mixed-model random-effect term, so that the decoder generates x = f(u) + Zb + ε. The authors modify the ELBO to include separate encoders for the fixed latent u and the random effect b, and they derive an explicit loss for the high-cardinality categorical case with an independent random-effect prior B ~ MN(0, I_q, D), D diagonal. They then extend this construction by analogy to longitudinal data (multiple polynomial random effects) and spatial data (B ~ MN(0, K, D) with an RBF kernel), reporting substantial improvements in squared reconstruction error and negative ELBO loss on unseen data across simulated and real tabular datasets, plus a CelebA image experiment and downstream classification gains.","tokens_in":28481,"tokens_out":5170,"duration_ms":46873,"significance":"The idea of separating fixed and random latent effects is a natural and practically motivated extension of VAEs to correlated tabular data, and the empirical evaluation is unusually broad: multiple correlation scenarios, six real tabular datasets, comparisons to PCA, VAE variants, VRAE, and SGP-VAE, and publicly available code. The categorical and longitudinal derivations in Appendices A and B are standard and appear correct. However, the spatial variant, one of the three headline scenarios, is not derived as an ELBO under the correlated prior, and the mini-batch averaging step is a heuristic. These gaps weaken the theoretical grounding of a substantial portion of the empirical claims and need to be addressed before the paper can be accepted.","major_comments":[{"comment":"The spatial variant changes the random-effect prior to B ~ MN(0, K, D) with an RBF kernel K but states, without derivation, that 'the loss function, therefore, remains the same as in (9)'. Equation (9) is derived in Appendix B only for B ~ MN(0, I_q, D) with D diagonal. For B ~ MN(0, K, D), the KL divergence between a diagonal Gaussian q(b|x) and N(0, K) contains terms involving K^{-1}, log|K|, and Tr(K^{-1}Σ), none of which appear in Eq. (9). As written, the spatial objective is not the ELBO of the spatial generative model, so the spatial experiments in Tables 1, 13–16, and 23–25 train a heuristic rather than a variational lower bound. The paper must either derive the spatial ELBO correctly or explicitly reframe the spatial method as an approximation with its own objective.","section":"Section 3.3.3, Eq. (9)"},{"comment":"The model samples one p-dimensional b vector per observation in a mini-batch, but then averages these vectors by group level to form the B matrix before computing ZB. This averaging is a heuristic that does not follow from the ELBO in Eq. (8), where each observation x_ij has its own b_j in the generative model. The paper does not show that this averaging step is equivalent to optimizing a coherent variational objective, either for the categorical/longitudinal cases or for the spatial case. Since the group-level B is used both in training and in the prediction procedure, a derivation or a clear statement that this is an approximation is needed.","section":"Section 3.1, architecture description"},{"comment":"For unseen locations, the paper uses the Cholesky factor Ψ*^{1/2} computed from training locations and left-multiplies the sampled B. This is a kriging-style heuristic: no variational posterior for new random effects is defined, and the loss for test observations is not the ELBO conditional on the training posterior. The Unknown-mode results in Tables 14 and 23 should either be accompanied by a derivation or explicitly labeled as a predictive heuristic, so that the reader can interpret the reported gains appropriately.","section":"Section 3.3.3, 'Unknown mode' experiments"},{"comment":"The claim that the matrix formulation in Eq. (6) is 'general enough' to cover categorical, longitudinal, and spatial scenarios is stronger than what is actually derived. The loss derivations cover the independent-prior categorical case (Appendix B); the longitudinal extension changes the prior to MN(0, Φ ⊗ I_q, D) without computing the corresponding KL, and the spatial case changes the row covariance to K without re-deriving the KL. Please state precisely which scenarios have a fully derived ELBO and which are heuristic extensions, and adjust the abstract and Section 3.1 accordingly.","section":"Section 3.1, Eq. (6) and abstract"}],"minor_comments":[{"comment":"The paper uses 'negative log-likelihood (NLL)' to refer to the negative ELBO, but the reported quantities are bounds on the marginal likelihood, not the marginal likelihood itself. Please use 'negative ELBO' consistently or define the terminology explicitly.","section":"Throughout"},{"comment":"The column header 'n_j' is confusing for datasets with multiple categorical features (e.g., News and Spotify), where 'q' is reported as two or more values and 'n_j' is not a single number. Consider clarifying the column meaning or splitting it into group-size descriptions per feature.","section":"Table 2"},{"comment":"The symbol Ψ*^{1/2} is used without specifying the Cholesky convention (lower vs. upper triangular). For reproducibility, please state the convention used in the implementation.","section":"Section 3.3.3"}],"recommendation":"major_revision","confidential_remarks":"The paper is likely to be of interest to the journal's audience, but the spatial gap and the mini-batch averaging issue are load-bearing for the claimed generality. The relationship to the authors' prior work (Simchoni & Rosset, NeurIPS 2021 and JMLR 2023) should be clarified; the novelty statement 'first work to take inspiration from LMM for dimensionality reduction' may need to be softened or positioned relative to those earlier papers. With a proper spatial derivation or an explicit repositioning of the spatial experiments as heuristic, the paper could be acceptable."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The paper does something real: it takes the random-effects idea from linear mixed models and grafts it onto VAEs, giving a decoder of the form f(u) + Zb. That combination is new as far as I know, and the empirical work is substantial. On simulated and real data with categorical or longitudinal structure, the reconstruction gains over vanilla VAE, VAE with one-hot or entity embeddings, and GP-based competitors are large and consistent. The categorical/longitudinal derivation in Appendix B is standard but clean, and the authors ship code. That part deserves credit.\n\nThe soft spots are concentrated in the spatial variant. Section 3.3.3 changes the random-effects prior from MN(0, Iq, D) to MN(0, K, D) with a general kernel K, then says the loss is unchanged because they left-multiply the posterior samples by Ψ*^{1/2}. That does not work. The KL divergence in the ELBO changes when the prior has a nondiagonal row covariance; you get extra terms involving K^{-1} and the log determinant of K. The paper never derives those. So the spatial model is trained on an objective that is not the ELBO of the generative model it claims to fit. It may still be a useful heuristic — the empirical results in the spatial tables are good, especially in unknown-location mode — but it is not a principled variational method as presented. The same applies to the \"Unknown mode\" prediction, which replaces the posterior with a kriging-style plug-in rather than a variational posterior for new locations.\n\nTwo smaller issues. The mini-batch averaging of per-observation b samples (the ⊙ operation) is a heuristic, not a consequence of the ELBO; the paper does not analyze its bias. And the NLL numbers are compared across models with different latent structures, which makes them not directly comparable as likelihoods; they are more like a reconstruction-plus-regularization score. The abstract's claim that Eq. (6) is general is also stronger than the derivations support, since only the diagonal-D case is actually derived.\n\nMy overall read: the categorical and longitudinal results are believable and useful; the spatial variant needs either a proper derivation (even an approximate bound) or a clear statement that it is a heuristic. This is a solid conditional accept, not a desk reject. The paper deserves a serious referee, and the reviewer should push on the spatial loss and the NLL comparison. I would probably cite it for the categorical/longitudinal framework, but I'd steer around the spatial claim until it is fixed.","headline":"LMMVAE is a genuinely useful idea with strong empirical results for categorical and longitudinal data, but the spatial variant's loss is not a valid ELBO under the correlated prior and should be re-derived or qualified.","tokens_in":29026,"tokens_out":1424,"would_cite":true,"duration_ms":17028,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Adding random effects to variational autoencoders makes them far better at reconstructing correlated data, the paper reports.","keywords":["variational autoencoder","random effects","linear mixed models","correlated data","dimensionality reduction","reconstruction error","latent representations","longitudinal and spatial data"],"falsifier":"Compute the exact KL divergence between the correlated prior $B \\sim \\mathrm{MN}(0, K, D)$ and the factorized variational posterior used by the spatial model, and compare the resulting ELBO with equation (9); if the extra terms involving $K$ and $\\Psi_*$ are non-negligible, the reported spatial gains are not coming from maximizing the stated lower bound.","tokens_in":27888,"feed_emoji":"📉","tokens_out":6180,"duration_ms":55364,"temperature":0.7,"pith_summary":"This paper argues that the standard variational autoencoder's assumption of independent observations is the reason it underperforms on datasets with clustering, temporal, or spatial structure. It proposes LMMVAE, a variational autoencoder whose decoder output is $f(U)+ZB$, where $U$ is the usual independent latent variable and $B$ is a matrix of random effects shared by observations in the same group, in the spirit of linear mixed models. On held-out data from simulations and real tabular datasets, the model reports substantially lower squared reconstruction error and negative log-likelihood than standard VAEs, PCA baselines, and Gaussian-process-prior VAE variants, and the learned representations improve a downstream classification task. If the paper is right, a small architectural change—adding a random-effects term with its own encoder—is a general way to make deep generative dimensionality reduction respect known correlation structure.","feed_headline":"Random effects make VAEs far better on correlated data","feed_subtitle":"Splitting the latent space into fixed and random parts exploits clusters, time, and location for reconstruction and downstream tasks.","key_machinery":"The load-bearing object is the matrix model $X = f(U) + ZB + E$, with $B \\sim \\mathrm{MN}(0, K, D)$. Here $Z$ is a user-specified $n\\times q$ design matrix: an indicator matrix for categorical levels, polynomial-in-time columns for longitudinal data, and location indicators for spatial data; $D$ is a diagonal $p\\times p$ prior variance matrix for the $p$ feature-wise random effects; and $K$ is a $q\\times q$ between-group correlation kernel, the identity matrix for categorical and longitudinal settings and an RBF kernel for spatial data. The architecture uses an encoder that outputs the mean and variance of $B$, samples $B$ by the reparameterization trick, averages the samples per group to form $B$ for the mini-batch, and adds $ZB$ to the fixed decoder output. The loss in equation (9) is the ELBO with the KL term decomposed into fixed and random parts, so the only change relative to a standard VAE is one extra KL term per random-effect source. For spatial data, the paper enforces between-row correlation by left-multiplying the sampled $B$ by the Cholesky factor of the posterior covariance $\\Psi_*$ obtained from the best linear unbiased predictor formula.","core_discovery":"The central claim is that the generative model $X = f(U) + ZB + E$, where $U$ are ordinary independent latent variables and $B$ is a matrix of random effects with a matrix-normal prior $\\mathrm{MN}(0, K, D)$, is a strictly better VAE-type model for correlated data. The fixed part $f(U)$ is the usual non-linear decoder; the design matrix $Z$ encodes which cluster, subject, location, or time polynomial each row belongs to, and $B$ captures the shared deviations for those groups. The paper derives a modified ELBO whose KL divergence splits into a fixed and a random part under the assumption that $q_\\phi(u,b\\mid x)=q_\\phi(u\\mid x)q_\\phi(b\\mid x)$, yielding the loss in equation (9), and implements the model with separate encoders for $u$ and $b$. Across simulated high-cardinality categorical, longitudinal, and spatial settings, and on real datasets such as news shares, store sales, census income, and used-car listings, LMMVAE reports the best test reconstruction error and NLL, often by a wide margin, and better $k$-NN classification on the latent space for small latent dimension. On CelebA images with face identity as the random effect, the method matches but does not beat a VAE that ignores the identity, a result the paper attributes to the need for additional external features.","pith_inferences":["If the spatial loss in equation (9) is not the exact ELBO for the correlated prior $B \\sim \\mathrm{MN}(0, K, D)$, the spatial results may still stand as a heuristic, but the claim of maximizing a variational lower bound in that setting would need re-derivation; a direct computation of the KL divergence would settle it.","The additive random-effects decomposition could be transferred to other deep generative models, such as normalizing flows or diffusion decoders, whenever rows are exchangeable within known groups, because the argument relies only on the additive structure in the decoder.","Treating $D$ and $K$ as learnable variance components rather than fixed priors might improve performance further and would connect LMMVAE to empirical-Bayes mixed-model estimation; the paper fixes $D=\\sigma_b^2 I$ and the RBF length scale.","The independence assumption between $U$ and $B$ given $X$ could be relaxed by sharing encoder features, and the paper's own observation that a single double-output encoder is slightly inferior suggests the separation itself, not just the architecture, is doing the work."],"forward_implications":["On datasets with known grouping structure, practitioners can expect LMMVAE to reconstruct held-out rows markedly better than a VAE that ignores the grouping, because the random-effect term carries the group-level shared signal.","The learned random-effects matrix $\\hat{B}$ can be read as per-group offsets for each feature, giving a built-in diagnostic such as mapping location effects onto a map without fitting a separate model.","The same architecture covers categorical, longitudinal, spatial, and combined scenarios by changing only $Z$ and the prior on $B$, so one implementation handles several correlation types.","For downstream tasks, the fixed latent representation $U$ from LMMVAE supports better supervised classification at small latent dimension than standard VAE or PCA representations on the tested datasets.","On image data with a categorical random effect, the paper's experiments show negligible gains over ignoring the effect, suggesting the benefit concentrates on tabular data or requires additional external features."],"supporting_citations":[{"why":"Supplies the probabilistic PCA formulation $X = WU + \\varepsilon$ that LMMVAE generalizes by replacing the linear map with a non-linear function and adding the random-effects term.","marker":"[40]"},{"why":"Defines the variational autoencoder and the ELBO that LMMVAE modifies.","marker":"[18]"},{"why":"Supplies the linear mixed model framework, the $N(0,D)$ random-effects prior, and the best linear unbiased predictor formulas used for spatial correlations.","marker":"[23]"},{"why":"Provides the $\\beta$-weighted KL term in the loss that the paper adopts for both the fixed and random parts.","marker":"[13]"},{"why":"Introduces the Gaussian-process-prior VAE baseline whose independence limitation LMMVAE addresses and which the experiments compare against.","marker":"[4]"},{"why":"Provides the scalable Gaussian-process VAE baseline used as a comparison on longitudinal and spatial tasks.","marker":"[15]"},{"why":"Provides the variational recurrent autoencoder baseline for longitudinal data, comparing LSTM-based sequence modeling to LMMVAE.","marker":"[10]"},{"why":"Gives the matrix-normal distribution machinery used to enforce between-row correlation in the spatial variant via Cholesky multiplication.","marker":"[11]"}],"fun_headline_variants":["VAE latent splits into fixed and random for correlated data","Random effects in VAE improve reconstruction on clusters and time","LMMVAE: mixing random effects into VAEs cuts error on correlated data","Spatial, temporal, clustered data? Random effects boost VAE"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"For spatial data, the paper assumes that the loss derived for independent random effects remains the correct variational objective when the random effects are instead correlated across locations through a kernel matrix; if that loss is not the true lower bound, the spatial experiments train a heuristic rather than the stated model.","fun_headline_variants_meta":{"raw":{"variants":["VAE latent splits into fixed and random for correlated data","Random effects in VAE improve reconstruction on clusters and time","LMMVAE: mixing random effects into VAEs cuts error on correlated data","Spatial, temporal, clustered data? Random effects boost VAE"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00025,"raw_usage":{"total_tokens":1584,"prompt_tokens":1004,"completion_tokens":580,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":620,"completion_tokens_details":{"reasoning_tokens":505}},"tokens_in":620,"tokens_out":580,"duration_ms":5834,"temperature":1.0,"reasoning_tokens":505,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T06:00:16.927585+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Compute the exact KL divergence between the correlated prior $B \\sim \\mathrm{MN}(0, K, D)$ and the factorized variational posterior used by the spatial model, and compare the resulting ELBO with equation (9); if the extra terms involving $K$ and $\\Psi_*$ are non-negligible, the reported spatial gains are not coming from maximizing the stated lower bound.","supporting_citations":[{"cited_title":"Tipping and Christopher M","cited_arxiv_id":null,"evidence_quote":"Supplies the probabilistic PCA formulation $X = WU + \\varepsilon$ that LMMVAE generalizes by replacing the linear map with a non-linear function and adding the random-effects term."},{"cited_title":"McCulloch, Shayle R","cited_arxiv_id":null,"evidence_quote":"Supplies the linear mixed model framework, the $N(0,D)$ random-effects prior, and the best linear unbiased predictor formulas used for spatial correlations."},{"cited_title":"Gaussian process prior variational autoencoders","cited_arxiv_id":null,"evidence_quote":"Introduces the Gaussian-process-prior VAE baseline whose independence limitation LMMVAE addresses and which the experiments compare against."},{"cited_title":"Scalable gaussian process variational autoencoders","cited_arxiv_id":null,"evidence_quote":"Provides the scalable Gaussian-process VAE baseline used as a comparison on longitudinal and spatial tasks."}],"review_version":1}