Kinetic Langevin Diffusion for Crystalline Materials Generation

Conference: ICML 2025
arXiv: 2507.03602
Code: None
Area: Computational Biology
Keywords: Crystal material generation, Langevin diffusion, hypertorus, equivariance, structure prediction

TL;DR

KLDM proposes using Kinetic Langevin Diffusion to address the issue of fractional atomic coordinates residing on a hypertorus in crystal material generation. By introducing an auxiliary velocity variable, the diffusion process is shifted to a flat Euclidean space while preserving periodic translational symmetry, achieving competitive performance on crystal structure prediction and de novo generation tasks.

Background & Motivation

Background: Use of diffusion models for crystalline material generation has become an important direction in materials science. Crystal structures are defined by lattice parameters, atom types, and the fractional coordinates of atoms within the unit cell.

Limitations of Prior Work: The distribution of crystal data possesses inherent symmetries and involves multiple modalities (discrete atom types, continuous coordinates, and lattice parameters). A key challenge lies in the fractional coordinates residing on a hypertorus \(\mathbb{T}^d\) (with periodic boundary conditions), requiring special treatment.

Key Challenge: Although performing Riemannian diffusion directly on the hypertorus is mathematically correct, the training objective is complex and struggles to handle the periodic translational symmetry of crystals. Conversely, simply ignoring periodicity and performing diffusion in Euclidean space leads to the generation of invalid structures.

Goal: Design a diffusion model that respects the hypertorus geometry while being efficiently trainable in a flat space.

Key Insight: Drawing inspiration from Kinetic Langevin Dynamics in physics (which introduces velocity as an auxiliary variable), the diffusion process on the torus is "lifted" to the tangent space (Euclidean space).

Core Idea: Use the Kinetic Langevin method to couple the diffusion process on the hypertorus to a flat velocity space, allowing the training objective to naturally account for periodic symmetry.

Method

Overall Architecture

• Input: Material composition information (atom types) and optional target properties.
• Diffusion Space: Joint space of coordinates on the hypertorus + auxiliary velocity in the Euclidean space.
• Reverse Process: Learning the joint reverse diffusion denoising to generate coordinates and lattice parameters.
• Output: Complete crystal structure (lattice + atomic positions + atom types).

Key Designs

1. Generalization of Trivialized Diffusion Model (TDM):

   • TDM originally performs diffusion on manifolds using tangent spaces. KLDM generalizes it to handle the unique periodic translational symmetry of crystals.
   • Fractional coordinates \(\mathbf{x} \in \mathbb{T}^{3N}\) are coupled with velocity variables \(\mathbf{v} \in \mathbb{R}^{3N}\) in the tangent space via exponential mapping.
   • Design Motivation: The hypertorus lacks a global coordinate system; TDM provides an elegant way to perform diffusion in the tangent space (flat space).

2. Handling of Periodic Translational Symmetry:

   • The true data distribution of crystals satisfies periodic translational invariance: translating all atoms simultaneously by a lattice vector does not change the physical structure.
   • The training objective of KLDM naturally accounts for this symmetry.
   • Through the coupling of velocity variables, gradient information correctly reflects periodic boundary conditions.
   • Design Motivation: Ignoring this symmetry would cause the model to treat physically equivalent structures as distinct, reducing data efficiency.

3. Multi-modal Joint Generation:

   • Simultaneously handles continuous variables (coordinates, lattice parameters) and discrete variables (atom types).
   • Coordinates are modeled with Kinetic Langevin diffusion, lattice parameters with standard Euclidean diffusion, and atom types with discrete diffusion.
   • Design Motivation: Crystal generation is inherently a multi-modal generation problem.

Loss & Training

• Denoising score matching loss for the Kinetic Langevin version.
• Joint denoising objective for positions and velocities.
• Standard diffusion loss for lattice parameters.
• Cross-entropy loss for atom types.

Key Experimental Results

Main Results

Task	Dataset	KLDM	Prev. SOTA	Description
Crystal Structure Prediction (CSP)	Perov-5	Competitive	DiffCSP/FlowMM	Match Rate metric
CSP	MP-20	Competitive	DiffCSP++	Match Rate
De Novo Generation (DNG)	Perov-5	Competitive	CDVAE/DiffCSP	Validity + Diversity
DNG	MP-20	Competitive	Existing methods	Stability

Ablation Study

Configuration	Key Metric	Description
Standard Euclidean diffusion	Decreased	Ignoring periodicity leads to invalid structures
Direct Riemannian diffusion	Comparable but more complex	Training objective is harder to optimize
Without velocity coupling	Decreased	Velocity variables provide crucial gradient information
Without symmetry handling	Decreased	Periodic symmetry is crucial for data efficiency

Key Findings

• The Kinetic Langevin method effectively transforms toroidal diffusion into flat space operations.
• Proper handling of periodic symmetry is essential for generating valid crystal structures.
• KLDM achieves performance competitive with current state-of-the-art methods on both CSP and DNG tasks.
• The method holds a solid theoretical foundation rather than being purely empirical.

Highlights & Insights

1. Physics-inspired: Obtains inspiration from Kinetic Langevin Dynamics, using velocity variables to resolve geometric diffusion issues.
2. Symmetry-aware: The training objective naturally encodes the physical symmetries of crystals.
3. Theoretically rigorous: Not a heuristic solution, with clear mathematical derivations.
4. Meaningful Generalization of TDM: Extension of Trivialized Diffusion to handle group symmetries.

Limitations & Future Work

1. Introducing velocity variables increases the dimensionality of the diffusion space, potentially affecting sampling efficiency.
2. Currently validated only on relatively small datasets (Perov-5, MP-20).
3. Lacks a direct comparison with the latest Flow Matching methods.
4. Capability for conditional generation (given target properties) is not fully demonstrated.

Related Work & Insights

• The DiffCSP series serves as the main baseline for crystal diffusion generation.
• FlowMM handles crystal generation using Flow Matching.
• Insight: The Kinetic Langevin method might also be valuable for other generation tasks with periodic/symmetrical structures, such as proteins.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Kinetic Langevin to handle hypertorus diffusion is a novel and elegant idea.
• Experimental Thoroughness: ⭐⭐⭐⭐ Covers two standard tasks, but dataset scales are limited.
• Writing Quality: ⭐⭐⭐⭐ Theoretical derivations are clear.
• Value: ⭐⭐⭐⭐ Provides a new tool for diffusion on manifolds.

Related Papers

• [ICML 2025] CFP-Gen: Combinatorial Functional Protein Generation via Diffusion Language Models
• [ICML 2025] PepTune: De Novo Generation of Therapeutic Peptides with Multi-Objective-Guided Discrete Diffusion
• [CVPR 2025] Multimodal Protein Language Models for Enzyme Kinetic Parameters: From Substrate Recognition to Conformational Adaptation
• [CVPR 2025] DiffVsgg: Diffusion-Driven Online Video Scene Graph Generation
• [ICLR 2026] Controllable Diffusion-based Generation for Multi-channel Biological Data