Technical Summary: Transport Novelty Distance (TNovD)

Problem Statement

The rapid advancement of generative machine learning for materials discovery has outpaced the development of rigorous evaluation metrics. Existing approaches, such as the SUN metrics (Stability, Uniqueness, Novelty), evaluate generated structures on an instance-by-instance basis and fail to capture distributional properties. Conversely, distributional metrics like the Fréchet Inception Distance (FID) and its molecular analog (FCD) rely on Gaussian assumptions that compress data into mean and variance, rendering them incapable of detecting "memorization" (where a model simply reproduces training data). Furthermore, current methods struggle to simultaneously assess both the quality (statistical similarity to the training distribution) and novelty (dissimilarity from specific training samples) of generated materials. The paper highlights that distinguishing between high-quality generalization and dangerous memorization is critical, particularly as newer models like diffusion models become more prone to overfitting.

Methodology

The authors propose the Transport Novelty Distance (TNovD), a distributional metric grounded in Optimal Transport (OT) theory and contrastive learning. The methodology consists of three main components:

1. Feature Extraction via Equivariant GNN:

   • Crystal structures are embedded into a feature space using an $SE(3)$-equivariant Graph Neural Network (GNN).
   • The GNN is trained using InfoNCE contrastive loss on the MP20 dataset.
   • Augmentation Strategy: Positive pairs are created by applying rotations, translations, and supercell expansions to the same crystal structure. Negative pairs consist of different materials. This training ensures the feature space is invariant to rigid transformations and supercell size while distinguishing distinct chemical compositions.

2. Optimal Transport Coupling:

   • TNovD calculates the coupling between the feature distributions of the training set ($P_D$) and the generated set ($P_G$) using the Wasserstein distance framework.
   • Unlike standard Wasserstein distances, TNovD introduces a threshold $\tau$ to separate the cost function into two regimes:
     • Memorization Regime ($x < \tau$): If a generated sample is too close to a training sample (distance $x < \tau$), it is penalized heavily using a power function $(\tau - x)^M$. This detects exact or near-exact replicas.
     • Quality Regime ($x \ge \tau$): If a sample is sufficiently novel ($x \ge \tau$), the cost is linear $(x - \tau)$, measuring how far the sample deviates from the training distribution.

3. Hyperparameter Heuristics:

   • $\tau$ (Threshold): Determined heuristically by analyzing the 1-nearest neighbor and 2-nearest neighbor distances within the training set, approximating the condition where a point is considered "memorized."
   • $M$ (Exponent): Derived from an equilibrium condition where the derivative of the TNovD with respect to $\tau$ is zero for two splits of a "good" distribution, ensuring the metric balances the penalty for memorization against the penalty for poor quality.

Key Contributions

• A Unified Distributional Metric: TNovD is a metric designed to jointly evaluate quality and novelty in material generative models without relying on Gaussian assumptions. While prior work [8] proposed a GNN-based distributional metric, TNovD distinguishes itself by explicitly training the GNN for novelty detection and avoiding Gaussian assumptions to better capture memorization.
• Memorization Detection: By utilizing a non-linear cost function for small distances, TNovD explicitly penalizes the reproduction of training data, a capability lacking in FID and FCD.
• Domain-Agnostic Framework: While applied to crystals, the authors note the framework is adaptable to other domains (e.g., images, molecules) by swapping the feature extractor.
• Robust Feature Learning: The use of contrastive learning with specific crystallographic augmentations (supercells, rotations) creates a feature space that respects physical symmetries and distinguishes between true novelty and structural variations of the same material.

Results

The authors validated TNovD through several experiments:

1. Toy Experiments (MP20 Validation Set):

   • Memorization: When validation data was progressively replaced with training data, TNovD showed a near-linear increase, successfully detecting memorization.
   • Noise Injection: Adding Gaussian noise to atomic coordinates or shearing the lattice resulted in increased TNovD scores, indicating the detection of low-quality structures.
   • Invariance: The metric remained stable when supercells were generated, demonstrating the GNN's ability to handle size variations without penalizing them as "novelty."
   • Substitutions: Random and in-group atom substitutions were correctly identified as lower-quality materials, with random substitutions penalized more heavily.

2. WBM Substitution Dataset:

   • Applied to the WBM dataset (structures created by iterative atom substitution), TNovD showed that as structures diverged from the MP20 training set, the memorization component dropped to near zero, while the quality component increased slightly, reflecting the growing dissimilarity.

3. Benchmarking Generative Models:

   • Six popular material generative models (including MatterGen, CDVAE, DiffCSP, and ADiT) were evaluated on 10,000 generated samples.
   • ADiT achieved the best overall score, driven by high quality, though the authors hypothesize it may still memorize data in a way not fully captured by the metric.
   • MatterGen performed similarly to the validation set baseline.
   • CDVAE performed poorly, showing high quality penalties (low quality) and no memorization, indicating a failure to match the training distribution.
   • Chemeleon showed the highest memorization component.

Significance and Claims

The paper claims that TNovD offers a necessary "complementary perspective" to existing instance-level metrics. Its primary significance lies in its ability to detect memorization while simultaneously assessing distributional quality, addressing a critical gap in the evaluation of modern generative models.

The authors modestly note that:

• TNovD is not a training objective, as a generator could theoretically "game" the metric; it is strictly an evaluation tool.
• The metric does not directly capture material stability (e.g., energy above hull) or synthesizability, though it can be extended with pre-trained energy models in future work.
• The results depend heavily on the chosen feature representation (MP20-based GNN) and augmentations.
• The low memorization scores observed in many models suggest that the primary challenge for current material generative models is fitting the correct data distribution (quality) rather than avoiding memorization.

The work establishes a modular blueprint for holistic evaluation, emphasizing that as generative models advance, the ability to distinguish between genuine discovery and training set replication becomes paramount.