SAI
← All ICML 2026 orals

Towards Sub-second Biological Foundation Model Infrastructure: A Quantized Consistency Diffusion Framework for Molecular Docking

Kexin Zhang, Weichen Qin, Yue Teng, Jiale Yu, Yuanyuan Ma, Jinyu Lin, Liping Sun, Jie Zheng, Jingyi Yu

OralReplication score 4%Paper PDFCode repoOpenReview

Towards Sub-second Biological Foundation Model Infrastructure: A Quantized Consistency Diffusion Framework for Molecular Docking

SAI replication review · Referee report

Summary

The paper reframes molecular docking as a sub-second structural primitive for agent-centered drug discovery by co-designing three parts of a co-folding pipeline: a streamlined MSA-free Pairformer–DiT backbone with a single recycling pass, Progressive Consistency Regularization (PCR) for reliable few-step generation, and Residual-Safe Quantization that keeps residual streams and geometry-sensitive operations in BF16 while quantizing selected compute-heavy linears. The conceptual move — treating few-step molecular generation as a trajectory-geometry problem, and treating quantization as a pathway-selective operation rather than a uniform compression — is well-scoped. The root-NSR alignment in Section 3 / Appendix A is a helpful scaffold: it makes precise what is genuinely different between EDM-style denoising and Linear (rectified-flow) transport (the target and the update rule) versus what can be normalized away. The paper's second methodological contribution — an explicit information-level-matched evaluation protocol — is timely and, in principle, an improvement on the way co-folding docking is currently benchmarked. The headline numbers (0.17 s for 5 conformations on H20, >300x vs AlphaFold3, 83.18% top-5 SR on PDBBind2020 under an interface-informed protocol) would be exciting if they hold.

Three tensions weaken the current evidence. First, the SOTA-accuracy configuration and the sub-second configuration are different systems — 83.18% is at NFE=20 in high precision on 25 candidates, 0.17 s is Linear-PCR at NFE=5, W8A8+JIT on 5 conformations — and the deployed system's SR is never reported. Second, Table 1 mixes information-level regimes and bolds the top row, exactly the apples-to-oranges ranking the paper's own principle warns against; the blind claim is asserted without a matched-baseline table. Third, when the code was actually executed here, essentially every distinguishing contribution (PCR, Linear sampler, W8A8, JIT, latency harness, four-level driver, torsion metric, PDBBind2020 split, confidence-based ranking) was missing or fell back to a different benchmark, so the headline latency was never measured and headline SR was only exercised on an off-distribution PoseBusters subset. The intellectual contribution stands; the empirical scaffolding does not yet.

Strengths

  • Clear conceptual framing. The paper is unusually explicit about why sub-second docking matters (agent-centered loops, screening-scale callability) and structures its methods against the specific latency sources (recycling, sampling NFE, per-step arithmetic). This gives the accuracy–latency co-design a coherent narrative.
  • Root-NSR analytical scaffold. Section 3.1 and Appendix A provide a careful common Gaussian-interpolation coordinate that isolates what is genuinely different between EDM, VP, cosine, and Linear parameterizations. This lets the paper make a substantive rather than rhetorical argument that Linear parameterization is more discretization-friendly under coarse solvers.
  • Explicit information-level protocol. Making the amount of structural prior an explicit axis of evaluation (blind -> noisy pocket -> exact pocket -> interface-informed) is a genuine methodological contribution to a subfield that has been sloppy about this axis. Table 4 and Appendix D.2 give a clean picture of how much of docking difficulty is set by input priors.
  • Residual-Safe Quantization design principle. The decision to keep residual streams, triangle operations, softmax, normalization, and structure heads in BF16 while quantizing only compute-heavy inner linears is well-motivated by geometry sensitivity, and the reported robustness of global pose under W8A8 (with χ14\chi_{1-4}-MAE degrading first) matches the paper's own theory.
  • Careful protocol reporting. Sample counts (N=25 for SR, N=5 for latency), timed region (excludes preprocessing and JIT compile), and hardware (single H20) are all stated clearly, which is better than most co-folding papers.
  • Honest negative results. Appendix D reports thermodynamically-informed priors, 3D-RoPE coordinate modulation, adaptive modulation projection, and Muon as ablations that did not help; these are correctly interpreted as motivating a simpler design rather than dressed up as improvements.

Weaknesses

  • Deployment configuration is not the SOTA configuration. The abstract, contributions, and conclusion frame the paper as delivering both state-of-the-art accuracy and sub-second inference, but Table 1's 83.18% is the diffusion-pretrained model at NFE=20 in high precision on N=25 candidates, while 0.17 s is Linear-PCR at NFE=5 with W8A8 quantization on 5 conformations. The SR of the actual sub-second deployment (Linear-PCR, NFE=5, W8A8) is never reported. The reader is asked to hold both endpoints in mind as if they were the same system.
  • Accuracy and latency protocols use different candidate counts. SR uses N=25, latency uses N=5. Under Eq. (27) the accuracy-matched latency would be materially larger than 0.17 s. Either the SR should be reported at N=5 or the latency at N=25, so the two headline numbers describe the same pipeline.
  • Table 1 violates the paper's own principle. Blind (Vina, DiffDock, DiffDock-L, Uni-Mol V2) and interface-informed (SurfDock, the paper's rows) methods are placed in a single ranked SR column with the top (interface-informed) row bolded — exactly the cross-protocol ranking Section 5.3 warns against.
  • Blind-docking claim unquantified. 'Remains competitive under blind docking' is repeated in the abstract, introduction, and contributions, but no matched blind-docking comparison table is provided. Table 4 gives 52.6% at Level 1 for the paper's model, without any Level-1 baselines.
  • >300x speedup is not decomposed. The headline speedup bundles R=1 recycling, PCR NFE=5, and W8A8+JIT against baselines matched to the paper's protocol only 'whenever supported'. The reader cannot see how much of the frontier is driven by the two new methods versus R=1 alone, or which baselines were and were not matched.
  • No per-seed variance or CIs on SR. Linear (83.18%) vs EDM (82.48%) vs SurfDock (81.07%) gaps map to a handful of targets out of 428. Without per-seed variance or a Wilson interval on the binomial SR, the ranking cannot be judged as statistically meaningful.
  • Table 4 SR criterion under-specified. Table 4 states only 'RMSD < 2.0 Å' without top-1 vs top-5, sample count, or pocket alignment. This blocks direct comparison to Table 1 and prevents the Level-4 number from serving as the Table-1 interface-informed row.
  • Scaling limitation stated ambiguously. 'Does not grow linearly with the number of input tokens' is directionally ambiguous. The intended message is super-linear (O(N2)O(N^2) pair, O(N3)O(N^3) triangle) — a stronger and more accurate limitation.
  • Muon 'smoother trajectory' inference over-claimed. Lower raw gradient norms under Muon vs AdamW cannot support a smoothness claim because Muon rescales updates; the appendix reads the norm difference too causally.
  • 'Validates' overstated. Table 4's monotonic trend supports the empirical premise behind level-matched evaluation, but does not validate the methodological principle itself.
  • Framework naming is confusing. The proposed system has no name; the released repository is called PhysDock, referring to the earlier Zhang et al. 2025 work. It is not clear which components are new here (PCR, Residual-Safe Quantization, R=1) and which are inherited, which matters because the released code contains only the inherited pieces.
  • Layout scramble in Section 4.1 / Preliminaries. The Section 3 deterministic-update material is interleaved with Section 4.1's cost-equation exposition, so 'Eq. (27)' sits next to unrelated equations. Multiple display equations render as blank in the manuscript, which makes several load-bearing claims hard to verify from prose alone.
  • Noisy-pocket distribution unspecified. The 6 Å radius shift for Level 2 does not identify the distribution (uniform-in-ball vs uniform-on-sphere vs truncated Gaussian), yet Table 4's 59.1% depends on this choice.

Reproducibility & code

The executed replication attempt did not confirm any headline number, and the released artifact is much thinner than 'part of the implementation' suggests. Grounding in the veritas run:

  • Headline latency and speedup were never measured. A harness matching the paper's protocol was authored (bench_latency.py, 100 timed runs after warmup, CUDA-synchronized, 256-token / 5-conformation, excluding preprocessing and compile), but the orchestrator terminated after step 3 and the benchmark never ran; no latency*.json exists in the run outputs. The available hardware was also an A100 80GB rather than the paper's H20, but this is moot because no value was produced. Baseline latency scripts for AlphaFold3, Chai-1, Protenix, Interformer, SurfDock, and DiffDock are also absent from the released code, so the >300x speedup cannot be reconstructed even with a working Ours harness.
  • Headline accuracy is not runnable on the paper's benchmark. The repository contains no PDBBind2020 data or split manifest; it ships PoseBusters, DeepDockingDare, and PhiBench instead. The one interface-informed configuration that completed (EDM, NFE=20, 64-system PoseBusters subset with pocket + key-residue + physics-correction + KMeans ranking) scored 95.31% — well above the paper's 82.48% and on a different, easier benchmark, so it does not confirm the paper's number. The Linear variant is not runnable at all because the Linear sampler is absent from the released code.
  • Table 4 is only partially runnable. In the executed reproduction, only Levels 1 and 4 could be attempted, on 32/64 PoseBusters systems respectively; both diverged materially from the paper (blind 62.5% vs 52.6%, interface-informed 95.3% vs 82.48%). Levels 2-3 could not be exercised because the noisy-pocket sampler and the four-level driver are not scripted.
  • Distinguishing contributions absent from the artifact. The released code has no PCR consistency loss, no EMA teacher θ\theta^{-}, no progressive M(k)M(k) schedule, no Linear parameterization (interpolation, velocity target, Euler transport), no INT8 weight/activation quantization or activation calibration, no dequantization-before-residual routine, no JIT/torch.compile/TensorRT toolchain, no χ14\chi_{1-4}-MAE evaluator, no pocket-residue extractor under the 10 Å heavy-atom criterion, and no NFE-sweep driver. Every panel of Figure 2 and every bar of Figure 3 depends on one or more of these.
  • SR protocol relies on a disabled module. The confidence module is present in source but explicitly marked 'not used in released version' and its import is commented out. The shipped redocking.py samples up to max_samples=5 (not 25) and ranks samples by KMeans over inter-sample RMSD distances that require the ground-truth structure — so even for the pieces that do run, the paper's top-5-by-confidence SR is not implementable from the artifact as-is.
  • Under-specified hyperparameters. The PCR curriculum M(k)M(k) has only endpoints stated (10 -> 1000), not a functional form; the EMA decay for θ\theta^{-} is not stated; the PCR-stage optimizer settings (Adam betas/epsilon, whether the 100k-step tuning stage inherits the pre-training schedule) are not stated; the PTQ backend is not named; and the calibration-set size, split, and scale-estimation method are not specified. These are load-bearing for the W8A8 SR and χ14\chi_{1-4}-MAE numbers.

Recommended Changes

Essential

  • Report SR of the actual sub-second deployment. Add a row (or a Figure 3 companion table) giving docking SR for Linear-PCR at NFE=5 with W8A8+JIT, generating 5 conformations, so that the headline accuracy and headline latency describe the same system. Addresses the Deployment-configuration-is-not-SOTA-configuration weakness.
  • Split Table 1 by information level and drop cross-protocol boldface. Present matched-protocol sub-tables (blind vs interface-informed) so that the strongest row is not a cross-regime rank. Addresses the Table-1-violates-paper's-own-principle weakness.
  • Add a matched blind-docking comparison table. Provide numeric comparison against blind DiffDock / DiffDock-L / blind AlphaFold3 (and any co-folding baseline configurable to blind) under the Level-1 protocol, so 'remains competitive' becomes evidence-grounded. Addresses the Blind-docking-claim-unquantified weakness.
  • Decompose the >300x speedup. Add a per-intervention waterfall on the paper's own model: baseline -> R=1 -> +PCR NFE=5 -> +W8A8+JIT, and state per-baseline whether AlphaFold3, Chai-1, Protenix, Interformer, SurfDock, DiffDock were actually matched to the 5-conformation, 256-token setting or estimated. Addresses the >300x-not-decomposed weakness.
  • Reconcile the SR (N=25) and latency (N=5) protocols. Report SR at N=5 or latency at N=25 in the headline, so the two numbers describe the same pipeline. Addresses the mismatched-candidate-counts weakness.
  • Release the missing components. Ship the PCR loss (consistency term with EMA teacher, progressive M(k)M(k) schedule), the Linear-parameterization sampler and NFE-sweep driver, the Residual-Safe Quantization pass (weight quantization, activation calibration, per-layer selection), the JIT deployment path, the latency harness, the baseline invocation scripts, the four-level information-conditioning driver (including the noisy-pocket sampler), the χ14\chi_{1-4}-MAE evaluator, the curated PDBBind2020 428-complex split as a PDB-ID manifest with a curation script, and enable the confidence module and its training loss. Addresses every bullet in Reproducibility & code.
  • Report per-seed variance or Wilson CIs on SR. With 0.7-pp gaps on a 428-target denominator, provide per-seed standard deviations across the sample-and-rank randomness and a Wilson binomial interval on each headline SR. Addresses the No-per-seed-variance weakness.

Suggested

  • Specify the PCR curriculum M(k)M(k), EMA decay for θ\theta^{-}, and PCR-stage optimizer schedule. Render the omitted equations and state whether the 100k-step PCR stage inherits the pre-training schedule. Addresses the PCR-hyperparameters-under-specified point.
  • Name the quantization backend and calibration protocol. State the backend (e.g., torch.ao, TensorRT-LLM, custom), the calibration-set source and size, and the scale-estimation method (min-max, percentile, MSE). Addresses the PTQ-backend-unnamed point.
  • Pin down the noisy-pocket distribution. Specify whether the 6 Å shift is uniform-in-ball, uniform-on-sphere, or truncated Gaussian, plus the RNG seeding convention. Addresses the noisy-pocket-distribution weakness.
  • Specify the Table 4 SR criterion. State top-1 vs top-5, sample count per target, and whether pocket alignment was applied, so Table 4 can be read alongside Table 1. Addresses the Table-4-SR-criterion-under-specified weakness.
  • State scaling as O(N2)O(N^2) / O(N3)O(N^3). Replace 'does not grow linearly' with the specific super-linear complexity. Addresses the Scaling-limitation-stated-ambiguously weakness.
  • Delineate PhysDock vs this-paper contributions. Add a short paragraph naming which components (backbone, EDM sampler, MMFF, dataset scripts) are inherited from PhysDock and which are novel here. Addresses the Framework-naming-is-confusing weakness.
  • Soften the 'validates' claim in D.2. Change to 'motivates' or 'necessitates'. Addresses the 'validates'-overstated weakness.
  • Repair the Section 3 / 4.1 layout. Move the Linear/EDM deterministic-update material back into Section 3.1 and render the blank equations, so 'Eq. (27)' can be located unambiguously. Addresses the Layout-scramble weakness.
  • Temper the Muon interpretation. Note that raw gradient-norm magnitudes across differently-normalized optimizers do not straightforwardly imply landscape smoothness. Addresses the Muon-interpretation weakness.
  • Restore the level-matched qualifier in the conclusion. Replace 'preserving high docking accuracy' with the SOTA-under-interface-informed / competitive-under-blind phrasing used earlier. Addresses the Conclusion-overstates-blind-accuracy weakness.