{"id":"7210e1f3-a6bd-4092-aa3d-fb3a7b704914","arxiv_id":"2608.12084","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"A conditional surrogate loss that always picks the gradient estimate aligned with the reconstruction loss improves flow autoencoder training and reaches state-of-the-art generative performance on molecules, tabular data, and images.","lead":"This paper proposes a new loss for autoencoders that double as normalizing flows, choosing between two gradient tricks based on which one agrees with the reconstruction error. It reports better generation quality than prior methods on molecules, tabular data, and images.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"§3.4 drops all off-diagonal probe terms; with K=1 the mask sign(c_ii) need not align with the true reconstruction gradient, so the claimed proof and mechanism are not established.","rationale":"The paper's contribution is a conditional surrogate loss whose selection rule is derived from a truncated probe expansion of the reconstruction gradient. Re-deriving §3.4 and Algorithm 1 shows that the mask condition surrD>s_k is algebraically equivalent to sign(v_k^T(JG−I)v_k), so the authors are using the diagonal coefficient c_kk as the selection signal. Everything rests on whether the off-diagonal terms in the double sum can be ignored. They provide no bound, no empirical diagnostic, and all experiments use a single Hutchinson probe, so the per-probe sign is a noisy scalar rather than a basis-resolved diagonal. This is exactly the reader's weakest_assumption. I could not find a more fundamental flaw: the empirical comparisons are extensive, the method is simple to implement, and the toy alignment plots (Fig. 2) are suggestive but not decisive because they measure alignment of each surrogate with the reconstruction gradient separately, not the conditional update's alignment with the full gradient. Since the concern is addressable by a small exact-Jacobian experiment and does not by itself disprove the empirical results, the reader's CONDITIONAL verdict is appropriate; I would not move it.","tokens_in":22167,"tokens_out":8014,"duration_ms":73018,"concrete_test":"On the 4D toy dataset (d=2 and d=4), for several checkpoints along training and 100 random probe resamples per checkpoint, compute (i) the Algorithm-1 conditional-loss gradient with K=1 and (ii) the exact gradient of the Jacobian-regularization objective ∥I_d−N∥_F^2 from Eq. (14), obtained by differentiating through exact J and G Jacobians. Report the mean cosine similarity between (i) and (ii), and the same quantity for the encoder-surrogate and decoder-surrogate gradients. If the conditional gradient's mean cosine similarity is not clearly positive and above the single-surrogate baselines, the diagonal truncation in §3.4 is load-bearing and the claimed alignment mechanism is not supported.","verdict_should_be":"UNCHANGED","load_bearing_attack":"In §3.4, Eq. (15) expands ∇Lrec into a double sum over probe indices i,j. The paper then truncates to i=j ('Finally, focusing on the term with equal indices') before defining the conditional selection. This truncation is load-bearing: Algorithm 1 line 10 sets mask←1[surrD>s_k], and since surrD=v_k^TJGv_k and s_k=||v_k||^2, the mask is exactly sign(v_k^T(JG−I)v_k)=sign(c_kk). The derivation only justifies choosing by sign(c_kk) if the off-diagonal terms v_i^T(N−I)v_j and v_i^T(∇N)v_j are negligible. No such argument or diagnostic is given; for d>1 there are more off-diagonal than diagonal terms, and all experiments use K=1, i.e., a single random probe rather than a full basis. If off-diagonal terms are non-negligible, sign(c_kk) does not determine whether the selected surrogate's gradient is aligned with ∇Lrec, so the central 'alignment' motivation and the claimed proof of suboptimality of prior losses are unsupported. The SOTA empirical results may still hold, but they do not validate the mechanism as stated.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper studies 'flow autoencoders,' i.e., autoencoder architectures with separately parameterized encoder and decoder trained as approximate inverses, and proposes the Normalizing Autoencoder (NAE). The central contribution is a conditional surrogate loss that, per Hutchinson probe, selects either the encoder surrogate or the decoder surrogate based on a scalar mask intended to align the chosen surrogate gradient with the reconstruction-loss gradient. The authors first derive a relation between the reconstruction loss and the latent-space Jacobian residual, then use a probe-basis expansion to justify the mask, and finally report extensive experiments on molecular generation (DW4, LJ13, LJ55, QM9), tabular benchmarks (Power, Gas, HEPMASS, MiniBooNE), and image benchmarks (CelebA, MNIST, CIFAR-10). NAE achieves state-of-the-art or competitive results on most benchmarks, often with substantially better likelihoods or FID scores than prior flow-autoencoder baselines.","tokens_in":22447,"tokens_out":4311,"duration_ms":42825,"significance":"If the empirical results hold, NAE is a practically strong generative model: it improves over prior flow autoencoders on standard benchmarks while inheriting their architectural flexibility and sampling speed, and it does so with a minimal modification to the training objective. The paper is well positioned in the literature and includes a useful ablation study on the reconstruction weight and latent dimension, as well as honest discussion of limitations. The theoretical framing is less established: the central derivation in Section 3.4 rests on a diagonal truncation of a probe expansion and on a single-probe estimator whose justification is not supplied. Because the empirical gains are claimed to follow from this alignment mechanism, the gap is load-bearing rather than cosmetic. A revision that either proves the needed diagonal dominance, supplies empirical diagnostics for it, or appropriately reframes the theoretical claims would make the paper's contribution more solid.","major_comments":[{"comment":"The derivation expands the reconstruction-loss gradient as a double sum over probe indices i,j: Σ_i Σ_j (v_i^T(N−I_d)v_j)(v_i^T(∇N)v_j), then truncates to the equal-index terms with 'Finally, focusing on the term with equal indices' before defining c_ii and the conditional mask. This truncation is load-bearing: Algorithm 1 line 10 sets mask = 1[surrD > s_k], which with K=1 is exactly sign(c_kk), and the selection rule is only justified if the off-diagonal products are negligible. No argument or diagnostic for this diagonal dominance is provided; for d>1 the off-diagonal terms outnumber the diagonal ones. Moreover, all experiments use a single probe (K=1), so the completeness relation VV^T = d I_d used in the expansion does not hold for the realized estimator; a single random v satisfies E[vv^T] = I only in expectation, not for the actual sample. The central claim that the conditional loss aligns the surrogate gradient with the reconstruction gradient is therefore not established by the derivation as written.","section":"§3.4, Eq. (15) and Algorithm 1"},{"comment":"The paper states in the abstract that it 'prove[s] that the proposed loss used by existing approaches is suboptimal,' but the argument in Section 3.4 is an approximate, heuristic derivation resting on the diagonal truncation and on the small-perturbation expansion of Eq. (6). No formal theorem is stated, and no verified set of conditions identifies when one surrogate is suboptimal. Please either replace 'prove' with a weaker claim clearly labeled as an analysis/heuristic, or supply a precise theorem with assumptions and proof.","section":"Abstract and §3.4"},{"comment":"The equivalence between minimizing reconstruction loss in data space and minimizing ||I_d − N||_F^2 in latent space is approximate and inherits the bound in Eq. (11), which depends on δ = ||E_D||_F + ||E_d||_F. During training, before the approximate-inverse regime is reached, δ need not be small, and no empirical check of its magnitude is reported. Because this equivalence is the bridge to Eq. (14) and hence to the conditional-loss derivation, the paper should specify the regime in which the approximation is valid and, ideally, verify it experimentally (e.g., by reporting δ over the course of training).","section":"§3.3, Eq. (13)"}],"minor_comments":[{"comment":"The affiliation 'University of Hildeshiem' appears to be misspelled; it should be 'Hildesheim'.","section":"Title page"},{"comment":"The mask notation 1[condition] is used in line 10 but is not defined in the main text; please define it explicitly as an indicator function.","section":"Algorithm 1"},{"comment":"The caption reads 'dim 1 and 2 corresponds to Two-moons'; the verb should agree with the plural subject ('correspond').","section":"Figure 3 caption"},{"comment":"The SiLU-based results are marked as 'not considered for direct comparison,' which is clear, but the main text could state more explicitly which of these entries form the basis for the headline claims, since some tables include both ReLU and SiLU NAE results.","section":"Tables 8 and 11"}],"recommendation":"major_revision","confidential_remarks":"The empirical contribution is substantial and the paper is likely to be of interest to the journal's audience. However, the theoretical claim of a proof is currently overreaching, and the specific diagonal-truncation/single-probe gap should be resolved or the claims reframed. I would be receptive to a revised version that addresses these points; the experimental protocol and comparisons are otherwise sound."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Mate, here's my read on NAE (arXiv:2608.12084). The empirical story is genuinely strong: NAE beats FIF/FFF across molecular, tabular, and image benchmarks, often by a wide margin, while adding minimal computational overhead. The idea of dynamically choosing between encoder and decoder surrogates based on local Jacobian behavior is new and sensible, and the toy experiments measuring alignment ratios are a nice diagnostic. If the results hold up to reproduction, this could become the default loss for flow autoencoders.\n\nBut the theory is not as solid as the abstract suggests. Section 3.4 derives ∇Lrec and then explicitly truncates the double sum to the diagonal terms (i=j). The conditional mask in Algorithm 1 then selects by sign(c_ii), which is exactly sign(v_i^T(JG−I)v_i). The paper gives no argument that off-diagonal terms v_i^T(N−I)v_j and v_i^T(∇N)v_j are negligible. With K=1, you're using a single random probe, so the mask is based on one diagonal entry while the sum has order d^2 terms. This doesn't invalidate the empirical results, but it does mean the 'proof' that prior losses are suboptimal and the claim that the conditional selection aligns with reconstruction gradients are not established. The derivation is a heuristic; it should be presented as such. The reader's stress-test note on this is correct.\n\nOther soft spots: no code is released, which matters for a method whose main evidence is empirical. Baselines are mostly borrowed from prior papers, and for QM9 they changed batch size and Adam beta2 relative to FFF, so the comparison is not perfectly controlled. The ablation on Power is good, though. The SiLU modification is reasonable, but it makes the comparison to ReLU-based baselines less direct. Also, the paper duplicates the observation that denoising reconstruction regularizes the Jacobian without citing earlier work (e.g., denoising autoencoders and score matching literature); that's a minor omission but worth fixing.\n\nThe central empirical claim—that NAE improves NLL/FID across many benchmarks—looks credible and is more than enough to justify sending this to reviewers. I'd recommend conditional acceptance with a required revision to tone down the theory claims and add a diagnostic showing off-diagonal terms are small, or at least acknowledging the truncation. The paper deserves a serious referee.","headline":"Strong empirical paper with a heuristic theory; the conditional loss works, but the proof of alignment doesn't survive contact with off-diagonal terms.","tokens_in":22916,"tokens_out":3151,"would_cite":true,"duration_ms":25300,"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":"Flow autoencoders should follow reconstruction-loss gradients—NAE does this and sets new state-of-the-art results on molecular, tabular, and image benchmarks.","keywords":["normalizing autoencoder","flow autoencoder","surrogate gradient","reconstruction loss alignment","injective flows","Hutchinson trace estimator","conditional loss","generative modeling"],"falsifier":"On a trained NAE, compute the full probe expansion of $\\nabla\\|I_d - JG\\|_F^2$ on held-out data and compare the diagonal sum to the off-diagonal sum, or compute the cosine similarity between the conditional-loss gradient and the exact reconstruction-loss gradient on random batches; a frequently negative similarity or a large off-diagonal contribution would falsify the proposed mechanism.","tokens_in":21978,"feed_emoji":"🎯","tokens_out":6416,"duration_ms":56354,"temperature":0.7,"pith_summary":"Flow autoencoders train separate encoder and decoder networks as approximate inverses while optimizing a likelihood term through cheap surrogate gradients. This paper argues that the standard practice of using just one surrogate is suboptimal, because within each Hutchinson probe subspace one surrogate's gradient aligns with the reconstruction-loss gradient while the other opposes it. The authors derive that the reconstruction-loss gradient contains both surrogate terms weighted by a scalar measuring whether the composite Jacobian is locally expanding or contracting the probe direction. They propose Normalizing Autoencoder (NAE), which per probe selects the aligned surrogate via a conditional mask, and report improved negative log-likelihood on molecular benchmarks and improved FID on tabular and image benchmarks. A sympathetic reader would take the core claim to be that optimizing encoder and decoder surrogates in alignment with reconstruction loss is what makes flow autoencoder training stable and effective.","feed_headline":"Loss alignment makes flow autoencoders state of the art","feed_subtitle":"NAE chooses, probe by probe, the surrogate whose gradient matches reconstruction loss, improving likelihood and FID.","key_machinery":"The load-bearing object is the composite Jacobian $M = G_\\phi J_\\theta$ in data space and its latent-space counterpart $N = J_\\theta G_\\phi$. The reconstruction loss under isotropic noise is shown to approximately equal the squared deviation $\\|I_d - N\\|_F^2$, and its gradient is expanded in the scaled orthonormal Hutchinson probe basis $V$. The machinery is the scalar prefactor $c_{ii} = v_i^\\top(JG - I_d)v_i$ for each probe: it measures whether the composite map is locally expansive ($c_{ii} > 0$) or contractive ($c_{ii} < 0$) along that probe direction, and it determines which of the two surrogate terms, $-v_i^\\top J\\,SG(Gv_i)$ or $SG(v_i^\\top J)Gv_i$, is aligned with the reconstruction-loss gradient. Algorithm 1 uses this sign information to set a per-probe mask and build the conditional surrogate loss.","core_discovery":"On the paper's own terms, the central discovery is that the reconstruction loss already encodes the correct training signal for both networks: under isotropic input perturbations, minimizing reconstruction loss is approximately minimizing $\\|I_d - JG\\|_F^2$, and the gradient of this term decomposes over probe vectors into the encoder and decoder surrogate gradients scaled by $c_{ii} = v_i^\\top(JG - I_d)v_i$. When $c_{ii} > 0$ the decoder surrogate aligns with the reconstruction gradient and the encoder surrogate opposes it; when $c_{ii} < 0$ the roles reverse. Existing methods commit to one surrogate and therefore systematically fight the reconstruction objective in a subset of directions. NAE instead computes both surrogates per probe and, with the mask $\\mathbb{1}[\\mathrm{surr}_D > s_k]$, keeps only the term whose gradient is aligned. The paper reports that this conditional loss reaches the approximate-inverse regime at lower reconstruction weight $\\beta$, stabilizes training, and achieves state-of-the-art negative log-likelihood on DW4, LJ13, LJ55, and QM9, best FID-like scores on Power, HEPMASS, and MiniBooNE, and best FID on CelebA in most settings.","pith_inferences":["The diagonal-dominance assumption in Section 3.4 could be checked directly by measuring the magnitude of off-diagonal terms $v_i^\\top(N - I_d)v_j \\cdot v_i^\\top(\\nabla N)v_j$; if they are not negligible, a weighted combination of both surrogates with coefficient $c_{ii}$ might outperform the hard mask.","The same alignment argument transfers to any approximate-inverse training objective beyond flows, such as regularized autoencoders or learned forward/inverse maps in inverse problems, wherever a reconstruction loss and a surrogate likelihood-like term coexist.","The empirical finding that larger latent dimensions need larger $\\beta$ suggests an automatic scheduling rule: scale $\\beta$ with the latent dimension or with the intrinsic dimensionality of the data instead of tuning it per dataset."],"forward_implications":["Training with the conditional loss instead of the encoder surrogate alone should lower negative log-likelihood on any flow-autoencoder setup, not just the benchmarks reported, because it removes directions in which the surrogate opposes reconstruction.","The reconstruction weight $\\beta$ needed to reach the approximate-inverse regime is smaller with NAE, reducing sensitivity to this hyperparameter.","Because NAE and Free-Form Flows share architecture and inference procedure, the improved likelihoods on DW4, LJ13, LJ55, and QM9 come at no additional sampling cost.","The near-equal FID between standard-normal and GMM samplers on CelebA suggests the latent space learned by NAE is regular enough to sample with a plain Gaussian prior.","SiLU activations, whose second derivative is nonzero, preserve second-order surrogate gradients and further improve results; this points to activation choice as a first-class component of Jacobian-based training."],"supporting_citations":[{"why":"Introduces flow autoencoders, the encoder/decoder surrogate trace estimates, the reconstruction penalty, and the injective-flow baselines that NAE builds on and compares against.","marker":"Sorrenson et al. (2023)"},{"why":"Extends the surrogate approach to full-dimensional flows; NAE adopts its E-FFF architecture and molecular evaluation protocol as the direct baseline.","marker":"Draxler et al. (2024)"},{"why":"Provides the FID-like metric for tabular evaluation and the Rectangular Flows baseline.","marker":"Caterini et al. (2021)"},{"why":"Supplies the Pythae benchmark, architectures, and training protocol used for the image experiments.","marker":"Chadebec et al. (2022)"},{"why":"Supplies the MCMC training samples and equivariant flow-matching baselines for the Boltzmann generator benchmarks.","marker":"Klein, Krämer, and Noe (2023)"},{"why":"Provides the equivariant diffusion baseline and its reported stability and sampling-time numbers on QM9.","marker":"Hoogeboom et al. (2022)"}],"fun_headline_variants":["Aligning loss gradients lifts flow autoencoders to SOTA","NAE: Conditional loss aligns surrogates for better generative models","Reconstruction-aligned surrogates fix flow autoencoder training","Autoencoder flow gains from gradient-aligned conditional loss","Surrogate alignment: key to state-of-the-art flow autoencoders"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The conditional mask in Algorithm 1 is valid only if the diagonal terms of the probe expansion dominate and the sign of $\\mathrm{surr}_D - s_k$ correctly identifies gradient alignment; otherwise the selected surrogate can oppose the true reconstruction gradient and the method's motivation collapses.","fun_headline_variants_meta":{"raw":{"variants":["Aligning loss gradients lifts flow autoencoders to SOTA","NAE: Conditional loss aligns surrogates for better generative models","Reconstruction-aligned surrogates fix flow autoencoder training","Autoencoder flow gains from gradient-aligned conditional loss","Surrogate alignment: key to state-of-the-art flow autoencoders"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000162,"raw_usage":{"total_tokens":1239,"prompt_tokens":943,"completion_tokens":296,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":559,"completion_tokens_details":{"reasoning_tokens":209}},"tokens_in":559,"tokens_out":296,"duration_ms":2857,"temperature":1.0,"reasoning_tokens":209,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T00:17:23.835899+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"On a trained NAE, compute the full probe expansion of $\\nabla\\|I_d - JG\\|_F^2$ on held-out data and compare the diagonal sum to the off-diagonal sum, or compute the cosine similarity between the conditional-loss gradient and the exact reconstruction-loss gradient on random batches; a frequently negative similarity or a large off-diagonal contribution would falsify the proposed mechanism.","supporting_citations":[],"review_version":1}