FK-ReMask
Multi-property Feynman-Kac Steering of Crystal Diffusion Models
Abstract
Pretrained diffusion models provide powerful priors for generating new materials, including inorganic crystals, but generating useful materials requires satisfying specific chemical and physical properties, from mechanical stiffness to bandgaps. Existing approaches either retrain models for new target properties or steer only the atom coordinates and lattice during generation, leaving composition unchanged despite its fundamental role in determining physical properties. We present FK-ReMask, an inference-time, gradient-free method that steers both the continuous and discrete channels of a pretrained crystal diffusion model. During generation, FK-ReMask maintains multiple denoising trajectories and replicates promising candidates using Feynman–Kac particle selection. However, selection alone is insufficient: in the atom-type channel, committed elements cannot be revised, so cloning propagates existing compositions and limits the population's ability to explore new ones. We address this with an atom re-masking mechanism, which reopens committed sites during the generation process and allows alternative compositions to emerge. Across four electronic, photovoltaic, and mechanical objectives, our method improves the yield of SUN (stable, unique, novel) and property-on-target materials. We further demonstrate joint steering of three photovoltaic properties, showing that multiple physical objectives can be specified entirely at inference time without retraining.
Method
FK-ReMask steers a frozen, pretrained crystal diffusion model (MatterGen) toward target properties at inference time, with no retraining and no gradients. It combines two ideas:
- Feynman–Kac selection. Several candidate crystals (particles) are denoised in parallel. At regular steps, a property predictor scores each one; promising particles are cloned and weak ones are discarded.
- Atom re-masking. Once an atom type is chosen, the model normally never changes it, so cloning only copies existing compositions. Re-masking reopens committed atom sites, so the population can keep exploring new compositions.
Overview of FK-ReMask. (a) Within a single trajectory, the model denoises the crystal while atom re-masking reopens committed atom sites, so the composition can still change. (b) Across trajectories, reward-guided Feynman–Kac selection prunes weak candidates and clones promising ones.
Atom re-masking under selection. Four particles go through denoising, selection and re-masking; cells show the element at each atom site. (a) Without re-masking, selection hands the winning particles more slots, which reduces the diversity of the final compositions. (b) With re-masking, committed sites reopen inside the re-masking window, yielding four distinct formulas.
Single-Objective Steering
Goal: generate stable, novel crystals that hit one requested property: bandgap, solar-cell efficiency, hole effective mass or bulk modulus.
Result: FK-ReMask produces the most stable, on-target crystals on all four properties, clearly ahead of Best-of-N at the same compute. Re-masking keeps the population diverse after selection.
Distribution of target properties. Property histograms of crystals generated by each method, with the on-target region shaded.
Bandgap steering with and without re-masking. Dashed lines mark where selection and re-masking begin. (a) Selection collapses the population; with re-masking on, diversity is restored. (b) Re-masking reaches a higher mean particle reward.
Multi-Objective Steering
Goal: design solar-cell materials by hitting three properties at once: bandgap, solar-cell efficiency and hole effective mass.
Result: MatterGen almost never meets all three targets. FK-ReMask meets them far more often than Best-of-N, while also producing more stable crystals.
| on target | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| method | valid | unique | novel | stable | SUN | Eg | SLME | m*h | all | stable & all | SUN & all |
| MatterGen | 86.3 | 99.6 | 59.2 | 75.0 | 37.5 | 1.0 | 2.5 | 5.7 | 0.2 | 0.0±0.0 | 0.0±0.0 |
| + Best-of-N | 83.8 | 99.8 | 66.8 | 66.6 | 36.7 | 7.6 | 21.3 | 26.0 | 1.4 | 1.0±0.4 | 0.4±0.3 |
| + FK-ReMask | 84.0 | 97.5 | 66.4 | 77.5 | 46.1 | 27.0 | 68.8 | 58.6 | 11.5 | 9.2±1.3 | 5.3±1.0 |
Multi-property results (%). Jointly targeting three photovoltaic objectives. all denotes the conjunction of the three on-target conditions.
Example crystals generated by FK-ReMask under joint steering of three photovoltaic properties (bandgap, solar-cell efficiency and hole effective mass). Property values are predicted with pretrained ML models.
BibTeX
@article{mousai2026fkremask,
title = {Multi-property Feynman-Kac Steering of Crystal Diffusion Models},
author = {Mousai, Oded and Yiflach, Sapir Esther and Davidovitch, Barak Itzchak and
Mualem, Yuval and Kraus, Sarit and Barad, Hannah-Noa and Kreis, Karsten and
Chechik, Gal},
journal = {arXiv preprint arXiv:2610.24828},
year = {2026},
url = {https://arxiv.org/abs/2610.24828}
}