🤖 AI Summary
Existing generative models struggle to jointly model small molecules, crystals, and proteins due to the heterogeneous structural priors inherent in these three domains within 3D space. This work proposes SinAE—the first cross-domain atomic system generation framework based on a standard Transformer architecture—which leverages a flow-matching autoencoder to delegate geometric reconstruction to an iterative decoder, eliminating the need for graph networks, equivariant operators, or domain-specific modules. SinAE achieves near-lossless reconstruction, reducing latent space errors by several orders of magnitude compared to prior baselines, and attains state-of-the-art performance across all three generation tasks. Furthermore, joint training demonstrates that a shared latent representation effectively facilitates cross-domain transfer and enhances overall performance.
📝 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 .