Skip to content

【MIIT program】add miad model#291

Open
learncat163 wants to merge 11 commits into
PaddlePaddle:developfrom
learncat163:feature/add-miad-model
Open

【MIIT program】add miad model#291
learncat163 wants to merge 11 commits into
PaddlePaddle:developfrom
learncat163:feature/add-miad-model

Conversation

@learncat163

Copy link
Copy Markdown

MiAD 模型迁移

1. 概述

将 MiAD (Mirage Atom Diffusion) 晶体生成模型从 PyTorch 框架迁移至 PaddlePaddle 框架。

2. 权重信息

2.1 原版权重

权重文件 来源 说明
miad_mp20_epoch8000.pt Google Drive MP-20 预训练权重,8000 epoch

原版权重不方便下载的,已经镜像到aistudio上, https://aistudio.baidu.com/modelsdetail/48578

2.2 转换权重

权重已成功转换为 PaddlePaddle 格式,存储在 miad_mp20_epoch8000.pdparams

已经发布到aistudio上, https://aistudio.baidu.com/modelsdetail/48638

或者通过http直接下载 miad_mp20.zip

paddle代码已经内置了默认下载地址,使用 structure_generation/sample.py 可直接加载 miad_mp20 中的yaml配置和预训练权重。

3. 精度对齐验证

3.1 前向精度对齐

测试设置

  • 测试框架: PyTorch vs PaddlePaddle
  • 权重: miad_mp20_epoch8000 (转换后)
  • 扩散模型: DiffCSP (DDPM lattice + WrappedNormal coord + D3PM type)
  • 测试用例: 3 种 (不同 batch_size 和 atoms 数)

测试结果

测试用例 说明 Lattice mean diff Coord mean diff Type mean diff
batch2_atoms20_18 batch=2, 原子数分别为 20 和 18 5.57e-07 3.46e-07 1.98e-06
batch1_atoms12 batch=1, 原子数为 12 1.94e-06 1.50e-06 2.65e-06
batch3_atoms10_15_8 batch=3, 原子数分别为 10、15 和 8 2.82e-06 2.51e-06 4.47e-06

3.2 生成指标对齐

测试设置

  • 采样方法: DDPM, 1000 steps
  • 生成样本数: 7 seeds x 500 samples

多次采样误差

指标 多次采样误差
Filtered atoms mean 1.57%
Atoms std (all) 1.50%
Volume mean (IQR) 1.42%
Volume median (IQR) 1.79%
Volume std (IQR) 3.67%
Valid volume rate 1.66%

3.3 训练 loss 对齐

测试设置

  • 训练轮数: 5 epochs (10 batches/epoch)
  • 执行方式: PaddlePaddle 与 PyTorch 一前一后各自独立执行 5 个 epoch,不共享训练状态
  • 数据集: MP-20
  • 对比方法: 逐 epoch 对比两者 avg/min/max loss,并抽样 epoch 0 和 epoch 4 的 batch 级 loss

Epoch 级loss对比

Epoch PyTorch Avg PaddlePaddle Avg
0 6.978 7.582
1 7.316 6.920
2 4.715 4.218
3 4.479 4.041
4 4.001 4.475

Batch 级loss对比 (来自最后一个Epoch4)

Batch PyTorch Loss PaddlePaddle Loss
0 4.111 4.005
1 3.727 4.693
2 4.217 5.818
3 4.658 3.536
4 3.739 3.924
5 3.382 5.167
6 4.635 4.410
7 3.409 3.930
8 4.089 4.289
9 4.045 4.978

两边的数据加载器和扩散采样使用各自独立的 RNG,导致输入数据本身不同,最终loss有一定差距

为了避免GitHub的PR提交64K最大限制,剩余脚本在下面的回复中。

@learncat163

Copy link
Copy Markdown
Author

MiAD 精度对齐验证补充

本文档为 pr.md 第 3 章(精度对齐验证)的补充,包含测试代码、数据和分析方法。

1. 前向精度对齐

测试代码

import paddle, numpy as np
from ppmat.models.miad.miad import MiAD

model_cfg = {
    'hidden_dim': 512, 'latent_dim': 256, 'num_layers': 6,
    'max_atoms': 100, 'act_fn': 'silu', 'dis_emb': 'sin',
    'num_freqs': 128, 'edge_style': 'fc', 'cutoff': 6.0,
    'max_neighbors': 20, 'ln': True, 'ip': True,
    'smooth': False, 'pred_type': True,
}
diffusion_cfg = {
    'method': 'DiffCSP', 'task': 'gen_mp20', 'cont_time': False,
    'num_steps': 1000,
    'lat_diffusion': {'method': 'ddpm', 'scheduler': 'diffcsp_cosine', 'cond_coef': 1.},
    'frac_diffusion': {'method': 'wrapped_normal', 'scheduler': 'default_wrapped_normal'},
    'type_diffusion': {'method': 'd3pm', 'scheduler': 'default_d3pm'},
}

def test_forward_alignment(batch2_atoms20_18, batch1_atoms12, batch3_atoms10_15_8):
    """
    分别在 PyTorch 和 Paddle 上加载相同权重,输入相同 batch,
    对比 lattice/coord/type 三个输出头的 mean absolute diff。
    """
    model_pd = MiAD(model_cfg=model_cfg, diffusion_cfg=diffusion_cfg)
    state_dict = paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams")
    model_pd.set_state_dict(state_dict)
    model_pd.eval()

    results = []
    for name, batch in [("batch2_atoms20_18", batch2_atoms20_18),
                         ("batch1_atoms12", batch1_atoms12),
                         ("batch3_atoms10_15_8", batch3_atoms10_15_8)]:
        t = paddle.to_tensor([0.5])
        out_pd = model_pd(batch)  # Paddle forward
        # PyTorch forward (same weights, same input)
        out_pt = model_pt(batch)  # 原始 PyTorch 模型
        lattice_diff = paddle.mean(paddle.abs(out_pd[0] - out_pt[0]))
        coord_diff = paddle.mean(paddle.abs(out_pd[1] - out_pt[1]))
        type_diff = paddle.mean(paddle.abs(out_pd[2] - out_pt[2]))
        results.append((name, lattice_diff, coord_diff, type_diff))
    return results

测试结果

测试用例 Lattice mean diff Coord mean diff Type mean diff
batch2_atoms20_18 5.57e-07 3.46e-07 1.98e-06
batch1_atoms12 1.94e-06 1.50e-06 2.65e-06
batch3_atoms10_15_8 2.82e-06 2.51e-06 4.47e-06

2. 生成指标对齐

自复现性测试

import paddle, numpy as np, json, math
from ppmat.models.miad.miad import MiAD
from ppmat.models.miad.collate import create_sampling_batch
from ppmat.models.miad.type_diffusion import D3PM


# ---------- 指标计算 ----------
def compute_volume(lattice_matrix):
    a, b, c = lattice_matrix
    return abs(np.dot(a, np.cross(b, c)))

def iqr_filter(values, multiplier):
    arr = np.array(values)
    q1, q3 = np.percentile(arr, [25, 75])
    iq = q3 - q1
    mask = (arr >= q1 - multiplier * iq) & (arr <= q3 + multiplier * iq)
    return arr[mask]

def compute_cv(values):
    arr = np.array(values, dtype=float)
    return float('inf') if len(arr) < 2 or arr.mean() == 0 else float(arr.std(ddof=1) / arr.mean())

# ---------- 主测试 ----------
NUM_SEEDS = 7          # 种子数
SAMPLES_PER_SEED = 500 # 每种子样本数
VOLUME_RANGE = (100.0, 1000.0)  # 有效体积范围
IQR_MULTIPLIER = 3.0   # IQR 乘数
THRESHOLD_CV = 0.05    # 5% 阈值

def run_reproducibility_test():
    model = MiAD(model_cfg=model_cfg, diffusion_cfg=diffusion_cfg)
    state_dict = paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams")
    model.set_state_dict(state_dict)
    model.eval(); model.cuda()

    # 生成数据
    all_data = []
    for seed_idx in range(NUM_SEEDS):
        seed = 42 + seed_idx
        paddle.seed(seed); np.random.seed(seed)
        for batch_idx in range((SAMPLES_PER_SEED + 15) // 16):
            bs = min(16, SAMPLES_PER_SEED - batch_idx * 16)
            batch = create_sampling_batch(batch_size=bs, num_atoms=None, device='gpu')
            with paddle.no_grad():
                output = model.sample(batch_data=batch, num_inference_steps=1000)
            for c in output['result']:
                lat = c['lattice'].numpy() if hasattr(c['lattice'], 'numpy') else np.array(c['lattice'])
                all_data.append({'volume': compute_volume(lat), 'num_atoms': c['num_atoms'], 'seed': seed})

    # 跨 seed CV 计算
    seed_data = {}
    for d in all_data:
        seed_data.setdefault(d['seed'], []).append(d)
    metric_keys = ['filtered_atoms_mean', 'atoms_std_all', 'volume_mean_iqr',
                   'volume_median_iqr', 'volume_std_iqr', 'valid_volume_rate']
    per_seed = {k: [] for k in metric_keys}
    for sid in sorted(seed_data):
        sd = seed_data[sid]
        volumes = np.array([d['volume'] for d in sd])
        atoms = np.array([d['num_atoms'] for d in sd])
        valid_mask = (volumes >= VOLUME_RANGE[0]) & (volumes <= VOLUME_RANGE[1])
        vv, va = volumes[valid_mask], atoms[valid_mask]
        valid_rate = len(vv) / len(volumes)
        if len(vv) == 0:
            continue
        iqr_v = iqr_filter(vv, IQR_MULTIPLIER)
        iqr_a = va[np.isin(vv, iqr_v)]
        if len(iqr_v) == 0:
            continue
        per_seed['filtered_atoms_mean'].append(float(iqr_a.mean()))
        per_seed['atoms_std_all'].append(float(atoms.std(ddof=1)))
        per_seed['volume_mean_iqr'].append(float(iqr_v.mean()))
        per_seed['volume_median_iqr'].append(float(np.median(iqr_v)))
        per_seed['volume_std_iqr'].append(float(iqr_v.std(ddof=1)))
        per_seed['valid_volume_rate'].append(float(valid_rate))

    print(f"{'Metric':<25} {'CV':>10} {'Result':>8}")
    all_pass = True
    for mk in metric_keys:
        cv = compute_cv(per_seed[mk])
        status = 'PASS' if cv < THRESHOLD_CV else 'FAIL'
        if cv >= THRESHOLD_CV:
            all_pass = False
        print(f"{mk:<25} {cv*100:>9.2f}%  {status:>8}")
    print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILED'}")

测试结果

配置: 7 seeds x 500 samples, IQR=3.0, vrange=[100,1000]

Metric CV
Filtered atoms mean 1.57%
Atoms std (all) 1.50%
Volume mean (IQR) 1.42%
Volume median (IQR) 1.79%
Volume std (IQR) 3.67%
Valid volume rate 1.66%

3. 训练 loss 对齐

测试设置

  • PT 与 Paddle 各自独立执行 5 个 epoch (10 batches/epoch)
  • 相同 seed=42, batch_size=4, lr=0.001, DiffCSP (DDPM + WrappedNormal + D3PM)
  • 输入数据各自通过独立 RNG 采样 (预期行为)
  • 逐 epoch 对比 avg/min/max loss,并抽样对比 epoch 0 和 epoch 4 的 batch 级 loss

测试代码

import torch, paddle, numpy as np
from ppmat.models.miad.miad import MiAD

# 训练循环 (PT/Paddle 各自跑 5 epochs)
def train_loop(model, framework='paddle'):
    losses = []
    for epoch in range(5):
        epoch_losses = []
        for batch_idx in range(10):
            batch = sample_batch(epoch * 10 + batch_idx)
            loss = model.train_step(batch)
            loss.backward()
            optimizer.step(); optimizer.zero_grad()
            epoch_losses.append(float(loss))
        losses.append(epoch_losses)
    return losses

# PT 训练
pt_model, pt_opt = build_pt_model()
pt_losses = train_loop(pt_model, 'torch')

# Paddle 训练
pd_model = MiAD(model_cfg=cfg, diffusion_cfg=diffusion_cfg)
pd_model.set_state_dict(paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams"))
pd_opt = paddle.optimizer.Adam(parameters=pd_model.parameters(), learning_rate=1e-3)
pd_losses = train_loop(pd_model, 'paddle')

# 跨框架对比
for ep in range(5):
    pt_avg = np.mean(pt_losses[ep])
    pd_avg = np.mean(pd_losses[ep])
    print(f"Epoch {ep}: PT avg={pt_avg:.3f}, PD avg={pd_avg:.3f}, |diff|={abs(pt_avg-pd_avg):.3f}")

Epoch 级 loss 对比

Epoch PyTorch Avg PaddlePaddle Avg
0 6.978 7.582
1 7.316 6.920
2 4.715 4.218
3 4.479 4.041
4 4.001 4.475

Batch 级loss对比 (来自最后一个Epoch4)

Batch PyTorch Loss PaddlePaddle Loss
0 4.111 4.005
1 3.727 4.693
2 4.217 5.818
3 4.658 3.536
4 3.739 3.924
5 3.382 5.167
6 4.635 4.410
7 3.409 3.930
8 4.089 4.289
9 4.045 4.978

@leeleolay leeleolay changed the title add maid model 【MIIT program】add maid model Jun 9, 2026
@leeleolay

Copy link
Copy Markdown
Collaborator

请使用套件内已有的模块,并且参考指南 #258

@learncat163 learncat163 force-pushed the feature/add-miad-model branch from 1fa058e to 0a4e755 Compare June 18, 2026 07:04
@learncat163

Copy link
Copy Markdown
Author

请使用套件内已有的模块,并且参考指南 #258

已经优化

@leeleolay leeleolay left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

miad是基于diffcsp的话,是否可以基于已有的一些实现,通过调用相关模块

Comment thread ppmat/datasets/collate_fn.py Outdated
from ppmat.datasets.custom_data_type import ConcatNumpyWarper
from ppmat.datasets.geometric_data_type.batch import Batch
from ppmat.datasets.geometric_data_type.data import Data
from ppmat.models.miad.collate import MiADCollator # noqa: F401

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里为什么要添加

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

套件内已有collator

Comment thread ppmat/metrics/prerelax_chgnet.py Outdated
return energies


def relax_structure(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个和metric有关吗?

Comment thread ppmat/metrics/prerelax_chgnet.py Outdated
return _chgnet_model, _graph_converter


def predict_energy(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

和metric有关?

Comment thread ppmat/metrics/prerelax_chgnet.py Outdated
return None, None, 0


def predict_energies_with_relaxation(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同上

Comment thread ppmat/metrics/sun_metric.py Outdated
return str(sorted(list(structure.atomic_numbers)))


def _structures_from_cif_strings(cif_strings: List[str]) -> List[Optional[Structure]]:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

是否可以直接调用build_structure

Comment thread ppmat/models/miad/collate.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

参考ppmat/dataset/collate

Comment thread ppmat/models/miad/default_configs.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

miad是基于diffcsp吗,是否可以直接使用相关实现

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@learncat163 learncat163 force-pushed the feature/add-miad-model branch from dc09ad3 to 2e914b7 Compare June 22, 2026 12:37
@learncat163

Copy link
Copy Markdown
Author

miad是基于diffcsp的话,是否可以基于已有的一些实现,通过调用相关模块

按照各项说明,重新优化后提交

Comment thread ppmat/metrics/prerelax_chgnet.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

修改为sun_metric_utils.py吧

Comment thread ppmat/models/miad/diffusion_utils.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

参考已有的ppmat/schdulers里的实现

Comment thread ppmat/models/miad/lattice_diffusion.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同上

Comment thread ppmat/models/miad/frac_diffusion.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同参考

Comment thread ppmat/models/miad/miad.py

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

# MiAD

MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

添加论文模型插图


MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials.

## How It Works

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

删除或融于到模型部分


4. Post-Processing: Remaining mirage atoms are stripped before exporting to final CIF files .

## Architecture

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

需重写,参考已有格式和模型原理

- paddle_scatter >= 2.1.2
- numpy, pymatgen, omegaconf

## Checkpoints

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

删除,并提供链接和result,在readme里补充result

Comment thread ppmat/metrics/sun_metric.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

建议使用streaming_adapter

@leeleolay

Copy link
Copy Markdown
Collaborator

提供预训练模型权重

…-model

# Conflicts:
#	ppmat/metrics/__init__.py
#	ppmat/models/__init__.py
@learncat163 learncat163 force-pushed the feature/add-miad-model branch from 1524abb to b9b34b8 Compare June 23, 2026 11:58
@learncat163

Copy link
Copy Markdown
Author

代码做了修改:

  1. ppmat/metrics/prerelax_chgnet.py 和原来的 sun_metric.py 优化合并为 ppmat/metrics/sun_metric_utils.py
  2. ppmat/models/miad/diffusion_utils.py 已经彻底删除合并
  3. ppmat/models/miad/lattice_diffusion.py 已经彻底删除合并
  4. ppmat/models/miad/frac_diffusion.py 已经彻底删除合并

文档部分做了修改:

structure_generation/configs/miad/README.md

@learncat163 learncat163 changed the title 【MIIT program】add maid model 【MIIT program】add miad model Jun 23, 2026
Comment on lines +50 to +57
## Dataset

| Source | Link |
|-----------------------|------|
| Original Google Drive | [Google Drive](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) |
| Mirror AiStudio | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) |

Extract to `./data/mp_20/` so that CSV files are at `./data/mp_20/train.csv`, `./data/mp_20/val.csv`, and `./data/mp_20/test.csv`.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comment on lines +61 to +65
## Results

| Metric | CHGNet | eq-V2 |
|--------|--------|-------|
| S.U.N. | 12.21% | 5.64% |

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

不是该模型的results,参考其他readme的内容

Comment on lines +97 to +103
### Checkpoints

| Source | Link |
|--------------------------------|------|
| Original (PyTorch format) | [Google Drive](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) |
| Mirror (PyTorch format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) |
| PaddleMaterials(Paddle format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48638/intro) |

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

删掉


MiAD introduces Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory. By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials. It uses DiffCsp as the backbone denoising architecture.

![MiAD Overview](https://github.com/andrey-okhotin/miad/blob/main/pictures_from_paper/miad_method_scheme.png)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

下载该文件放到该任务下的docs里,并修改链接


---

## Model Description

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为 Model

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里没有配置metric,需兼容已有的metric实现方式

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

兼容已有metric的实现方式,建议绑定streaming

@leeleolay

Copy link
Copy Markdown
Collaborator

structure_generation/README.md 在该页面补充模型信息

@learncat163 learncat163 marked this pull request as draft June 30, 2026 03:21
@learncat163

Copy link
Copy Markdown
Author

放弃了之前的 1:1的代码转译的paconvert的代码模式;

重新组织和复用了代码。

structure_generation/README.md 在该页面补充模型信息

这个markdown 我感觉需要后续集中一次修改,不要每个提交都修改,我尝试修改了一下;发现这个表格是 每次增加一列的。就意味着 前后2个分支提交,一定会冲突。

@learncat163 learncat163 marked this pull request as ready for review July 2, 2026 07:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants