XQueryer Lightweight 算法说明书

项目概述

本项目是 XQueryer 的轻量化改良版本。原项目主页： https://github.com/Bin-Cao/XQueryer。新版本保留 100315 分类任务，但训练数据不再直接导入预生成 XRD，而是从晶体结构数据库 MP500.db 动态模拟粉末 XRD。所有模拟和实验谱图都会对齐到 10-90 度范围内的 3500 个点。

数据格式与下载

请从本项目 GitHub Release 下载 MP500.db，放到 data/MP500.db。该文件是 ASE 数据库，共 100315 条结构。 ASE 的 entry id 从 1 开始；模型内部标签使用 entry_id - 1，因此分类编号为 0..100314。

数据按 XRD 模拟层面划分。train、validation、test 都覆盖全部 100315 个结构，但各 split 使用不同的模拟参数随机种子区间，因此同一结构会在三个 split 中以不同参数组合生成不同 XRD。

动态 XRD 模拟

训练时，数据集会为每个结构随机抽取样品状态和外部环境参数，然后调用 Pysimxrd.generator.parser(database, entry_id, ...)。默认 train 每个结构模拟 20 条 XRD，validation 每个结构 1 条， test 每个结构 2 条。

grainsize: 10-100
prefect_orientation: 两个值均在 -0.3 到 0.3 之间
thermo_vibration: 0-2
zero_shift: 0-2
dis_detector2sample: 300-600
half_height_slit_detector: 3-8
half_height_sample: 1-4
xrd: reciprocal
background_order: 6
background_ratio: 0.02-0.2
mixture_noise_ratio: 0.01-0.1

模型结构

当前模型保留原始 XQueryer 的核心框架：FFT、CNN、Cross-Attention 和分类头。改动重点不是简单删掉注意力，而是把原来非常大的元素到序列映射改成轻量的元素 query，让模型仍然能用元素信息引导 XRD token 匹配。

1. FFT 频域分支

输入 XRD 先归一化，然后做 real FFT。模型构造多组低通频率视图，再通过 inverse FFT 回到实空间，并与原始谱图拼接。这样保留了 XQueryer 对频域信息的利用，可以增强模型对背景、峰展宽和高频噪声的鲁棒性。

2. CNN 峰特征编码器

FFT 得到的多通道谱图进入 residual depthwise-separable 1D CNN。较大的卷积核负责捕捉峰宽、局部峰形和小角度偏移；步长卷积逐层压缩序列，降低后续注意力计算量。相比普通宽卷积，这种结构更轻，但仍保留 XRD 局部峰形信息。

3. Residual Cross-Attention

CNN 特征先池化成固定数量的 XRD tokens。92 维元素存在向量会被编码成多个结构 query，这些 query 通过 multi-head Cross-Attention 去读取 XRD tokens。这样化学组成会参与决定哪些衍射峰更重要。相比原始模型中 element -> 768 x 3500 的大展开，本模型用少量 query 实现相同的元素引导注意力思想。每层注意力都包含残差 Cross-Attention、query self-attention 和 FFN refinement，让 query 可以多轮修正峰匹配结果。

4. 分类头

Cross-Attention 输出的 query 会和 CNN 全局 mean/max 特征、FFT 全局统计量融合，再经过 MLP neck 和 normalized cosine classifier 预测 100315 个 MP 结构类别。相比普通线性分类头，cosine classifier 更强调特征方向的类间间隔，对大量相似结构的检索式分类更合适。默认配置为 base_channels=64、 attn_dim=192、num_heads=6、 num_tokens=96、num_queries=4、 num_attn_layers=2，约 2801 万可训练参数。其中大部分参数来自 100315 分类的最后一层，这是该任务本身带来的主要开销。

与原始 XQueryer 的主要不同

保留的核心思想包括 FFT 增强、CNN 峰特征、元素引导 Cross-Attention 和 MP 结构直接分类。主要改进包括：从 MP500.db 动态模拟训练数据、统一 10-90 度谱图范围、用紧凑元素 query 替代巨大元素序列展开、残差注意力 refinement、CNN/FFT 全局特征融合、以及更适合大类别检索的 cosine 分类头。

训练方法

pip install torch ase scipy tqdm Pysimxrd

torchrun --nproc_per_node=4 trainer.py \
  --db_path data/MP500.db \
  --epochs 100 \
  --batch_size 64 \
  --num_workers 8 \
  --simulations_per_entry 20 \
  --output_dir outputs/lightweight

训练脚本会保存 checkpoint_0000.pth、 latest.pth、每轮 checkpoint 和验证集最优 best.pth。 checkpoint 中包含 EMA 权重，推理时会优先使用 EMA 模型。本地调试可使用 --max_train_entries 和 --max_val_entries 限制结构数量。

模型结构可调参数包括 --base_channels、 --attn_dim、--num_heads、 --num_tokens 和 --num_queries。增大这些参数通常会提高容量，但显存和训练时间也会增加；减小参数适合做消融或更轻量部署。

新增训练策略包括 warmup-cosine 学习率、AdamW、gradient clipping、 label smoothing、top-1/top-5/top-10 指标、可选弱 XRD mixup 和 EMA 验证/保存。这些改动主要是为了在动态模拟噪声和 100315 大分类任务下稳定训练。

推理方法

实验 XRD 使用 CSV 格式，必须包含 angle 和 intensity 两列。推理脚本会把输入谱图线性插值到 10 到 90 度范围内的 3500 个点，然后输出 top-k 预测结构。

python inference.py \
  --checkpoint outputs/lightweight/checkpoints/best.pth \
  --inputs "exp_data/*.csv" \
  --topk 5

文献引用

@article{cao2025xqueryer,
  title={XQueryer: an intelligent crystal structure identifier for powder X-ray diffraction},
  author={Cao, Bin and Zheng, Zinan and Liu, Yang and Zhang, Longhan and Wong, Lawrence WY and Weng, Lu-Tao and Li, Jia and Li, Haoxiang and Zhang, Tong-Yi},
  journal={National Science Review},
  volume={12},
  number={12},
  pages={nwaf421},
  year={2025},
  publisher={Oxford University Press}
}