cs.LGJul 14, 2026

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

Authors: Yuxuan Ren, Fan Yang, Jianhua Yao, Yatao Bian

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 .

Explore similar work

CardsList
Sep 8, 2026physics.chem-ph

Fixed-Dimensional Latent Flow for Generating Variable-Size 3D Molecules

Molecular size is coupled to composition, structure, and function, yet most 3D molecular generators require a predefined atom count. We introduce Equivariant-Free Transformer-Autoencoded Latent Flow Matching, a two-stage framework that samples a fixed-dimensional latent vector using flow matching and uses an autoregressive Transformer to determine molecular size, atom types, coordinates, and chemical attributes. Canonical atom ordering and rigid-pose alignment enable Transformers without equivariant layers, while decoded attributes guide bond reconstruction. On PCQM4Mv2, unconditional generation yields 87.9% unique, novel molecules passing sanitization and PoseBusters checks, exceeding baselines with lower end-to-end training and sampling time and higher end-to-end throughput. Across ten target HOMO-LUMO gaps, internal ranking retains 30% of screened candidates and increases the density functional theory-verified hit rate within 0.1 eV from 25.0% to 52.4%, while largely preserving novelty and diversity. These results demonstrate fixed-dimensional latent generation with autoregressive decoding as a practical approach to molecular design without prespecifying size.
Weichi Yao, Cameron Gruich, Bryan R. Goldsmith +1
Molecular GenerationMolecular Structure ElucidationGenerative Flow Networks3D Object UnderstandingLatent Representation
AskPDF
Aug 12, 2026cs.LG

NAE: Normalizing AutoEncoder

We consider the setting of Normalizing flows with approximate inverses, an established paradigm spanning both full-dimensional (d=Dd=D) and bottleneck (d<Dd<D) settings, and group these models under the term flow autoencoders. We present a theoretical investigation into their training dynamics and prove that the proposed loss used by existing approaches is suboptimal; specifically, both encoder and decoder surrogates must be optimized in alignment with reconstruction loss. Guided by these insights, we propose Normalizing Autoencoder (NAE), which employs a novel conditional loss that aligns the surrogate loss gradient with that of reconstruction loss, directly improving upon the current standard. Extensive experiments across molecule generation, tabular data, and image benchmarks demonstrate that NAE achieves state of the art performance. Our work highlights the importance of loss alignment in flow autoencoders and establishes NAE as a powerful generative framework.
Muhammad Abdur Rafae, Niels Landwehr
Masked AutoencodersTime-Step-Wise Learnable Surrogate Gradient FunctionMolecular GenerationNormalization
AskPDF
Jul 10, 2026cs.LG

Autoregressive latent diffusion for 3D molecule generation

Three-dimensional (3D) molecule generation has been dominated by diffusion models, which achieve strong generation quality but typically require the molecular size to be specified a priori. Recent autoregressive approaches have substantially narrowed the performance gap while naturally supporting variable-length generation and conditioning on partial molecular context. However, balancing unconditional and context-conditioned generation remains challenging. We introduce KRONOS, a latent autoregressive diffusion framework that generates molecules in the latent space of a pre-trained autoencoder, jointly modeling molecular graph topology and geometry, while retaining the flexibility of autoregressive generation. We further introduce a mixed training strategy inspired by Fill-in-the Middle (FIM) paradigm, enabling both unconditional and fragment-conditioned molecular generation within a single left-to-right autoregressive model. Experiments on QM9 and GEOM-Drugs demonstrate that KRONOS achieves leading unconditional generation performance among autoregressive methods, while remaining competitive with diffusion models. Moreover, fragment-conditioned generation is achieved with negligible impact on unconditional generation performance, demonstrating that both generation paradigms can be supported within a single architecture.
Federico Ottomano, Gaopeng Ren, Yingzhen Li +2
Molecular GenerationAutoregressive Diffusion3D Latent Diffusion Model3D Object UnderstandingChronos
AskPDF