Technische Samenvatting: SinAE – Een Single-Architecture Flow-Matching Autoencoder voor Cross-Domain Atomaire Systemen

1. Probleemstelling

De generatie van 3D-atomaire structuren is cruciaal voor medicijnontdekking, materiaalkunde en proteïne-engineering. De huidige generatieve pipelines zijn echter gefragmenteerd over drie afzonderlijke domeinen: kleine moleculen, kristallen en proteïnen. Elk domein vertrouwt op gespecialiseerde architecturen die zijn afgestemd op hun specifieke structurele priors:

• Moleculen: Gebruiken doorgaans equivariante message-passing netwerken om bindingsconnectiviteit en valentieregels te verwerken.
• Kristallen: Vereisen het afhandelen van periodieke randvoorwaarden en rooster-basis invariantie.
• Proteïnen: Vertrouwen vaak op residue-niveau parametrisaties (bijv. SE(3)-frames) of driehoeks-attention mechanismen.

Deze fragmentatie verergert de schaarste aan data, aangezien per-domein datasets (bijv. QM9 met 134k moleculen, MP-20 met 45k kristallen) naar moderne maatstaven klein zijn. Hoewel cross-domain training dit zou kunnen verzachten, worden bestaande benaderingen geconfronteerd met twee primaire hindernissen:

1. Beperkingen van Directe Generatie: Het direct genereren in 3D-coördinatieruimte dwingt een enkel model om simultaan geldige composities, precieze geometrie en heterogene structurele priors (periodiciteit, residue backbones, valentie) te leren, wat vaak complexe, domeinspecifieke tokenisatie of gefaseerde generatie vereist.
2. Beperkingen van Latente Autoencoders: Bestaande cross-domein latente autoencoders lijden aan reconstructiefidelity problemen (lossy decoders introduceren geometrische fouten die priors niet kunnen corrigeren) en trainingscomplexiteit (het vereist domeinspecifieke loss-combinaties en gewichtstuning). Geen enkele voorafgaande methode heeft een simultaan verliesloze, architecturaal algemene en met één enkel doel verbonden autoencoder bereikt over alle drie de domeinen.

2. Methodologie: SinAE

De auteurs introduceren SinAE, een single-architecture flow-matching autoencoder ontworpen voor moleculen, kristallen en proteïnen.

Architectuur

• Unified Backbone: SinAE maakt gebruik van een vanilla Transformer voor zowel de encoder als de decoder. Het vermijdt expliciet equivariante operatoren, graph message-passing of domeinspecifieke modules (bijv. rooster-bewuste lagen of driehoeks-attention).
• Input Featurization:
  • Coördinaten: Ruwe Cartesiaanse coördinaten worden gebruikt zonder invariante beschrijvers (bijv. paarafstand-afstanden) of canonieke frame-constructie.
  • Periodiciteit: Voor kristallen worden de drie roostervectoren toegevoegd als "virtuele atoom-tokens" aan de sequentie, waardoor het model de rooster-matrix identiek aan atomaire coördinaten kan verwerken.
  • Equivariantie: In plaats van architecturale equivariantie, vertrouwt het model op random SO(3) rotatie- en translatie-augmentatie tijdens de training om oriëntatie-robuustheid te garanderen.
• Tokenization: Atomen worden gerepresenteerd door de som van geleerde element-embeddings en lineaire projecties van hun coördinaten. Voor proteïnen worden residuen behandeld als enkele tokens op hun $C_\alpha$ posities.

Trainingsdoelen

Het model wordt getraind in twee stadia:

1. Flow-Matching Autoencoder:
   • Encoder: Een stochastische Transformer brengt de inputstructuur $G=(X, A)$ in kaart naar een per-token latente code $z$.
   • Decoder: Een niet-equivariante Transformer reconstrueert de structuur vanuit $z$ met behulp van iteratieve flow matching. In tegen tegenstelling tot single-pass VAE decoders, verfijnt de decoder de structuur over meerdere ODE-stappen.
   • Loss Functie: Een verenigd doel combineert continue conditional flow matching voor coördinaten ($L_{coord}$) en discrete flow matching voor atoomtypen ($L_{atom}$), plus een lichte KL-regularisator ($L_{KL}$). Geen domeinspecifieke loss-weging of hulp-losses zijn vereist.
2. Latent Flow Matching Prior:
   • Een vanilla Diffusion Transformer (DiT) leert de distributie $p(z)$ over de encoder-outputs.
   • Omdat de KL-term de posterior richting een standaard Gaussische verdeling duwt, opereert de prior op een gladde, ongeveer Gaussische latente ruimte zonder equivariante beperkingen.
   • Tijdens de generatie wordt een domeinlabel gebruikt als conditie (via classifier-free guidance) om de DiT te sturen naar moleculen, kristallen of proteïnen.

3. Belangrijkste Bijdragen

1. Single-Architecture Tokenization: SinAE is de eerste architectuur die moleculen, kristallen en proteïne-backbones behandelt met een enkele vanilla Transformer zonder equivariante of domeinspecifieke operatoren.
2. Bijna Verliesloze Reconstructie: Door de reconstructie-last te verschuiven naar een iteratieve flow-matching decoder, bereikt SinAE een sub-milli-Ångström RMSD over alle domeinen:
   • Moleculen: 0.0002 Å op QM9, 0.0006 Å op GEOM-Drugs.
   • Kristallen: 0.0017 Å op MP-20.
   • Proteïnen: ~0.013 Å backbone CA-RMSD op CASP15.
     Dit vertegenwoordigt een 10–100× verbetering in reconstructiefidelity ten opzichte van eerdere latente baselines, met 100% structurele match-rates op moleculen en kristallen.
3. Positieve Cross-Domain Transfer: Gezamenlijke training op QM9 en MP-20 verbetert de reconstructie en generatieprestaties in beide domeinen strikt ten opzichte van single-domein training, wat bewijs levert dat een gedeelde atomaire latent overdraagbare geometrische regulariteiten bevat (bijv. bindingslengte-distributies, coördinatiepatronen).

4. Experimentele Resultaten

De auteurs hebben SinAE geëvalueerd op standaard benchmarks (QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF).

• Reconstructiefidelity: SinAE presteert beter dan alle baselines (inclusclusief GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE) in RMSD en structurele match-rates. De iteratieve decoder lost de plaatsing van atomen vroeg op en verfijnt de bindingsgeometrie/steriek later, waardoor hoge precisie wordt bereikt zonder post-hoc relaxatie.
• Moleculaire Generatie: Op QM9 en GEOM-Drugs bereikt de latente DiT prior een hoge validiteit (PB Valid scores van respectievelijk 0.97 en 0.93), waarmee het gespecialiseerde equivariante generatoren (zoals EQGAT-diff, SemlaFlow) evenaart of overtreft.
• Kristalgeneratie: Op MP-20 bereikt het gezamenlijke SinAE-model de hoogste S.U.N. (Stable, Unique, Novel) en M.S.U.N. scores onder de vergeleken methoden, inclusief die met rooster-invariante architecturen. Het vertoont superieure thermodynamische stabiliteit (metastabiele structuren) ondanks het gebruik van een generieke Transformer.
• Proteïnegeneratie: SinAE bereikt tegelijkertijd een hoge designbaarheid (0.96 scRMSD ≤ 2.0 Å) en diversiteit, waarbij het ProteinAE en LatentDiff overtreft. In taken voor sequentie-structuur co-design behaalt het concurrerende scTM en scRMSD scores bij lengtes van 100 en 200.

5. Betekenis en Claims

Het artikel claimt dat SinAE aantoont dat cross-domein atomaire generatie verenigd kan worden onder een enkele, eenvoudige architectuur zonder in te boeten op fidelity.

• Eenvoud vs. Prestaties: De auteurs stellen dat complexe domeinspecifieke operatoren niet strikt noodzakelijk zijn voor hoogwaardige generatie als de decoder voldoende krachtig is (iteratieve flow matching) en de latente ruimte goed geregulariseerd is.
• Data-efficiëntie: De resultaten leveren direct bewijs dat gedeelde atomaire latents ervoor zorgen dat verschillende chemische populaties kunnen profiteren van elkaars data, wat de prestaties verbetert in scenario's met weinig data.
• Ontkoppeling van Geometrie en Priors: Het werk suggereert dat geometrische complexiteit kan worden afgehandeld door de iteratieve verfijning van de decoder, waardoor de encoder en de prior kunnen opereren op een vereenvoudigde, domein-agnostische latente ruimte.

Erkende Beperkingen:

• De evaluatie van proteïnen is momenteel beperkt tot backbone-niveau generatie; all-atom modellering wordt nog niet behandeld.
• Het huidige model ondersteunt geen property-, compositie- of lengte-geconditioneerde generatie in één enkel checkpoint over alle drie de domeinen tegelijkertijd.
• Langere proteïneketens (>256 residuen) vereisen hertraining met hogere lengte-budgetten.