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

Abstract: Small molecules, crystals, and proteins all reduce to atoms in 3D space, yet their generative pipelines remain fragmented across domains, each with its Small molecules, crystals, and proteins all reduce to atoms in 3D space, yet their generative pipelines remain fragmented across domains, each with its own graph, equivariant, or frame-based architecture. Cross-domain training would mitigate per-domain data scarcity, but direct generation in 3D coordinate space cannot easily handle the heterogeneous structural priors of all three domains, and no prior latent autoencoder is simultaneously lossless and architecturally general across all three. We introduce SinAE, a single-architecture flow-matching autoencoder for molecules, crystals, and proteins, with vanilla Transformer encoder and decoder and no equivariant, graph, or domain-specific operators. Rather than requiring the encoder to capture fine-grained geometry, SinAE shifts the reconstruction burden into an iterative flow-matching decoder, achieving near-lossless reconstruction across domains and reducing reconstruction errors by orders of magnitude relative to prior latent baselines. The same per-token latent supports a standard Diffusion Transformer prior that reaches strong performance on molecular, crystal, and protein generation benchmarks. Joint molecule--crystal training strictly improves both domains, providing direct evidence of cross-domain transfer through a shared atomic latent. Code is available at https://github.com/BlueWhaleLab/SinAE .

• The paper introduces SinAE, a unified Transformer-based autoencoder that uses iterative flow matching for cross-domain reconstruction of molecules, crystals, and proteins.
• It achieves near-lossless reconstruction with sub-milliångström RMSD and superior generative validity, recording a PB Valid of 0.97 on molecular benchmarks.
• The efficient ODE-based decoder and joint latent space enable rapid inference and effective transfer of geometric regularities across diverse atomic systems.

SinAE: A Unified Flow-Matching Autoencoder for Cross-Domain Atomic Systems

Motivation and Context

Existing generative pipelines for molecules, crystals, and proteins have developed along distinct lines, largely due to domain-specific structural conventions and physical priors. Despite sharing an atomic-level representation, previous approaches have relied on specialized architectures—graph and equivariant networks for molecules, lattice-aware operators for crystals, and residue-level frames for proteins. This fragmentation exacerbates per-domain data scarcity and complicates the development of universal generators. The introduction of SinAE ("A Single-Architecture Flow-Matching Autoencoder for Cross-Domain Atomic Systems" (2607.12380)) addresses this by proposing a domain-general, vanilla-Transformer autoencoder combined with flow-matching for both reconstruction and generation, requiring no specialized geometric or equivariant modules.

Figure 1: SinAE overview: A flow-matching autoencoder architecture for molecules, crystals, and proteins, leveraging a shared per-token latent representation.

Methodology

SinAE comprises two core components: a vanilla Transformer encoder mapping input atomic structures ($\mathbf{X}, \mathbf{A}$) to a compact, per-token latent ($\mathbf{z}$), and a Transformer decoder implementing iterative flow matching. The encoder is stochastic, aligning its output distribution with a Gaussian via light KL regularization. The decoder reconstructs atomic coordinates and types by minimizing a unified flow-matching objective, dispensing with domain-specific loss terms and architectural modules. For crystals, periodicity is encoded via virtual lattice tokens processed as "atoms," ensuring architectural uniformity across all domains.

During training, random SO(3) rotation augmentations maintain orientation-invariant representations without equivariant layers. The decoder refines its output using an ODE trajectory, progressively improving atomic placement and resolving fine geometric details.

A latent DiT prior (Diffusion Transformer) is trained over the encoder outputs, enabling de novo generation by sampling a latent sequence and decoding via the flow-matching decoder.

Figure 2: SinAE architecture and training procedures: (a–d) show shared embedding, encoding, decoding, and prior learning; (e–g) outline reconstruction, prior training, and generation.

Empirical Evaluation

Reconstruction Fidelity

SinAE achieves sub-milli-Å RMSD across all domains, with structure match rates of $100\%$ on QM9 (small molecules) and MP-20 (crystals), demonstrating near-lossless recovery compared to baselines, which exhibit substantially higher RMSD and lower match rates. For proteins, backbone CA-RMSDs of $\sim 0.01$–$0.02$ Å represent an order-of-magnitude improvement over prior latent autoencoders (e.g., ProteinAE, DPLM-2).

Figure 3: Reconstruction RMSD across domains summarized on log scale; SinAE achieves $20$–$100\times$ tighter RMSD than prior baselines in crystals and proteins.

Strong cross-domain transfer is evidenced by joint training: pooling QM9 and MP-20 reduces RMSD further, from $0.0007$ to $0.0002$ Å (QM9) and $0.0041$ to $\mathbf{z}$0 Å (MP-20), strictly improving both domains. The decoder converges to optimal fidelity after only six function evaluations, confirming computational efficiency.

Generative Performance: Molecules

SinAE attains competitive or superior PoseBusters validity (PB Valid), which enforces stringent checks on geometric and chemical plausibility. On QM9, the joint model yields PB Valid $\mathbf{z}$1, surpassing state-of-the-art equivariant models (EQGAT-diff $\mathbf{z}$2, FlowMol $\mathbf{z}$3). On GEOM-Drugs, it achieves $\mathbf{z}$4 PB Valid, markedly higher than ADiT ($\mathbf{z}$5) and SemlaFlow ($\mathbf{z}$6).

Figure 4: Generated samples from SinAE's latent DiT prior decoded across domains; small molecules, crystals, and protein backbones illustrated.

Generative Performance: Materials

On MP-20, SinAE sets new benchmarks in S.U.N. (Stable, Unique, Novel) and M.S.U.N. (Metastable S.U.N.) metrics, reaching $\mathbf{z}$7 and $\mathbf{z}$8 respectively, outperforming FlowLLM ($\mathbf{z}$9) and FlowMM ($100\%$0). Structure and composition validity approach $100\%$1, though SinAE's main improvement lies in thermodynamic stability, as evaluated by DFT-relaxed energies.

Figure 5: Unconditional crystal samples from SinAE's joint molecule-crystal prior, decoded with lattice vectors to produce inorganic crystal structures.

Generative Performance: Proteins

Using the same architecture as the molecule-crystal model, the protein variant achieves backbone reconstruction RMSDs of $100\%$2–$100\%$3 Å on CASP14/15. Unconditional generation yields designability $100\%$4 and diversity $100\%$5 clusters at temperature $100\%$6, outperforming ProteinAE and matching RFdiffusion.

Co-design tasks (sequence+structure generation) show competitive self-consistency scores: scRMSD $100\%$7 Å, scTM $100\%$8 (length 100), matching MultiFlow and ProteinGenerator, and below APM (a protein-specific model) in scTM.

Figure 6: Unconditional protein backbone samples generated from the protein-only SinAE prior, visualized as C$100\%$9 cartoon traces.

Latent Space Analysis and Cross-Domain Transfer

t-SNE projections of joint latent spaces reveal domain separation and clustering without explicit conditioning, with shared low-density boundary regions suggestive of transferable geometric regularities. Interpolations between latents decode into valid intermediate structures across molecules, crystals, and proteins, evidencing a smooth and continuous latent manifold.

Figure 7: t-SNE of SinAE latents shows distinct yet internally coherent regions for molecules and crystals, colored by domain-specific clusters.

Figure 8: Latent-space interpolations between pairs of SinAE latents decode into valid, continuous intermediate structures across all domains.

Decoder Trajectory Insights

Analysis of decoder ODE trajectories (generation from DiT-sampled latents) shows rapid RMSD convergence ($\sim 0.01$0) followed by continued energy minimization, indicating early geometric placement and late-stage local refinement—bond planarity, dihedral angles, sterics. This applies even to prior-sampled latents unseen by the encoder, confirming the physical relevance and robustness of the flow-matching decoder.

Figure 9: Decoder ODE trajectory during generation: (a) physical energy declines logarithmically; (b) RMSD quickly plateaus, with late steps refining energetics.

Generalization and Efficiency

Protein length generalization (tested up to 260 residues) exhibits sub-linear scaling of CA-RMSD with chain length. Generation NFE ablation demonstrates rapid convergence of reconstruction RMSD, saturating by $\sim 0.01$1. In wall-clock terms, SinAE's decoder achieves fast inference, requiring only $\sim 0.01$2 ms per molecule, confirming practical deployability.

Figure 10: Reconstruction RMSD versus decoder NFE; curves plateau rapidly, evidencing efficient convergence.

Figure 11: SinAE protein model generalizes reconstruction fidelity across full length range; generative scRMSD is stable across sampled lengths.

Implications and Future Directions

SinAE's architecture shifts the reconstruction burden entirely onto an iterative, flow-matching decoder, enabling domain-general latent spaces without equivariance or domain-specific modules. This level of generality and fidelity facilitates cross-domain transfer, robust generation, and scale-up across small molecules, materials, and proteins. It opens avenues for length-conditioned generative modeling, all-atom protein design, and single-checkpoint cross-domain training.

Potential advances include integrating all-atom protein modeling, conditional generation (property, composition, or length), and hybrid training across more extensive domain pools, with attendant benefits for generative discovery pipelines in biochemistry and materials science. Limitations are noted in protein evaluation granularity and checkpoint scalability.

Conclusion

SinAE demonstrates that flow-matching autoencoders with vanilla Transformer backbones can achieve near-lossless reconstruction and competitive generative performance across diverse atomic domains, with joint training enabling strict improvements via cross-domain transfer. The iterative decoder, unified objective, and latent-space diffusion prior offer a scalable, architecture-neutral foundation for cross-domain atomic structure modeling. SinAE advances the possibilities for atomic-scale generative AI, bridging molecular, material, and protein design.