Thrilled to announce Zatom-2, an atomistic generative model spanning molecular, material, and protein science! Special thanks to @MirunaCretu2 for leading this effort. (1/11) 🌐 Project page: zatom-ai.github.io/nucleus

Interested in keeping up with advancements in AI for Science? Me too. That's why I started a monthly blog covering new papers in the field. Feel free to check it out! Substack: amorehead.substack.com/

🧬 If you’d like to contribute to the next generation of open-source atomistic generative models in chemistry and biology, please reach out! (11/11)

After finetuning, Zatom-2 can generate all-atom proteins by way of SCOPe. (10/11)

Zatom-2 also supports generative sampling of periodic materials via OMat24. (9/11)

Zatom-2 enables generative sampling of diverse molecular conformations thanks to OMol25. (8/11)

Analyzing representations, scaling behavior, and generation quality, our results show that joint pretraining across molecular+materials data improves generative modeling and representation transfer, pointing toward general-purpose generative models for atomistic systems. (7/11)

Furthermore, we demonstrate that conditioning on atomic forces enables control over the degree of geometric relaxation in generated structures. (6/11)

Zatom-2 establishes strong performance on QM9, GEOM-Drugs, MP20, OMat24, and OMol25. When transferring to as few as 2,000 protein training examples, Zatom-2 generates designable backbones and extrapolates strongly to sequence lengths beyond the training distribution. (5/11)

This multitask pretraining leverages the rich geometric information in quantum-chemical calculations to shape the model's internal representations. (4/11)

Uniquely, Zatom-2 is pretrained on the large-scale OMol25 and OMat24 datasets using a multitask objective that combines flow matching generation and structure prediction with auxiliary machine learning interatomic potential (MLIP) tasks predicting energies and forces. (3/11)

Atomistic machine learning models typically target specific domains, such as small molecules, crystals, or proteins. Here, we introduce a unified atomistic generative model across these domains thanks to Zatom-2’s new atom1+atom14 tokenization scheme for chem+bio data. (2/11)

📢 Only two days left to submit to LoG 2026! We welcome submissions on all areas broadly related to learning on graphs and geometry. 🔗 logconference.org/cfp/ 📅 Key dates (AoE): • July 29 – Abstract deadline • August 1 – Full paper deadline

🧬 MegaFold: Efficient AlphaFold-3 Training Across NVIDIA and AMD GPUs MegaFold will be presented at ISC High Performance 2026! 🗺️Research Paper Session: ML Systems and Frameworks 🗓️ Wednesday, June 24 · 1:00–2:00 PM (CEST) 📍Hall E, 2nd Floor, Hamburg, Germany Come find out how we train AlphaFold-3 on 3.36× longer sequences, up to 1.73× faster, on both NVIDIA and AMD GPUs. Using MegaFold, we can successfully train on H200s, MI250 and MI325X. 🚀 Think training AlphaFold-3 (AF3) is just like training an LLM? Think again. Despite having only ~500M parameters, AF3 is up to two orders of magnitude slower to train than a size-equivalent LLM and OOMs on sequences as short as 590 residues. Why? One operator: 3D attention over 2D pairwise representations. It makes compute and memory scale cubically with sequence length, blowing up activation memory, fragmenting execution into tens of thousans of tiny kernels per iteration, and starving the GPU while the CPU grinds through evolutionary-context preprocessing. In our latest work, MegaFold, we diagnose exactly why AF3 training is slow and build an end-to-end system that fixes it across the entire stack. The results: ✅ 3.36× longer trainable sequences on 32 GPUs ✅ Up to 1.73× faster per-iteration training (1.62× on AMD) ✅ 23% lower peak memory ✅ Runs on both NVIDIA H200, AMD MI250 & AMD MI325x, fully cross-platform 🔧 Four key innovations: - EvoFlash-3D: a cross-platform Triton kernel that computes 3D attention without ever materializing the cubic attention map (cutting EvoAttention memory by 4.24×). - EvoSP-3D: a communication-efficient sharding strategy built for alternating-axis 2D pairwise representations - EvoFusion: a fused operator stack that tames AF3's kernel fragmentation - EvoPipe: a determinism-aware host-device pipeline that eliminates preprocessing stalls And it's verified: MegaFold matches baseline training loss step-for-step, so you get all the speed and memory savings with zero impact on convergence. Huge thanks to our students Hoa La and Ahan Gupta for working hard on this project, and to our collaborators Jianlin Cheng and Alex Morehead for their valuable contributions, feedback, and support throughout the work! MegaFold is fully open source. 🎉 💻 Code: github.com/Supercomputing-Sy… 🔗 Project page: supercomputing-system-ai-lab… 📄 Paper: arxiv.org/pdf/2506.20686 📝 Blog: supercomputing-system-ai-lab… See you at ISC, Hall E! 👋

A review of flow matching and its role in CompBio/virtual cell research, now in Nature Machine Intelligence! Congrats, team! (@lazar_atan, Akshata Hegde, Yanli Wang, @frimpongboadu1, @joelselvaraj95, @AlexanderTong7, @ask1729, @jianlincheng) Paper: tinyurl.com/3jj4k6c3

Introducing Zatom-1, the first end-to-end, fully open-source foundation model for 3D chemistry! This was a great collaborative effort with many brilliant scientists. I'm grateful to have played a small part. Paper: arxiv.org/abs/2602.22251 Code: github.com/Zatom-AI/zatom

Excited to announce that PoseBench is now published in Nature Machine Intelligence! Paper: rdcu.be/eW5oj Code: github.com/BioinfoMachineLea…

Interested in learning more about how flow matching has begun to advance bioinformatics and computational biology? And how it has already started making strides towards the development of an AI-based virtual cell? Paper: authorea.com/users/637193/ar… Code: github.com/amorehead/awesome…

🔬Interested in training AlphaFold3 faster, at scale, and beyond NVIDIA GPU? Now you can. AlphaFold3 is a major leap in biomolecular modeling, but behind the scenes, it introduces severe system bottlenecks: 🧠 2D EvoAttention spikes memory usage 📉 Retrieval-augmented training pipeline causes long GPU idle time ⛔ Frequent but memory-intensive ops slow everything down Today, I'm excited to announce MegaFold, a fully open-source system to make AlphaFold3 training fast, scalable, and cross-platform on both NVIDIA and AMD GPUs. MegaFold delivers: ⚡ Up to 1.73x / 1.62x faster training on NVIDIA H100 / AMD MI250 🧬 Up to 1.35× longer sequences compared to PyTorch baseline Key features: 🚀 Memory-Efficient EvoAttention via portable Triton kernels 💡 Ahead-of-Time Caching to eliminate GPU idle time in retrieval pipelines 🔗 DeepFusion for reducing overhead of small but frequent memory-intensive AF3 ops 📘 Project page: supercomputing-system-ai-lab… 📄 Paper: arxiv.org/pdf/2506.20686 💻 Code: github.com/Supercomputing-Sy… 🤝 MegaFold is developed in collaboration between UIUC SSAIL Lab and researchers from University of Missouri and Lawrence Berkeley National Laboratory. Kudos to the brilliant team: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng #AlphaFold3 #AI #ProteinFolding #Bioinformatics #AMD #Triton #CrossPlatform #OpenSource