One Path to Model Them All: Learnable-Time Flow Matching for Crystal Structure and Energy Prediction

Abstract

Crystals are cornerstone materials for semiconductors and renewable energy, yet their discovery is hindered by the prohibitive cost of Density Functional Theory (DFT). While geometric graph neural networks have advanced Crystal Structure Prediction (CSP) and energy estimation, existing methods treat these tasks as disparate problems with task-specific architectures, failing to exploit cross-task synergies. In this paper, we propose UniPath, a novel framework based on Learnable-Time Flow Matching that unifies crystal structure and energy prediction within a single probability path. Our key insight is modeling unstable crystals as an intermediate state along the probability path from a prior to the stable distribution. We introduce two pivotal mechanisms: (1) a learnable time parameter that dynamically identifies the intermediate state, and (2) an auxiliary velocity field that parameterizes tailored flows to synchronize the intermediate and unstable distributions. We theoretically prove that under optimal velocity fields, these paths converge to a consistent trajectory, enabling a unified backbone to maximize shared representation learning capabilities. Extensive evaluations on CSP and energy prediction benchmarks demonstrate that UniPath significantly outperforms task-specific baselines. Notably, UniPath achieves over 50% higher Match Rate on MPTS-52 and over 30% lower energy MAE on MPT-MPTS-52. Furthermore, UniPath generalizes effectively to small molecules, surpassing conventional models in molecular conformation generation. The source code is publicly available at https://github.com/GLAD-RUC/UniPath.