All-Atom Structural Knowledge for Omni Biomolecular Design

About

Omni biomolecular models aim to support peptides, small molecules, antibodies, and other molecular classes through one common generative interface. AnewOmni does this with a shared block-level latent state. The tradeoff is that fine-grained all-atom structure is compressed before diffusion.

Core question. Can we keep the unified omni generator while adding structural knowledge learned by a stronger all-atom model?

Keep the omni latent

The shared block-level state remains the generative space.

Add all-atom knowledge

Frozen PairFormer features contribute token-level and relational structural information.

Preserve native geometry

The transfer path only changes scalar conditioning; coordinate updates remain native.

Geometry-compatible representation transfer

PairFormer receives the same visible design condition as the generator and produces final-layer single and pair representations. These representations are computed once per design task, aligned to the corresponding AnewOmni blocks, and reused throughout sampling.

Frozen donor

PairFormer

Extract single-token and pairwise all-atom structural representations.

→

Trainable interface

Cross-attention

Native scalar conditions query the donor knowledge. Pair features shape both attention and messages.

→

Frozen generator

AnewOmni

Original latent diffusion model still predicts feature noise and all 3D coordinate updates.

Why this preserves geometry

The transferred information is treated as fixed scalar side information. Rotation acts on AnewOmni’s native coordinates, not on the donor knowledge. Because the interface does not introduce a geometric direction, the augmented coordinate prediction keeps the native transformation law.

\epsilon_{\mathrm{aug},q}\!\left(Rz_t^G,t\mid Rz^C,s,p\right)=R\,\epsilon_{\mathrm{aug},q}\!\left(z_t^G,t\mid z^C,s,p\right)

Representation guidance

At inference time, the base and augmented predictions can be interpolated with a scalar guidance parameter γ.

\epsilon_\gamma=\epsilon_{\mathrm{base}}+\gamma\left(\epsilon_{\mathrm{aug}}-\epsilon_{\mathrm{base}}\right)

Guidance strength

γ = 1.00

Use the learned PairFormer-conditioned model at training strength.

Base prediction

All-atom correction

All-atom knowledge
sᵢ, pᵢⱼ

Omni generator
same latent, richer condition

Theory: conditional geometry compatibility

Assumptions. Fix the extracted PairFormer single and pair representations s and p. Native scalar conditioning a is rotation invariant. The native AnewOmni denoiser has an invariant scalar-noise prediction and a rotation-equivariant coordinate-noise prediction. Rotate generated and context coordinates jointly, including native geometric prompts where applicable.

Theorem — scalar transfer preserves native equivariance

\begin{aligned}\ell_{ij}&=\frac{(W_Qa_i)^\top W_Ks_j}{\sqrt{d}}+b_\phi(p_{ij})\\[5pt]\alpha_{ij}&=\operatorname{softmax}_j(\ell_{ij})\\[5pt]m_i&=\sum_j\alpha_{ij}\left(W_Vs_j+W_Pp_{ij}\right)\\[5pt]\widetilde a_i&=a_i+W_Om_i\end{aligned}

The transferred conditioning stays invariant. Consequently the augmented scalar prediction is invariant and the coordinate prediction is equivariant under every R ∈ SO(3).

\begin{aligned}\epsilon_{\mathrm{aug},h}(Rz\mid R\,\mathrm{context},s,p)&=\epsilon_{\mathrm{aug},h}(z\mid\mathrm{context},s,p)\\[6pt]\epsilon_{\mathrm{aug},q}(Rz\mid R\,\mathrm{context},s,p)&=R\,\epsilon_{\mathrm{aug},q}(z\mid\mathrm{context},s,p)\end{aligned}

Proof

With a, s and p fixed under rotation, the attention logits do not change. Softmax weights, projected values and their weighted sums therefore do not change. The residual conditioning ã is invariant. Substituting ã into the unchanged native denoiser invokes its original symmetry law, proving scalar invariance and coordinate equivariance.

Corollary — representation guidance

\epsilon_\gamma=(1-\gamma)\epsilon_{\mathrm{base}}+\gamma\epsilon_{\mathrm{aug}},\qquad\gamma\geq0

For a fixed scalar γ, both coordinate predictions obey the same rotation law. Their linear combination does too, including extrapolation at γ > 1.

Scope. This is conditional on fixed donor representations. It does not establish end-to-end equivariance when the raw input is rotated and PairFormer features are recomputed, nor does it prove empirical improvement.

See the theorem and detailed appendix proof in the paper, or inspect the proof source.

LNR benchmark results

Four validation-selected adapters and a historical AnewOmni reference, evaluated on the same 93 targets with 10 candidates per target. The adapters use the official generation and reconstruction procedure.

Candidate means within each target, then equal weight across targets. Lower is better except AAR. Bold indicates best; underline indicates second best.
Model	AAR ↑	Complex RMSD ↓Å	Peptide RMSD ↓Å	Intra-peptide clash ↓	Interface clash ↓	dG ↓REU	ddG ↓REU
AnewOmni · historical	6.517%	8.585	2.917	0.987%	0.877%	9.35	46.69
Single + Pair	6.435%	8.451	2.839	0.841%	0.586%	-6.84	30.50
Single-only	6.223%	8.588	2.870	0.934%	0.822%	-0.99	36.34
Pair-only	6.412%	8.542	2.858	0.921%	0.546%	-9.06	28.28
No pretrained features	6.214%	8.589	2.887	0.880%	0.879%	1.40	38.73

Summary CSV ↓Protocol & summary JSON ↗All 4,650 candidate scores ↓

Initial evidence. Single + Pair improves all seven point estimates over the no-pretrained-feature control. Pair-only has the lowest interface clash and mean energies. Against the historical AnewOmni reference, Single + Pair improves the structural and mean energy metrics, while AAR is slightly lower.

These are single-seed estimates, not a statistical-significance claim. The AnewOmni row uses historical candidate indices 0–9 and is not noise-paired with the adapters; it is not a paper-exact reproduction. BoltzGen is not included.

Energy distributions and interpretation

dG is the Rosetta interface energy after FastRelax (ref2015); ddG subtracts the corresponding relaxed reference energy. All five groups were scored with the same evaluator. Units are Rosetta energy units (REU), not experimentally measured free energies. Original structures remain the inputs for the five structural metrics.

Model	Median dG (REU)	Median ddG (REU)	dG < 0	ddG < 0
AnewOmni · historical	-18.24	17.47	89.46%	8.60%
Single + Pair	-18.50	16.51	91.18%	9.35%
Single-only	-17.72	17.47	89.25%	8.39%
Pair-only	-18.51	17.09	90.86%	9.89%
No pretrained features	-17.96	17.58	90.00%	9.89%

Means are sensitive to large positive outliers; medians and negative-energy fractions provide additional context. Because every group uses the same reference structures, mean dG and mean ddG have identical rankings and are not independent evidence.

Training results and checkpoint selection

Each adapter was trained from scratch on 10,000 sampled records for 10 epochs. The AnewOmni base and PairFormer donor remained frozen. All four correctness gates passed; fixed validation loss selected the checkpoints.

Adapter	Selected epoch	Initial train loss	Selected train loss	Validation loss
full	9	1.075492	1.061421	1.291728
single	7	1.075496	1.067028	1.294466
pair	9	1.075493	1.061867	1.290897
control	7	1.075498	1.066693	1.294433

Evaluation protocol and implementation scope

Official API target order, native-length peptide inputs, batch size 16, 100 diffusion steps, initial VAE decode followed by re-encoding and one fixed-topology reconstruction cycle, 10 decoding steps per stage, and FP32 highest precision. The existing 930 candidates per adapter were retained without quality-based resampling. Energy scoring adds no new generation.

The released adapter modifies scalar conditioning for all aligned blocks, including context blocks. The paper PDF is compiled from the latest author-supplied source ZIP; the verified results on this page are maintained separately. The data and protocol links above describe this completed 10-candidate evaluation.

BibTeX

@misc{allatomomni2026,
  title={All-Atom Structural Knowledge for Omni Biomolecular Design},
  author={AAA and Han Liu},
  year={2026}
}