Technical Summary: SinAE – A Single-Architecture Flow-Matching Autoencoder for Cross-Domain Atomic Systems

1. Problem Statement

The generation of 3D atomic structures is critical for drug discovery, materials design, and protein engineering. However, current generative pipelines are fragmented across three distinct domains: small molecules, crystals, and proteins. Each domain relies on specialized architectures tailored to its specific structural priors:

• Molecules: Typically use equivariant message-passing networks to handle bond connectivity and valence rules.
• Crystals: Require handling periodic boundary conditions and lattice-basis invariance.
• Proteins: Often rely on residue-level parameterizations (e.g., SE(3) frames) or triangle attention mechanisms.

This fragmentation exacerbates data scarcity, as per-domain datasets (e.g., QM9 with 134k molecules, MP-20 with 45k crystals) are small by modern standards. While cross-domain training could mitigate this, existing approaches face two primary hurdles:

1. Direct Generation Limitations: Generating directly in 3D coordinate space forces a single model to simultaneously learn valid compositions, precise geometry, and heterogeneous structural priors (periodicity, residue backbones, valence), often requiring complex, domain-specific tokenization or staged generation.
2. Latent Autoencoder Limitations: Existing cross-domain latent autoencoders suffer from reconstruction fidelity issues (lossy decoders introduce geometric errors that priors cannot correct) and training complexity (requiring domain-specific loss combinations and weight tuning). No prior method has achieved a simultaneously lossless, architecturally general, and unified-objective autoencoder across all three domains.

2. Methodology: SinAE

The authors introduce SinAE, a single-architecture flow-matching autoencoder designed for molecules, crystals, and proteins.

Architecture

• Unified Backbone: SinAE utilizes a vanilla Transformer for both the encoder and decoder. It explicitly avoids equivariant operators, graph message-passing, or domain-specific modules (e.g., lattice-aware layers or triangle attention).
• Input Featurization:
  • Coordinates: Raw Cartesian coordinates are used without invariant descriptors (e.g., pairwise distances) or canonical frame construction.
  • Periodicity: For crystals, the three lattice vectors are appended as "virtual atom tokens" to the sequence, allowing the model to process the lattice matrix identically to atomic coordinates.
  • Equivariance: Instead of architectural equivariance, the model relies on random SO(3) rotation and translation augmentation during training to ensure orientation robustness.
• Tokenization: Atoms are represented by the sum of learned element embeddings and linear projections of their coordinates. For proteins, residues are treated as single tokens at their $C_\alpha$ positions.

Training Objectives

The model is trained in two stages:

1. Flow-Matching Autoencoder:
   • Encoder: A stochastic Transformer maps the input structure $G=(X, A)$ to a per-token latent code $z$.
   • Decoder: A non-equivariant Transformer reconstructs the structure from $z$ using iterative flow matching. Unlike single-pass VAE decoders, the decoder refines the structure over multiple ODE steps.
   • Loss Function: A unified objective combines continuous conditional flow matching for coordinates ($L_{coord}$) and discrete flow matching for atom types ($L_{atom}$), plus a light KL regularizer ($L_{KL}$). No domain-specific loss weighting or auxiliary losses are required.
2. Latent Flow Matching Prior:
   • A vanilla Diffusion Transformer (DiT) learns the distribution $p(z)$ over the encoder outputs.
   • Since the KL term pushes the posterior toward a standard Gaussian, the prior operates on a smooth, approximately Gaussian latent space without equivariant constraints.
   • During generation, a domain label is used as a condition (via classifier-free guidance) to steer the DiT toward molecules, crystals, or proteins.

3. Key Contributions

1. Single-Architecture Tokenization: SinAE is the first architecture to handle molecules, crystals, and protein backbones using a single vanilla Transformer without equivariant or domain-specific operators.
2. Near-Lossless Reconstruction: By shifting the reconstruction burden to an iterative flow-matching decoder, SinAE achieves sub-milli-Ångström RMSD across domains:
   • Molecules: 0.0002 Å on QM9, 0.0006 Å on GEOM-Drugs.
   • Crystals: 0.0017 Å on MP-20.
   • Proteins: ~0.013 Å backbone CA-RMSD on CASP15.
   • This represents a 10–100× improvement in reconstruction fidelity over prior latent baselines, with 100% structure match rates on molecules and crystals.
3. Positive Cross-Domain Transfer: Joint training on QM9 and MP-20 strictly improves reconstruction and generation performance in both domains compared to single-domain training, providing evidence that a shared atomic latent captures transferable geometric regularities (e.g., bond-length distributions, coordination patterns).

4. Experimental Results

The authors evaluated SinAE on standard benchmarks (QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF).

• Reconstruction Fidelity: SinAE outperforms all baselines (including GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE) in RMSD and structure match rates. The iterative decoder resolves atomic placement early and refines bond geometry/sterics later, achieving high precision without post-hoc relaxation.
• Molecular Generation: On QM9 and GEOM-Drugs, the latent DiT prior achieves high validity (PB Valid scores of 0.97 and 0.93, respectively), surpassing or matching specialized equivariant generators (e.g., EQGAT-diff, SemlaFlow).
• Crystal Generation: On MP-20, the joint SinAE model achieves the highest S.U.N. (Stable, Unique, Novel) and M.S.U.N. scores among compared methods, including those with lattice-invariant architectures. It demonstrates superior thermodynamic stability (metastable structures) despite using a generic Transformer.
• Protein Generation: SinAE achieves high designability (0.96 scRMSD ≤ 2.0 Å) and diversity simultaneously, outperforming ProteinAE and LatentDiff. In sequence-structure co-design tasks, it achieves competitive scTM and scRMSD scores at lengths 100 and 200.

5. Significance and Claims

The paper claims that SinAE demonstrates that cross-domain atomic generation can be unified under a single, simple architecture without sacrificing fidelity.

• Simplicity vs. Performance: The authors argue that complex domain-specific operators are not strictly necessary for high-quality generation if the decoder is sufficiently powerful (iterative flow matching) and the latent space is properly regularized.
• Data Efficiency: The results provide direct evidence that shared atomic latents allow different chemical populations to benefit from each other's data, improving performance in low-data regimes.
• Decoupling Geometry and Priors: The work suggests that geometric complexity can be handled by the decoder's iterative refinement, allowing the encoder and prior to operate on a simplified, domain-agnostic latent space.

Limitations Acknowledged:

• Protein evaluation is currently limited to backbone-level generation; all-atom modeling is not yet addressed.
• The current model does not support property-, composition-, or length-conditioned generation in a single checkpoint across all three domains simultaneously.
• Longer protein chains (>256 residues) require retraining with higher length budgets.