TLDR: A new training-free caching strategy significantly speeds up molecular geometry generation using flow matching models, reducing inference time by up to 3x (and 7x with other optimizations) while maintaining high sample quality. This method predicts intermediate hidden states across solver steps, addressing a major computational bottleneck in drug discovery and enabling faster sampling of molecular candidates.
The field of molecular design is undergoing a significant transformation thanks to deep learning, particularly with advanced generative models that can create new molecular geometries. These models are crucial for tasks like drug discovery, where the ability to generate high-fidelity molecular structures directly impacts the development of new compounds.
Among the various generative methods, flow matching models have emerged as a leading technology for producing high-quality molecular geometries. Their iterative nature allows them to model complex geometric distributions effectively. However, a major challenge with these models is their computational cost during inference. Generating a single molecule can require hundreds of network evaluations, making the process prohibitively slow when sampling large numbers of molecular candidates, which is often necessary in practical applications like screening hundreds of thousands to over a million compounds.
Existing methods to accelerate molecular generation typically involve additional training, which introduces overhead in terms of data, computation, and time. These include techniques like trajectory re-parameterization, progressive distillation, and latent methods, all of which require modifying and retraining the model.
A recent research paper introduces a novel, complementary, and training-free approach to accelerate molecular geometry generation. This method, called predictive feature caching, is inspired by advancements in image generation and focuses on reusing and forecasting intermediate hidden states during the sampling process. By leveraging features computed in previous steps, the method significantly reduces redundant computations.
The core idea behind predictive feature caching is to exploit the smooth evolution of intermediate activations within the neural network over time. As a molecule’s geometry is generated step-by-step, the internal features of the model also change smoothly. Instead of recalculating these similar features from scratch at every solver step, the system stores them at specific ‘checkpoint’ times. Then, for nearby time steps, it reuses or predicts these features, avoiding a full forward pass through the network.
The paper investigates two main predictive caching strategies: TaylorSeer and Adams–Bashforth (AB) multi-step forecasting. TaylorSeer uses local Taylor expansions to predict future features based on previously computed ones, while Adams–Bashforth employs a linear multi-step recursion using a history of cached outputs. Both methods operate directly on the SE(3)-equivariant backbone of the model, meaning they respect the fundamental symmetries of molecular structures (rotations and translations), ensuring the generated molecules remain physically valid.
Experiments conducted on the GEOM-Drugs dataset demonstrate the effectiveness of this caching strategy. The results show that predictive feature caching can achieve a twofold reduction in wall-clock inference time while maintaining the same sample quality. In some configurations, it can even lead to a speedup of up to 3x compared to the base model, with only minimal degradation in sample quality. Crucially, these gains are cumulative with other general, lossless optimizations. When combined with techniques like graph compilation and TensorFloat-32 (TF32) matrix–multiply kernels, the overall speedup can reach as much as 7x. This means generating 10,000 molecules, which previously took over 14 minutes, can now be done in approximately 2 minutes without significant loss in quality.
Also Read:
- A Faster Approach to Designing Protein Structures
- New Methods Enhance Diffusion Model Fine-Tuning and Flow Model Quality
The research highlights that predictive feature caching offers a practical solution to the inference bottleneck in molecular geometry generation. It’s a ‘drop-in’ method, compatible with pre-trained models, and works alongside existing training-based accelerations and system-level optimizations. This work paves the way for more efficient and scalable molecular design workflows, enabling researchers to sample millions of molecular candidates much faster. You can read the full research paper for more technical details here: Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation.
