Experimenting with MatterGen and New Denoising Rewards

# magnetic_guidance.py
import torch
from typing import Optional
from mattergen.diffusion.sampling.classifier_free_guidance import GuidedPredictorCorrector
from mattergen.diffusion.sampling.classifier_free_guidance import BatchTransform
from mattergen.diffusion.sampling.reward_functions import BaseRewardFunction

class PropertyGuidedPredictorCorrector(GuidedPredictorCorrector):
    """
    Generic sampler for property-guided generation using custom reward functions.
    """

    def __init__(
        self,
        *,
        guidance_scale: float,
        reward_function: BaseRewardFunction,
        remove_conditioning_fn: BatchTransform,
        keep_conditioning_fn: Optional[BatchTransform] = None,
        **kwargs,
    ):
        """
        Args:
            guidance_scale: Controls strength of property guidance
            reward_function: Instance of BaseRewardFunction to compute rewards
            remove_conditioning_fn: Function that removes conditioning from the data
            keep_conditioning_fn: Function applied before evaluating conditional score
            **kwargs: Passed to parent class constructor
        """
        super().__init__(
            guidance_scale=guidance_scale,
            remove_conditioning_fn=remove_conditioning_fn,
            keep_conditioning_fn=keep_conditioning_fn,
            **kwargs
        )
        self.reward_function = reward_function

    def _get_score(self, batch: dict, t: torch.Tensor) -> dict:
        """
        Override parent method to incorporate property guidance via reward function.
        
        Args:
            batch: Dictionary containing structure information
            t: Timesteps tensor
            
        Returns:
            Dictionary containing scores
        """
        # Get base scores from parent class
        base_scores = super()._get_score(batch, t)
        
        # Compute reward using provided reward function
        reward = self.reward_function.compute_reward(batch)
        
        # Scale and reshape reward to match score dimensions
        for key in base_scores:
            if key in ['pos', 'cell']:  # Apply to continuous variables
                scaled_reward = (
                    reward.view(-1, 1, 1) *  # For pos
                    torch.ones_like(base_scores[key])
                )
                base_scores[key] = base_scores[key] + scaled_reward
                
        return base_scores

The implementation uses specialized predictors and correctors from MatterGen’s existing codebase, each designed for different structural components:

1. Atomic Positions (pos)

• Reuses MatterGen’s continuous variable sampler.

• Implements ancestral sampling with the update rule:

  $$xt−1=μθ(xt,t)+σtz,z∼N(0,I)$$

  where $μθ$ is the learned mean and $σt$ is the noise schedule.

• Applies Langevin dynamics in the correction step:

  $$xi+1=xi+ϵ2∇log⁡p(xi)+ϵzi$$

  where $ϵ$

2. Cell Parameters (cell)

• Leverages MatterGen’s specialized lattice predictor.

• Uses a modified ancestral sampling that preserves lattice constraints:

  $$Lt−1=symmetrize(μθ(Lt,t)+σtZ)$$

  where $L$ is the lattice matrix and $symmetrize$ enforces crystal symmetry.

3. Atomic Numbers (atomic_numbers)

• Utilizes MatterGen’s discrete diffusion implementation.

• Implements D3PM (Discrete Denoising Diffusion) with transition matrix:

  $$Qt(xt∣x0)=vtvt⊤+(1−∥vt∥22)I$$

  where $vt$ is the corruption vector at time $t$.

• Direct $x0$ prediction mode enables better categorical sampling:

  $$pθ(xt−1∣xt)=∑x0p(xt−1∣xt,x0)pθ(x0∣xt)$$

The guidance mechanism combines these samplers with our reward function:

$$scoreguided(x,t)=scorebase(x,t)+λR(x)$$

where:

• $scorebase$ is the original MatterGen score function.

• $λ$ is the guidance scale.

• $R(x)$ is our reward function.

• $t$ is the diffusion timestep.

Our reward in this approach leaned on CHGNet to calculate the magnetic density of the structure for each tilmestep $t$ . Because this density value depends on many variables that are constantly changing during the diffusion process (individual magnetic moments, unit cell volume, and moment alignment), we did not see any meaningful move towards higher densities in the resulting structures. There simply isn't enough information in our reward in this case to indicate to the model that maybe maximizing moments and minimizing volume is the way to get high rewards.

So the logical next step here was to really simplify the reward process just to see if we could meaningfully influence the denoising process.

The Pivot: Atom-Only Guidance

The categorical nature of atomic number selection meant that guidance on this component could be more direct and effective than continuous variables.

We simplified our approach to focus solely on atomic rewards, implementing a new MomentMaximizationReward class with progressive scaling:

The reward computation follows:

$$R(x)=(m5N4)⋅∏iBi(m)$$

where:

• $m$ is the total magnetic moment.

• $N$ is the number of atoms.

• $Bi(m)$ are threshold bonus functions:

  Bi(m)={kiif m>ti1otherwise

Simplicity was the goal here. Given we don't have the cash to train a new MatterGen from scratch, the aim was to see what we could build on top of this base model and evaluate how the generation process might be influenced.

As luck would have it, this seems to have worked. We were able to influence MatterGen towards generating structures with cumulative moments 20-40 times larger than the unaltered base model.

The rewards shown are the final values that we settled on following iterations where each new iteration stepped up the reward from the last until we stopped seeing meaningful increases in average moments.

Latest Results

• All-time best total moment: 42.08 μB (Ba₁Gd₆Te₁₃)

• Second best: 41.00 μB (Eu₆Hg₂P₈)

• Multiple structures exceeding 25 μB

• Highest per-atom moment: 5.62 μB/atom (Gd₄N₁)

2. Key Compositions

Parameters

We're not even close to a breakthrough...

Despite some promising results, the magnetic densities of these structures are less than what we see in Iron-based fridge magnets and are well within MatterGen's training data distribution. As much as it would be nice to simply crank the moment and minimize the volume of the unit cells for these structures, the NdFeB magnets rely on exchange coupling to align magnetic moments in parallel, and this behavior gives rise to their extreme magnetic strength and high coercivity.

Maybe we need to think through some more physics based reward landscapes, incentivize similar coupling behaviors, but keep the troublesome rare earths out of the equation. More to come.

2 references

• Diffusion Experiments for Generating Crystal Structures

  post

  Experiments with diffusion models to generate crystal structures, moving from noisy representations to concrete atomic arrangements. They describe learning how these models can learn structure without strict physical rules, and compare approaches that rely on fixed constraints to ones that let the model discover valid layouts. The author notes limitations in existing crystal generators, such as only producing tiny unit cells and struggling with complex, multi-atom systems like NdFeB. To address this, they explore modeling larger supercells with hundreds of atoms to improve detail and tolerance to errors, potentially revealing new properties through dopants. They keep a running experiment log in an AI notebook and plan to explore conditioning methods and the difference between flow and diffusion approaches in future work.

  @will

  #materials-science1mo

• General materials discovery pipeline

  post

  I wanted to formalize in writing the idea that I keep coming back to for end-to-end material discovery. The hardest part of this project has been actually optimizing towards materials that have some p

  @mmoderwell

  #materials-science1y