🤖 machine learning

Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation

This paper introduces a training-free predictive feature caching strategy for SE(3)-equivariant flow matching models that accelerates molecular geometry generation by up to 7x while maintaining high sample quality by predicting intermediate hidden states across solver steps.

Original authors: Johanna Sommer, John Rachwan, Nils Fleischmann, Stephan Günnemann, Bertrand Charpentier

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Johanna Sommer, John Rachwan, Nils Fleischmann, Stephan Günnemann, Bertrand Charpentier

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake a perfect cake, but instead of a recipe, you have to figure out the exact shape and texture of the batter at every single second of the baking process. You have to check the oven, stir the mix, and measure the temperature hundreds of times just to get one cake out. This is exactly what happens when computers try to design new 3D molecular structures (the "cakes" of the drug world) using current AI models. They are incredibly accurate, but they are also painfully slow because they have to do so many repetitive calculations.

This paper introduces a clever trick called Predictive Feature Caching that speeds up this process without needing to retrain the AI or change its brain. Here is how it works, using some everyday analogies:

The Problem: The "Over-Engineered" Chef

Current AI models for making molecules work like a chef who insists on tasting the soup, adjusting the salt, tasting it again, adjusting the pepper, and tasting it again for every single step of the cooking process. Even though the soup changes very slowly between tastes, the chef still goes through the whole ritual every time. In computer terms, the AI runs its entire "brain" (a neural network) hundreds of times to generate just one molecule. This takes a long time and uses a lot of energy.

The Solution: The "Smart Shortcut"

The authors realized that just like the soup changes slowly, the internal "thoughts" (hidden states) of the AI also change very smoothly and predictably as it works.

Instead of making the AI do the full, expensive calculation every single time, this new method acts like a smart assistant who knows the pattern.

The Checkpoint: Every few steps, the assistant lets the AI do the full calculation and writes down the result in a notebook (this is the "cache").
The Prediction: For the steps in between, instead of asking the AI to do the hard work again, the assistant looks at the last few notes and uses simple math (like drawing a smooth line through the dots) to guess what the next result would have been.
The Result: The assistant skips the heavy lifting for those middle steps, saving a massive amount of time.

The Two "Guessing" Strategies

The paper tests two different ways to make these guesses:

TaylorSeer (The "Future-Proof" Guess): This method looks at the recent history and uses a mathematical formula (Taylor series) to predict where the feature is heading. It's like looking at a car's speed and acceleration to guess exactly where it will be in two seconds.
Adams-Bashforth (The "Pattern" Guess): This method looks at the last few results and draws a line based on that pattern to predict the next one. It's like noticing a trend in your daily coffee consumption and guessing tomorrow's order based on the last three days.

The Results: Faster, Not Worse

The paper claims that by using these shortcuts, they achieved some impressive results:

2x Speedup: They can generate molecules twice as fast with zero loss in quality. The molecules are just as good as the slow method.
3x Speedup: They can go even faster (three times the speed) with only a tiny, almost unnoticeable drop in quality.
Better than Cutting Corners: Usually, if you want something faster, you tell the AI to do fewer steps. But the paper shows that simply cutting steps makes the molecules worse. Their "caching" method is smarter; it keeps the quality high while cutting the time.
The "Bonus" Effect: In some cases, using this caching method actually made the molecules better than the standard method, even when using the same amount of computing power. It seems the "guessing" helps smooth out the path to a better solution.

Why This Matters

Think of this like upgrading a car engine without changing the car's design. The AI models (the cars) are already great, but they are thirsty for fuel (computing power). This method is like installing a turbocharger that lets them go twice as fast without needing a new engine or a new driver's license (retraining).

The authors also note that this trick works alongside other speed-up methods. If you combine this "caching" with other optimizations, you can get up to a 7x speedup. This means that tasks that used to take hours could be done in minutes, allowing scientists to explore millions of potential drug candidates much faster than before.

In short: The paper found a way to let AI "skip the hard work" by predicting the middle steps of its thinking process. It's a free lunch for speed: you get the same (or better) results in half the time, without needing to teach the AI anything new.

Technical Summary: Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation

Problem Statement

Deep generative models, particularly flow matching and diffusion models, have become state-of-the-art for generating high-fidelity molecular geometries (3D coordinates, atom types, and bond orders). However, their practical application in drug discovery is hindered by significant computational costs during inference. Generating a single molecule often requires hundreds of neural network evaluations (steps) to integrate the underlying ordinary differential equation (ODE). When scaling to the hundreds of thousands or millions of samples required for combinatorial screening or virtual screening, this inference latency becomes the primary bottleneck.

Existing acceleration methods, such as trajectory reparameterization, progressive distillation, and latent space generation, generally require additional training, incurring data, compute, and time overhead. Furthermore, simply reducing the number of inference steps (uniform step reduction) often leads to immediate degradation in sample quality, such as invalid molecular structures or poor physical properties.

Methodology

The authors propose Predictive Feature Caching, a training-free strategy that accelerates molecular geometry generation by exploiting the temporal smoothness of hidden state trajectories during the sampling process.

Core Mechanism

The method operates directly on the SE(3)-equivariant backbone of pretrained flow matching models. It recognizes that as the solver steps through time ( $t$ ), the intermediate activations (features) of the neural network evolve smoothly. Instead of recomputing the full network forward pass at every step, the method:

Caches: Evaluates the full backbone at specific "checkpoint" time steps.
Predicts: For intermediate steps within a caching interval $D$ , it substitutes the expensive backbone evaluation with a cheap prediction function $\hat{F}$ based on previously computed features.
Preserves Equivariance: The prediction schemes are linear combinations of cached features and finite differences, which commute with the SE(3) group action, ensuring the generated molecules remain rotation and translation invariant.

Specific Caching Strategies

The paper investigates two predictive schemes adapted from image generation literature, applied specifically to the last block of the network (a "last-block forecast" to avoid recomputing the entire prefix):

TaylorSeer (TS): Uses local Taylor series expansions to forecast future features. Derivatives are approximated via finite differences of previous cached outputs.
Adams–Bashforth (AB): Employs a multi-step linear recursion using the last $j$ cached outputs to predict the current output.

The implementation ensures that the final inference step is always computed via the full backbone to guarantee convergence, while intermediate steps utilize the cached predictions.

Key Contributions

Training-Free Acceleration: The paper introduces the first training-free scheme to accelerate molecular geometry generation with little to no quality loss. It is a "drop-in" replacement for pretrained models, requiring no retraining or fine-tuning.
SE(3)-Equivariant Adaptation: The authors successfully transfer and analyze feature-caching strategies to operate directly on molecular geometries and SE(3)-equivariant representations, proving that the caching operations preserve the necessary geometric symmetries.
Complementarity: The work demonstrates that predictive caching is orthogonal to existing training-based accelerations and system-level optimizations (e.g., graph compilation, TensorFloat-32), allowing for compounded speedups.
Refinement Capability: Beyond acceleration, the authors show that caching can act as a refinement mechanism, improving sample quality relative to the base model when operating at the same effective inference budget.

Experimental Results

The method was evaluated on unconditional and structure-based (protein-pocket conditioned) molecular generation using models like SemlaFlow, Tabasco, FLOWR, and FLOWR.root across datasets including GEOM, SPINDR, and CrossDocked2020.

Speed vs. Quality Trade-off: Caching achieves a 2× reduction in wall-clock inference time while maintaining matched sample quality. With minimal degradation, speedups of up to 3× are achievable.
Superiority over Step Reduction: At lower step budgets (e.g., 1/3 or 1/4 of the original steps), predictive caching significantly outperforms uniform step reduction. For instance, the Tabasco model failed to produce valid molecules at 14 steps without caching, whereas cached variants maintained high validity and physical plausibility.
Quality Improvement: In some configurations, cached models with an effective budget of 100 steps (achieved via 200 total steps with caching) outperformed the standard 100-step base model in metrics like Vina Score and Posebusters Validity.
Compound Optimizations: When combined with graph compilation and TF32 kernels, the approach yielded up to a 7× total speedup, reducing the time to generate 10,000 molecules from >14 minutes to ~2 minutes.

Significance and Claims

The paper claims that predictive feature caching defines a new speed-quality Pareto front for molecular generation. By leveraging the smoothness of feature trajectories, the method allows for the scaling of generative models to the millions of samples required for practical drug discovery without the prohibitive costs of retraining or the quality losses associated with naive step reduction.

The authors position this work as a systematic step toward improving inference-time efficiency, highlighting that caching is not merely a heuristic but a robust mechanism that preserves the physical validity and geometric constraints of generated molecules. The code is made available to facilitate further research and adoption in the molecular design community.