{"id":"7c83d3da-9966-4a57-ab54-1802589b01ee","arxiv_id":"2501.19077","paper_version":2,"verdict":"ACCEPT","confidence":"HIGH","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":11,"one_line_summary":"High-temperature reverse-KL pretraining followed by reweighting-based annealing lets normalizing flows sample Boltzmann distributions of peptides up to hexapeptide without mode collapse and with fewer target energy evaluations.","lead":"Researchers trained normalizing flows to sample molecular shapes by first learning at very high temperature, then gradually cooling the model with a reweighting scheme. The method, TA-BG, sampled three peptide systems more accurately than prior baselines and used up to three times fewer energy evaluations on the smaller systems.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Training objective omits the internal-coordinate Jacobian, so the target density is not the Boltzmann distribution as defined; this must be clarified or fixed.","rationale":"The reader's weakest assumption concerns mode-collapse generalization to new systems, which is a practical limitation but does not threaten the reported results on the three alanine peptides, where the critical-temperature ablation (Fig. 6) shows no visible collapse at 1200 K. A more fundamental issue is that the manuscript's equations do not define the target density in the internal-coordinate space used by the flow. Section 3.1 states p_X(x) is the Boltzmann distribution with x being Cartesian coordinates, while Section 4.1 says the flow uses internal coordinates. Eq. 6 then evaluates p_X(g(z)) without any Jacobian correction for the coordinate transformation. For flexible bond lengths and angles, this Jacobian is not constant, so a literal reading of the paper gives a training objective that does not target the correct Boltzmann distribution in internal coordinates. This is a correctness risk for the central claim, not merely a generalization concern. I give credit for the strong empirical evaluation, released code and data, and the careful ablations; these make it likely that the implementation includes the Jacobian through the bgflow library even though the text omits it. The concern is therefore testable and fixable, warranting a conditional acceptance rather than a rejection. The concrete test of inspecting the code and, if needed, re-running the dipeptide experiment with the Jacobian term will settle whether the published claims are reproducible as written or whether the equations must be corrected.","tokens_in":27546,"tokens_out":15433,"duration_ms":150900,"concrete_test":"Inspect the released code (github.com/aimat-lab/TA-BG) to verify whether the target log-density in the reverse-KLD loss and in the importance weights includes the Jacobian determinant of the internal-coordinate transformation. If it is absent, add the term and rerun the alanine dipeptide experiment with the same hyperparameters to see whether NLL and RAM-KLD change materially. If the code already includes the Jacobian, amend Section 4.2 to define the internal-coordinate target density explicitly, which would resolve the concern.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"The paper's training objective (Eq. 6) is written for a target p_X(x) defined on Cartesian coordinates (Section 3.1), but the flow operates on internal coordinates (Section 4.1). The target density on the flow's support must include the Jacobian determinant |∂x_cart/∂x_int| of the Cartesian-to-internal transformation: p_int(r) ∝ exp(−E(x_cart(r))/kT) |det J(r)|. This Jacobian is never defined or mentioned anywhere in the manuscript, even though it is not constant when bond lengths and angles are flexible. If the implementation follows Eq. 6 literally, the reverse-KLD pretraining and the reweighted annealing both target the wrong distribution, so the central claim of accurate Boltzmann sampling is not supported. The released code may include this term via bgflow, but the manuscript as written is ambiguous or incorrect; a reader implementing from the text would obtain a biased sampler.","agreement_with_reader":"disagree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes temperature-annealed Boltzmann generators (TA-BG), a variational sampling method that first trains a normalizing flow with the reverse Kullback-Leibler divergence at an elevated temperature (1200 K) to avoid mode collapse, and then anneals the learned distribution to a target temperature (300 K) through iterative importance-weighted resampling followed by forward-KLD training. The method is evaluated on alanine dipeptide, tetrapeptide, and hexapeptide in implicit solvent, where it is compared against flow models trained with forward KLD on MD data, reverse KLD at 300 K, and the FAB baseline. The authors report that TA-BG matches or improves on FAB on most metrics while using up to about three times fewer target energy evaluations, and that for the hexapeptide it is the only variational method that accurately resolves the metastable states. The paper also includes extensive ablations, a 2D Gaussian mixture study, and a public implementation and ground-truth datasets.","tokens_in":27696,"tokens_out":5931,"duration_ms":61431,"significance":"If the central claims hold, this is a practically useful contribution to data-free variational sampling of molecular systems. The main idea is simple and plausible: high-temperature reverse-KLD training avoids mode collapse, and a sequence of short annealing steps with reweighting transfers the learned density to the target temperature. The empirical evaluation is careful and unusually thorough: four independent runs per condition, standard errors, hyperparameter ablations for starting temperature, annealing schedule, buffer size, and fine-tuning, a robustness check for FAB, and a fair comparison of wall-clock time in Appendix J. The release of code and ground-truth data is a concrete asset. The main reservation is that the mathematical formulation of the target density in internal-coordinate space is incomplete, which affects the central claim of accurate Boltzmann sampling; this is a fixable issue but must be resolved before the paper can be accepted.","major_comments":[{"comment":"The training objective is written for a target density p_X defined on Cartesian coordinates, while the flow operates on internal coordinates (bond lengths, angles, dihedrals). The correct target density in internal coordinates is p_R(r) ∝ exp(−E(x_cart(r))/k_B T) |det(∂x_cart/∂r)|, and this Jacobian is not constant because bond lengths and angles are flexible. The Jacobian is never defined or mentioned in the paper. As written, both the reverse-KLD pretraining and the importance weights in Section 4.2 target a different density, so the central claim that the method samples the Boltzmann distribution is not supported by the equations presented. Please state explicitly what density the flow is trained to match, include the Jacobian factor in the target density if the internal-coordinate frame is used, and document how the released implementation handles this term.","section":"Section 3.1 / 4.1, Eq. (6)"}],"minor_comments":[{"comment":"The sentence 'For FAB applied to the hexapeptide, even when using almost 3 times as many target evaluations compared to our approach, we still achieve a lower NLL value' is ambiguous; it should be rephrased to make clear that TA-BG achieves the lower NLL, while FAB uses more evaluations.","section":"Section 6"},{"comment":"Mode collapse in the starting-temperature ablation is defined by 'manual visual inspection of the Ramachandran plots'; please specify a more objective criterion or at least note the possible subjectivity of this threshold.","section":"Appendix F.3, Figure 6"},{"comment":"The statement 'mode collapse is not a problem during the annealing' is too strong: forward-KLD training can only fit the support represented by reweighted samples from the current flow, so a mode missed by the high-temperature pretraining cannot be recovered. The limitation is acknowledged in Appendix F.3, but the main text should state this caveat.","section":"Section 4.2"},{"comment":"The wall-time comparison is useful, but the main-text efficiency claim ('up to three times fewer target energy evaluations') should be explicitly distinguished from total wall-clock time, since the appendix shows that FAB currently has lower wall time on these systems.","section":"Appendix J, Table 14"}],"recommendation":"major_revision","confidential_remarks":"The empirical work is strong and the central idea is attractive, but the missing internal-coordinate Jacobian is a load-bearing mathematical omission. I would not reject: if the authors confirm that the implementation includes the Jacobian and add the missing mathematical statement to the paper, the manuscript would be acceptable. The paper's novelty relative to existing temperature-conditioned flows is modest but the empirical demonstration on the hexapeptide is valuable."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Short version: the two-phase idea here is new and works. Training a normalizing flow with the reverse KLD at 1200 K to dodge mode collapse, then annealing to 300 K with reweighted forward KLD, is a clean combination I haven't seen in the prior FAB/replica-exchange work. The hexapeptide result is the headline: FAB leaves part of the metastable region distorted; TA-BG resolves it. The evaluation is careful — four independent runs per condition, standard errors, ablations on starting temperature, schedule, and fine-tuning, plus honest wall-time numbers. Code and ground-truth data are public. That part earns the credit.\n\nThe real soft spot is in the written math. The training objective (Eq. 6) is stated for a target p_X on Cartesian coordinates, but the flow operates on internal coordinates. The target density in internal coordinates needs the Jacobian determinant of the Cartesian-to-internal map, and that term is absent from the manuscript. A reader implementing from the text would target a biased sampler. The agreement with MD ground truth suggests the released code handles it through bgflow, but the paper has to say so. This is fixable, but it's exactly what a careful referee needs to check against the code.\n\nSecondary concerns: no theoretical guarantee that the reweighted forward KLD anneals to the target; the method leans on the empirically monitored buffer ESS. Hyperparameters (starting temperature, iteration count, fine-tuning) are partially selected on the evaluation systems, though the critical-temperature ablation in the appendix gives some cover. The importance-sampling scaling problem is real and is acknowledged, with an AIS sketch for future work. None of this changes the central empirical claim for the three systems studied.\n\nBottom line: anyone working on data-free variational sampling for molecules should read this. It deserves a serious referee. Send it to review, with an explicit note to verify the Jacobian term against the code. If that checks out, it's an accept.","headline":"A genuinely new two-phase training strategy that delivers the best hexapeptide sampling I've seen, but the written objective omits the internal-coordinate Jacobian and needs a fix before the math matches the code.","tokens_in":28314,"tokens_out":4999,"would_cite":true,"duration_ms":45747,"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":"Training a flow at 1200 K, then annealing it to 300 K by reweighting its own samples, captures all metastable peptide states at a fraction of the leading baseline's energy cost.","keywords":["temperature-annealed Boltzmann generators","normalizing flows","reverse Kullback-Leibler divergence","mode collapse","importance sampling reweighting","Boltzmann distribution sampling","metastable states","alanine peptides"],"falsifier":"Run the TA-BG recipe on a molecule whose metastable basins remain separated by a free-energy barrier of several $k_BT$ even at 1200 K, with the true 300 K populations fixed by a long unbiased molecular-dynamics trajectory. Apply the paper's own starting-temperature protocol — the fraction of collapsed reverse-KLD runs as a function of $T_1$ — to that system: if the pretraining collapses at 1200 K, or the annealed 300 K distribution assigns near-zero weight to a basin that the trajectory visits, the claim that temperature annealing circumvents mode collapse fails for that regime. A quantitative companion check is the final negative log-likelihood and Ramachandran KLD on an independent test set compared with the FAB baseline at a matched target-energy budget.","tokens_in":27269,"feed_emoji":"🧪","tokens_out":21479,"duration_ms":171421,"temperature":0.7,"pith_summary":"The paper proposes temperature-annealed Boltzmann generators (TA-BG), a two-phase recipe for training a normalizing flow — an invertible neural transformation with an exactly computable density — to sample the Boltzmann distribution of a molecule, the equilibrium distribution of its conformations, without the mode collapse that usually afflicts energy-based training. In the first phase the flow is trained with the reverse Kullback-Leibler divergence at 1200 K, where free-energy barriers are low enough that the molecule's metastable conformations are connected and the mode-seeking loss covers all of them. In the second phase the flow is cooled to 300 K along a geometric ladder of about nine intermediate temperatures: at each rung the flow samples a buffer, every sample is reweighted to the next lower temperature, the buffer is resampled, and the flow is retrained on it with the mass-covering forward KLD. On alanine dipeptide, tetrapeptide, and hexapeptide, the method matches or beats the Flow Annealed Importance Sampling Bootstrap (FAB) baseline on nearly all metrics, while using roughly a third of the target-energy evaluations on the two smaller systems; on the hexapeptide the authors report that it is the only variational method tested that resolves all metastable states. If the claim holds, accurate equilibrium sampling of flexible molecules becomes substantially cheaper, which matters most when each energy evaluation is expensive.","feed_headline":"Train at 1200 K, sample at 300 K: flow learns all peptide states","feed_subtitle":"Annealing with reweighted samples matches the leading baseline on three peptides at a third of the energy cost.","key_machinery":"The load-bearing mechanism is a two-stage temperature protocol built on a normalizing flow. Stage one trains the flow with the reverse Kullback-Leibler divergence at an elevated temperature; the mode-seeking behaviour of that loss, which collapses the flow at 300 K, is neutralized at 1200 K because the barriers shrink and the high-probability regions interconnect, and a regularized energy function prevents diverging van der Waals terms from destabilizing training. Stage two cools the flow along a geometric temperature ladder $T_i = T_{\\mathrm{start}}\\left(T_{\\mathrm{target}}/T_{\\mathrm{start}}\\right)^{(i-1)/(K-1)}$, where each rung resamples a buffer of the flow's own samples using importance weights for the next temperature and retrains with the forward, mass-covering KLD, which is what keeps the cooling phase collapse-free. The flow itself is built from monotonic rational-quadratic spline coupling layers acting on internal coordinates (bond lengths, angles, dihedrals), with circular splines so that the periodic dihedral angles keep their correct topology. A final fine-tuning iteration at the target temperature, with $T_{i+1}=T_i$, raises the effective sample size of the training buffer and improves the final metrics.","core_discovery":"The central claim is that mode collapse in data-free normalizing-flow training is a temperature problem rather than an intrinsic failure of the reverse KLD. At 1200 K the free-energy barriers between metastable basins are low enough that reverse-KLD training covers all modes reliably, and the paper's starting-temperature ablation shows the fraction of collapsed runs dropping to zero there; at 300 K the same objective collapses, increasingly so for the larger peptides. The second claim is that an iterative reweighting anneal carries this coverage down to the target temperature: at each rung of a geometric temperature ladder, samples drawn from the flow at $T_i$ are weighted by $w(x)=p_{X,T_{i+1}}(x)/q_X(x;\\theta)$, resampled, and used for forward-KLD training at $T_{i+1}$, repeated until 300 K. On the three alanine systems the resulting 300 K Ramachandran free-energy plots match long molecular-dynamics ground truth, with better negative log-likelihoods than FAB on all three systems and $7.56\\times 10^7$ versus $2.13\\times 10^8$ target energy evaluations on the two smaller systems. For the hexapeptide, the authors state, TA-BG is the only method tested that accurately resolves the metastable states.","pith_inferences":["An automatic pretraining protocol suggests itself: scan the starting temperature upward until the fraction of collapsed reverse-KLD runs drops to zero, following the paper's Figure 6 procedure, and then begin annealing; this converts the per-system empirical heuristic into a design rule for new molecules.","The critical temperature at which reverse-KLD training stops collapsing should track the physical barrier heights of the target's free-energy landscape, so a cheap low-temperature barrier estimate could predict a safe starting temperature without retraining.","The annealing buffer's effective sample size could serve as a live thermostat: pick the next temperature on the fly to hold the buffer overlap at a target value, adapting the schedule to systems whose modes merge faster or slower than alanine peptides.","Because the anneal is driven by forward KLD on reweighted samples rather than by the target energy directly, the annealing phase should preserve whatever modes the pretraining found; a testable consequence is that swapping the pretraining objective, for instance to FAB's $\\alpha$-divergence, while keeping the anneal would leave the final mode coverage unchanged."],"forward_implications":["Reverse-KLD training, previously discounted as unavoidably mode-collapsing, is sufficient for molecular Boltzmann sampling whenever the training temperature is high enough that metastable basins merge.","The buffered reweighting step is a collapse-free fine-tuning recipe for pretrained flows, applicable beyond temperature annealing, for example to debias flows trained on biased or non-equilibrated simulation data.","When target energies become the dominant cost, as with learned foundation-model force fields or ab initio potentials, the roughly threefold reduction in energy evaluations translates into a near-proportional wall-clock saving despite a larger number of flow evaluations.","Data-free variational sampling scales past the alanine-dipeptide benchmark: on the hexapeptide the only variational method that resolves all metastable states is TA-BG, making larger and more flexible molecules plausible targets.","Inside the same annealing ladder, plain importance sampling's exponential effective-sample-size decay in high dimensions can be replaced by annealed importance sampling, keeping buffer overlap approximately constant as system size grows."],"supporting_citations":[{"why":"Supplies the FAB baseline that TA-BG is compared against, the alanine dipeptide benchmark setup, and the regularized energy function used for training.","marker":"Midgley et al., 2023b"},{"why":"Provides the rational-quadratic spline coupling layers that form the normalizing flow backbone for all experiments.","marker":"Durkan et al., 2019"},{"why":"Introduces Boltzmann generators and the importance-sampling estimator (Equation 7) that the reweighting-based annealing relies on.","marker":"Noé et al., 2019"},{"why":"Defines annealed importance sampling, used inside FAB and proposed as the route to scaling TA-BG beyond the limits of plain importance sampling.","marker":"Neal, 2001"},{"why":"Source of the geometric temperature progression (Equation 8) that sets the annealing ladder between 1200 K and 300 K.","marker":"Sugita & Okamoto, 1999"},{"why":"Contributes circular splines, which let the flow treat periodic dihedral angles with the correct topology.","marker":"Rezende et al., 2020"},{"why":"Earlier temperature-transfer flow work whose tetrapeptide force-field parameters are reused; the paper distinguishes its sample-free approach from this MD-data-dependent one.","marker":"Dibak et al., 2022"},{"why":"Supplies the effective sample size measure used to report sampling efficiency.","marker":"Martino et al., 2017"}],"fun_headline_variants":["Heat the flow, cool the sample: annealing beats mode collapse","Annealing normalizing flows: 3x fewer energy calls, full mode coverage","Start hot, finish cold: flow sampling without mode collapse","Reweighted anneal reaches rare states at a third of the cost","Hot start, cool anneal: Boltzmann generators see every basin"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that every new molecular system has some starting temperature at which reverse-KLD training reliably finds all metastable modes; the paper establishes this empirically for three alanine peptides, and if the pretraining misses a mode, the annealing steps can only reweight samples the already-collapsed flow covers and cannot recover it.","fun_headline_variants_meta":{"raw":{"variants":["Heat the flow, cool the sample: annealing beats mode collapse","Annealing normalizing flows: 3x fewer energy calls, full mode coverage","Start hot, finish cold: flow sampling without mode collapse","Reweighted anneal reaches rare states at a third of the cost","Hot start, cool anneal: Boltzmann generators see every basin"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.0011,"raw_usage":{"total_tokens":4615,"prompt_tokens":998,"completion_tokens":3617,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":614,"completion_tokens_details":{"reasoning_tokens":3526}},"tokens_in":614,"tokens_out":3617,"duration_ms":24880,"temperature":1.0,"reasoning_tokens":3526,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-09T21:25:06.264319+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the TA-BG recipe on a molecule whose metastable basins remain separated by a free-energy barrier of several $k_BT$ even at 1200 K, with the true 300 K populations fixed by a long unbiased molecular-dynamics trajectory. Apply the paper's own starting-temperature protocol — the fraction of collapsed reverse-KLD runs as a function of $T_1$ — to that system: if the pretraining collapses at 1200 K, or the annealed 300 K distribution assigns near-zero weight to a basin that the trajectory visits, the claim that temperature annealing circumvents mode collapse fails for that regime. A quantitative companion check is the final negative log-likelihood and Ramachandran KLD on an independent test set compared with the FAB baseline at a matched target-energy budget.","supporting_citations":[{"cited_title":"Neural Spline Flows","cited_arxiv_id":null,"evidence_quote":"Provides the rational-quadratic spline coupling layers that form the normalizing flow backbone for all experiments."},{"cited_title":"J., Papamakarios, G., Racaniere, S., Albergo, M., Kanwar, G., Shanahan, P., and Cranmer, K","cited_arxiv_id":null,"evidence_quote":"Contributes circular splines, which let the flow treat periodic dihedral angles with the correct topology."}],"review_version":1}