[Paper Review] Differentiable Molecular Simulations for Control and Learning

Wujie Wang, Simon Axelrod+1 more|arXiv (Cornell University)|Feb 27, 2020

Machine Learning in Materials Science115 references32 citations

TL;DR

The paper demonstrates differentiable molecular dynamics to learn and control Hamiltonians so that simulated observables match targets, enabling parameter inference and active control of molecular systems.

ABSTRACT

Molecular dynamics simulations use statistical mechanics at the atomistic scale to enable both the elucidation of fundamental mechanisms and the engineering of matter for desired tasks. The behavior of molecular systems at the microscale is typically simulated with differential equations parameterized by a Hamiltonian, or energy function. The Hamiltonian describes the state of the system and its interactions with the environment. In order to derive predictive microscopic models, one wishes to infer a molecular Hamiltonian that agrees with observed macroscopic quantities. From the perspective of engineering, one wishes to control the Hamiltonian to achieve desired simulation outcomes and structures, as in self-assembly and optical control, to then realize systems with the desired Hamiltonian in the lab. In both cases, the goal is to modify the Hamiltonian such that emergent properties of the simulated system match a given target. We demonstrate how this can be achieved using differentiable simulations where bulk target observables and simulation outcomes can be analytically differentiated with respect to Hamiltonians, opening up new routes for parameterizing Hamiltonians to infer macroscopic models and develop control protocols.

Motivation & Objective

• Motivate learning and control of molecular Hamiltonians to match macroscopic targets.
• Show that observables derived from MD trajectories can be differentiated with respect to Hamiltonian parameters.
• Demonstrate learning of a control Hamiltonian that biases dynamics toward target states.
• Illustrate learning from observables such as pair distribution functions through differentiable simulations.

Proposed method

• Apply automatic differentiation to molecular dynamics with a control Hamiltonian H_b.
• Use adjoint sensitivity methods to backpropagate through the forward ODE integration.
• Represent control Hamiltonians with neural architectures (e.g., Graph Neural Networks) to preserve physical symmetries.
• Demonstrate learning both for classical trajectory targets and quantum dynamics control scenarios.
• Implement differentiable histogram via Gaussian smearing to enable gradient-based fitting of pair distributions.

Experimental results

Research questions

• RQ1 Can differentiable simulations learn a control/bias Hamiltonian that steers a system toward a target state without predefined reaction coordinates?
• RQ2 Can target macroscopic observables derived from MD trajectories be matched by learning the environment/control Hamiltonian?
• RQ3 How effectively can adjoint-based backpropagation optimize both classical and quantum control tasks within differentiable MD?
• RQ4 How can differentiable simulations be used to learn light-driven quantum control protocols for molecular systems?

Key findings

• A differentiable MD framework can learn a control Hamiltonian to steer a system toward a target state in toy and polymer examples.
• A Graph Neural Network can parameterize H_b to bias a polymer chain into a helix while maintaining temperature via Nose–Hoover chains.
• Differentiable histograms enable learning to match target pair distribution functions of liquids like water.
• Backpropagation through MD trajectories can improve quantum yield optimization in light-driven retinal models.
• Quantum control experiments show improved yield via learned electric field profiles through differentiable simulations.