Sharing an idea for finding "adjacent crystals". Why? In the AI research agent I'm working on, we're trying to discover materials with a set of target properties.
We do this by letting an agent generate and mutate crystals in loop, trying to mutate or generate materials to meet our property requirements. The issue is that when we mutate a crystal (like jitter, scale lattice, add site), we're messing around in crystal space which can cause the resolved crystal (after relaxation and repair) to be very different in structure and properties than the previous generation.
It's hard to make a slight change in the input and see just a slight change in the output properties.
If we could go faster, generate more crystals and evaluate very fast, this may not be an issue as a true genetic algorithm would cover much more of the search space and let randomness do it's thing. But here we are trying to intelligently evolve, guided by the LLM's direction. This is because property evaluation is expensive and we want to act more like gradient decent, finding our ideal in fewer steps.
So I'm thinking of three possible approaches we could apply to Matra-Genoa that might help us do this task better. Check out some of my notes on the approach here:
Check out the paper here. It's a short read. I recommend checking it out. Although not very technical (just machine learning concepts that have been explored elsewhere), the creativity and simplicity
If we remember the tokenized representation of the crystal, we see how generation is started by first defining the structure motif, chemical system, stoichiometric ratio, space group, and so on.
Schematic overview of the invertible sequenced representation. (a) The structure is first decomposed into composition, stability, structure and lattice. (b) The structure is then further decomposed into a set of Wyckoff positions, uniquely identified by a set of Wyckoff identifiers. Optional free parameters are also included to make the representation coordinate-aware. (c) All previous information is gathered into a tokenenized and invertible sequence. The color of the tokens represent the type or the Wyckoff position for ease of visualization.
The idea is to hard-code some of these tokens (fitting our initial crystal), then modifying a part we want to change, maybe just a single Wyckoff position, then going back to hard-coding the other positions, then letting the model generate the rest of the sequence.
There are some properties in the sequence that should be mutable, while others we should definitely let the model predict. For example, we could set composition, SG, Wyckoff positions, but let the model predict the lattice.
The other part is choosing other top logits, next to the current structure's representation. I imagine this to be very useful for finding polymorphs in the same space group. You fix everything up until position, then look at the more probable tokens, and choose high probability yet different tokens than the input structure to mutate it slightly.
The model uses an encoder-decoder-style MLP + attention heads. What if we learned a system that could move us around in the latent space in meaningful ways, then decode the sequence from there.
While previously possible, it seems like this model makes this idea much more approachable than some of the other models out there doing CSP. You could build property predictors on the latent space (something I've done with MLIP latent space before), then find SHAP values and features that influence the property then adjust the latent space and decode from it, straight into crystal space.
We could even combine both approaches such that we decode with some fixed tokens from the new place in latent space.
Just as they trained the model with tokens for energy above the hull, we could retrain or fine tune the model with a set of our own properties. Most straightforward but also most effort (collect training data) and most expensive (GPU time).
This is not as relevant to finding adjacent crystals as it might eliminate the need for trying to evolve a crystal toward desired properties.