# APO: Unsupervised Atomic Policy Optimization for 3D Structure Prediction of Atomic Systems

arXiv: https://arxiv.org/abs/2607.28553 · HTML: https://mlsanity.com/paper/2607.28553/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 19 references that are in the map.
- First posted on arXiv 30 Jul 2026.

## Builds on (most cited first, 8 of 8)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Crystal Diffusion Variational Autoencoder for Periodic Material Generation](https://mlsanity.com/paper/2110.06197/) (2021)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Protein structure generation via folding diffusion](https://mlsanity.com/paper/2209.15611/) (2022)

## Abstract (arXiv)

Predicting the 3D structures of atomic systems is fundamental to advancing material science and drug discovery. While flow-matching models (, FlowDPO) have recently shown promise in this domain, their performance relies heavily on alignment with ground-truth coordinates via supervised preference learning. However, obtaining experimental labels for novel crystal phases or de novo proteins is prohibitively expensive, creating a bottleneck for structural modeling in data-scarce regimes. In this work, we propose (Atomic Policy Optimization), a fully unsupervised alignment framework that eliminates the need for ground-truth reference structures. APO adapts group-relative policy optimization to 3D atomic environments, utilizing a novel dual-reward mechanism: (i) a that reinforces the policy's dominant latent structural modes through eigen-decomposition of sample similarities, and (ii) a that enforces thermodynamic stability. Our framework enables the model to ``self-correct'' by identifying physically plausible configurations within sampled groups. Extensive benchmarks on crystal and antibody structure prediction demonstrate that APO consistently outperforms fully supervised baselines, achieving a new state-of-the-art in match rates and structural fidelity. Furthermore, we show that APO effectively straightens probability paths, significantly improving inference efficiency. Our results suggest that intrinsic physical consistency can serve as a superior guide for alignment compared to noisy, supervised coordinate matching.
