Ouro
  • Docs
  • Blog
Join for freeSign in
  • Home
  • Teams
  • Search
Assets
  • Quests
  • Posts
  • APIs
  • Data
  • Home
  • Teams
  • Search
Assets
  • Quests
  • Posts
  • APIs
  • Data
8mo
150 views

On this page

  • 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...
Loading compatible actions...

Experimenting with MatterGen and New Denoising Rewards

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)xt−1=μθ(xt,t)+σtz,z∼N(0,I)xt−1=μθ(xt,t)+σtz,z∼N(0,I)

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

  • Applies Langevin dynamics in the correction step:

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

    where ϵϵϵ is the step size and zi∼N(0,I)zi∼N(0,I)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)Lt−1=symmetrize(μθ(Lt,t)+σtZ)Lt−1=symmetrize(μθ(Lt,t)+σtZ)

    where LLL is the lattice matrix and symmetrizesymmetrizesymmetrize 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)IQt(xt∣x0)=vtvt⊤+(1−∥vt∥22)IQt(xt∣x0)=vtvt⊤+(1−∥vt∥22)I

    where vtvtvt is the corruption vector at time ttt.

  • Direct x0x0x0 prediction mode enables better categorical sampling:

    pθ(xt−1∣xt)=∑x0p(xt−1∣xt,x0)pθ(x0∣xt)pθ(xt−1∣xt)=∑x0p(xt−1∣xt,x0)pθ(x0∣xt)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)scoreguided(x,t)=scorebase(x,t)+λR(x)scoreguided(x,t)=scorebase(x,t)+λR(x)

where:

  • scorebasescorebasescorebase is the original MatterGen score function.

  • λλλ is the guidance scale.

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

  • ttt is the diffusion timestep.

Our reward in this approach leaned on CHGNet to calculate the magnetic density of the structure for each tilmestep ttt . 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)R(x)=(m5N4)⋅∏iBi(m)R(x)=(m5N4)⋅∏iBi(m)

where:

  • mmm is the total magnetic moment.

  • NNN is the number of atoms.

  • Bi(m)Bi(m)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...
    1 reference
    • 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

      8mo