Experimenting with MatterGen and New Denoising Rewards

2025-03-02T23:10:10.101+00:00

Added by will to #permanent-magnets.

Some Background

MatterGen employs a diffusion-based approach for crystal structure generation, utilizing classifier-free guidance to steer the generation process. The core of our modifications centers on the PropertyGuidedPredictorCorrector class in magnetic_guidance.py, which implements a new guidance mechanism:

python

# 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

python

from typing import Dict, Optional, Union
from mattergen.generator import CrystalGenerator
from mattergen.diffusion.sampling.magnetic_guidance import PropertyGuidedPredictorCorrector
from mattergen.diffusion.sampling.reward_functions import (
    BaseRewardFunction,
    MagneticRewardFunction,
    CompositeRewardFunction
)
from mattergen.property_embeddings import SetUnconditionalEmbeddingType, SetConditionalEmbeddingType
from mattergen.diffusion.wrapped.wrapped_predictors_correctors import (
    WrappedAncestralSamplingPredictor,
    WrappedLangevinCorrector
)
from mattergen.diffusion.d3pm.d3pm_predictors_correctors import D3PMAncestralSamplingPredictor
from mattergen.common.diffusion.predictors_correctors import LatticeAncestralSamplingPredictor

def create_property_guided_generator(
    generator: CrystalGenerator,
    reward_function: Union[BaseRewardFunction, Dict[str, BaseRewardFunction]],
    guidance_scale: float = 1.0,
    reward_weights: Optional[Dict[str, float]] = None
) -> CrystalGenerator:
    """
    Create a CrystalGenerator that uses property guidance during generation.
    
    Args:
        generator: Base CrystalGenerator instance
        reward_function: Either a single BaseRewardFunction or a dictionary of
            named reward functions to combine
        guidance_scale: Controls strength of property guidance
        reward_weights: Optional weights for combining multiple reward functions.
            Only used if reward_function is a dictionary.
        
    Returns:
        Modified CrystalGenerator with property guidance
    """
    # Handle multiple reward functions
    if isinstance(reward_function, dict):
        reward_function = CompositeRewardFunction(
            reward_functions=reward_function,
            weights=reward_weights
        )
    
    # Get the original sampler config
    cfg = generator.cfg
    
    # Create new property guided sampler
    sampler = PropertyGuidedPredictorCorrector(
        guidance_scale=guidance_scale,
        reward_function=reward_function,
        diffusion_module=generator.model.diffusion_module,
        predictor_partials={
            'pos': lambda corruption, score_fn: WrappedAncestralSamplingPredictor(corruption=corruption, score_fn=score_fn),
            'cell': lambda corruption, score_fn: LatticeAncestralSamplingPredictor(corruption=corruption, score_fn=score_fn),
            'atomic_numbers': lambda corruption, score_fn: D3PMAncestralSamplingPredictor(corruption=corruption, score_fn=score_fn, predict_x0=True)
        },
        corrector_partials={
            'pos': lambda corruption, n_steps, score_fn: WrappedLangevinCorrector(corruption=corruption, score_fn=score_fn, n_steps=n_steps, max_step_size=1e6, snr=0.4),
            'cell': lambda corruption, n_steps, score_fn: LatticeAncestralSamplingPredictor(corruption=corruption, score_fn=score_fn)
        },
        device=generator.model.device,
        n_steps_corrector=10,
        N=1000,
        eps_t=1e-3,
        remove_conditioning_fn=SetUnconditionalEmbeddingType(),
        keep_conditioning_fn=SetConditionalEmbeddingType()
    )
    
    # Update the generator's sampler
    generator._model.sampler = sampler
    
    return generator


def create_magnetic_guided_generator(
    generator: CrystalGenerator,
    target_magnetic_moment: float,
    guidance_scale: float = 1.0,
    chgnet_model=None
) -> CrystalGenerator:
    """
    Convenience function to create a magnetically-guided generator.
    This is equivalent to using create_property_guided_generator with a MagneticRewardFunction.
    
    Args:
        generator: Base CrystalGenerator instance
        target_magnetic_moment: Target magnetic moment (in μB) to guide towards
        guidance_scale: Controls strength of magnetic property guidance
        chgnet_model: Pre-loaded CHGNet model, will load default if None
        
    Returns:
        Modified CrystalGenerator with magnetic guidance
    """
    reward_function = MagneticRewardFunction(
        target_magnetic_moment=target_magnetic_moment,
        chgnet_model=chgnet_model
    )
    
    return create_property_guided_generator(
        generator=generator,
        reward_function=reward_function,
        guidance_scale=guidance_scale
    )

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

1. Atomic Positions (`pos`)

plaintext

'pos': WrappedAncestralSamplingPredictor(corruption, score_fn)

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 $ϵ$ is the step size and $zi∼N(0,I)$ .

2. Cell Parameters (`cell`)

python

'cell': LatticeAncestralSamplingPredictor(corruption, score_fn)

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`)

python

'atomic_numbers': D3PMAncestralSamplingPredictor(
    corruption, score_fn, predict_x0=True
)

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:

python

class MomentMaximizationReward(BaseRewardFunction):
    """Simple reward function focused solely on maximizing magnetic moments through atom selection."""
    
    def __init__(
        self,
        moment_weight: float = 1.0,
        chgnet_model=None
    ):
        """
        Args:
            moment_weight: Weight for moment maximization term
            chgnet_model: Pre-loaded CHGNet model
        """
        from chgnet.model.model import CHGNet
        self.moment_weight = moment_weight
        self.chgnet_model = chgnet_model or CHGNet.load()
        
        # Ultra-aggressive thresholds and multipliers
        self.thresholds = [
            (5.0, 3.0),     # 3x bonus above 5 μB
            (10.0, 5.0),    # 5x bonus above 10 μB
            (15.0, 8.0),    # 8x bonus above 15 μB
            (20.0, 12.0),   # 12x bonus above 20 μB
            (25.0, 16.0),   # 16x bonus above 25 μB
            (30.0, 20.0)    # 20x bonus above 30 μB
        ]
    
    def compute_reward(self, batch: dict) -> torch.Tensor:
        structures = self._batch_to_structures(batch)
        rewards = []
        
        for structure in structures:
            prediction = self.chgnet_model.predict_structure(structure)
            mag_moments = prediction['m']
            mag_moment = float(sum(mag_moments))
            num_atoms = len(structure)
            
            # More aggressive base reward with seventh power scaling
            base_reward = (mag_moment ** 7) / (num_atoms ** 6)
            
            # Progressive threshold bonuses
            for threshold, multiplier in self.thresholds:
                if mag_moment > threshold:
                    base_reward *= multiplier
            
            # Enhanced alignment bonus with proper alignment calculation
            alignment = compute_alignment(mag_moments)
            if alignment > 0.9:  # Stricter alignment requirement
                base_reward *= 4.0  # Higher alignment bonus
            elif alignment > 0.8:
                base_reward *= 2.0
            
            # Enhanced per-atom moment bonuses
            per_atom_moment = mag_moment / num_atoms
            if per_atom_moment > 2.0:
                base_reward *= 2.0
            if per_atom_moment > 3.0:
                base_reward *= 3.0
            if per_atom_moment > 4.0:  # New higher threshold
                base_reward *= 4.0
            
            rewards.append(base_reward)
        
        return torch.tensor(rewards, device=batch['pos'].device)

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

Composition	Magnetic Moment (μB)
Ba₁Gd₆Te₁₃	42.08
Eu₆Hg₂P₈	41.00
Fe₆O₃F₉	26.63
Eu₄Ga₂Ge₄	27.33
Gd₄N₁	28.10

Parameters

python

# Latest optimized parameter grid
param_grid = {
    'guidance_scale': [30000.0, 35000.0, 40000.0],  # Higher guidance scales
    'moment_weight': [3500.0, 4000.0, 4500.0],      # Higher moment weights
}

# Enhanced batch configuration
BATCH_SIZE = 16  # Increased for better statistics
NUM_BATCHES = 3  # Multiple batches per trial

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.

Loading comments...

posts

Experimenting with MatterGen and New Denoising Rewards

Some Background

1. Atomic Positions (pos)

2. Cell Parameters (cell)

3. Atomic Numbers (atomic_numbers)

The Pivot: Atom-Only Guidance

Latest Results

2. Key Compositions

Parameters

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

General materials discovery pipeline

1. Atomic Positions (`pos`)

2. Cell Parameters (`cell`)

3. Atomic Numbers (`atomic_numbers`)