AGeDi-WS2-Unconditional
An unconditional atomistic diffusion model for generating WS₂ defect structures.
This model was trained specifically for WS₂ defect generation using a fixed pristine WS₂ host and a learnable defect region. It serves as the unconditional Phase 2 baseline for a broader study of controllable generation of charged WS₂ defects.
Overview
The goal of this model is to generate atomistic WS₂ defect structures while preserving the underlying pristine WS₂ host structure.
The model operates on:
- atomic positions
- atomic species
- a fixed WS₂ host structure
- a localized defect region
The pristine host is kept fixed during sampling, while atoms within the defect region are generated by the diffusion model.
The model is unconditional, meaning that it does not explicitly receive the target defect type, charge state, or formation energy during generation.
Model Architecture
The model uses a PaiNN-based score network with diffusion over atomic positions and atomic types.
| Component | Configuration |
|---|---|
| Framework | AGeDi |
| Representation | PaiNN |
| Atom basis | 64 |
| Interaction blocks | 4 |
| Radial basis functions | 30 |
| Cutoff | 6.0 Å |
| Position noiser | ConfinedCellPositions |
| Atom-type noiser | Types |
| Atom-type classes | 27 |
| Prediction type | Score |
| Reverse sampler | Euler-Maruyama |
| Fully connected | False |
The model learns a joint distribution over:
Atomic positions + atomic species
within the free defect region.
Dataset
The model was trained on a 2D WS₂ defect dataset containing 1,706 structures.
The dataset contains a variety of defect configurations including substitutions, vacancies, and multi-component defects.
Dataset statistics
- 2D structures: 1,706
- Unique defect labels: 418
- Two-component defect labels: 366
- Single-component defect labels: 52
The dataset contains structures with variable numbers of atoms depending on the defect configuration.
Host and Defect Representation
The generation protocol separates each structure into:
Fixed pristine WS₂ host
+
Free defect region
A production defect-region mask with a 2.5 Å radius was used.
The fixed host atoms remain unchanged during sampling.
Only the atoms inside the defect region are generated.
This allows the model to learn localized defect reconstruction while preserving the global WS₂ lattice.
Dataset Split
The 1,706 2D structures were divided into:
| Split | Structures |
|---|---|
| Training | 1,365 |
| Validation | 170 |
| Held-out test | 171 |
The 171 held-out structures were kept completely outside the training trajectory and were used for evaluating the unconditional model after training.
Training
The model was trained for 1,000 epochs on an NVIDIA A100 80 GB GPU.
Training configuration
Epochs : 1000
Batch size : 8
Learning rate : 1e-4
Weight decay : 0
Feature size : 64
PaiNN blocks : 4
Cutoff : 6.0 Å
Training result
Training time : 3 h 13 min 23 s
Best validation loss: 0.028201
The released checkpoint corresponds to the best validation checkpoint from the 1,000-epoch training run.
Checkpoint
The main released checkpoint is:
best_model.ckpt
This checkpoint achieved the best validation loss during training:
0.028201
Repository Contents
AGeDi-WS2-Unconditional/
│
├── best_model.ckpt
├── hparams.yaml
├── ws2_phase2_unconditional.yaml
└── README.md
best_model.ckpt
The trained WS₂ unconditional diffusion model.
hparams.yaml
Model hyperparameters and reconstruction information required to load the trained AGeDi model.
ws2_phase2_unconditional.yaml
Configuration used for the Phase 2 unconditional training run.
README.md
This model card and documentation.
Sampling
The model can generate WS₂ structures by providing a fixed-host template.
The main evaluation sampling configuration was:
Diffusion steps : 500
Sampler : Euler-Maruyama
The z-coordinate confinement used during sampling was:
Minimum z : 4.87254786 Å
Maximum z : 12.67985539 Å
The number of generated atoms depends on the free-atom count of the input defect-region template.
Held-Out Evaluation
The model was evaluated on all 171 held-out structures.
Each held-out structure was provided as a fixed-host template, with the corresponding defect region used to determine the number of atoms generated.
Structural Preservation
| Metric | Result |
|---|---|
| Structures generated | 171 / 171 |
| Correct atom count | 171 / 171 |
| Fixed host preserved | 171 / 171 |
| Cell preserved | 171 / 171 |
| PBC preserved | 171 / 171 |
The model therefore reproduced the required structural infrastructure correctly for all held-out sampling cases.
Fixed-Host Accuracy
The fixed WS₂ host remains unchanged during sampling.
For the held-out evaluation:
Fixed-host preservation: 171 / 171
Cell preservation : 171 / 171
PBC preservation : 171 / 171
In the production sampling verification, the maximum observed fixed-host displacement was approximately:
6.31 × 10⁻⁷ Å
and the maximum cell difference was approximately:
6.75 × 10⁻⁷ Å
These values are effectively numerical zero and confirm preservation of the fixed host and simulation cell.
Physical Quality Evaluation
The generated structures were compared with the corresponding held-out reference structures.
Minimum Interatomic Distance
| Metric | Reference | Generated |
|---|---|---|
| Mean minimum distance | 2.1785 Å | 1.8313 Å |
| Median minimum distance | 2.2695 Å | 1.9723 Å |
| Minimum observed distance | 1.1152 Å | 0.0483 Å |
| Maximum observed distance | 2.4155 Å | 2.3763 Å |
The generated structures reproduce the overall WS₂ geometric scale reasonably well, but some samples contain very short interatomic distances in the generated defect region.
Close-Contact Statistics
For the 171 held-out generated structures:
| Criterion | Reference | Generated |
|---|---|---|
| dₘᵢₙ < 0.5 Å | 0 / 171 | 6 / 171 |
| dₘᵢₙ < 0.8 Å | 0 / 171 | 9 / 171 |
| dₘᵢₙ < 1.0 Å | 0 / 171 | 10 / 171 |
| dₘᵢₙ < 1.2 Å | 1 / 171 | 12 / 171 |
| dₘᵢₙ < 1.5 Å | 1 / 171 | 20 / 171 |
These results indicate that the unconditional model occasionally produces physically implausible close contacts.
Therefore, this model should be considered a research baseline rather than a chemically filtered production generator.
Global Geometric Distribution
The mean all-pair distance was:
Reference : 6.5700 Å
Generated : 6.5683 Å
This indicates that the generated structures maintain a global geometric scale close to that of the held-out WS₂ structures.
The generated z-coordinate statistics also remain close to the reference slab geometry.
Composition
This model is unconditional.
It does not receive the ground-truth:
- defect type
- defect composition
- charge state
- formation energy
as conditioning variables during generation.
Therefore, exact per-sample composition matching against the corresponding held-out reference is not interpreted as conditional prediction accuracy.
For the 171-structure evaluation:
Exact free-region composition match : 1 / 171
Mean composition recall : 0.7628
These values are reported as descriptive statistics of the unconditional generator.
1,000-Sample Benchmark
A larger benchmark of 1,000 generated WS₂ structures was produced using the trained unconditional model.
The benchmark used:
Held-out templates : 171
Generated structures : 1,000
Diffusion steps/sample : 500
All 1,000 structures were successfully generated:
Generated frames : 1,000 / 1,000
Successful templates : 171 / 171
The benchmark dataset is intended for evaluation of:
- structural validity
- close-contact rate
- composition statistics
- uniqueness
- diversity
- defect-space coverage
- novelty
Detailed statistics from the 1,000-sample benchmark will be added after the complete analysis.
Why This Model Matters
This unconditional model provides a baseline for studying whether diffusion models can learn the local structural space of WS₂ defects without explicitly specifying the target defect.
The experiment separates two questions.
1. Structural generation
Can the model generate valid atomistic structures while preserving the WS₂ host?
The held-out evaluation shows successful preservation of:
Host structure
Atom count
Cell
PBC
2. Defect-space generation
Can the model generate diverse and physically reasonable defect configurations?
The current baseline shows that this remains challenging, particularly in the generation of physically plausible local defect geometries.
This motivates the next stage of controlled conditional generation.
Intended Use
This model is intended for research in:
- Materials science
- Atomistic generative modeling
- 2D materials
- WS₂ defect generation
- Diffusion models for atomistic systems
- Defect-space exploration
- Generative AI for materials discovery
The checkpoint is intended primarily as an unconditional baseline.
Limitations
This model has several important limitations.
No explicit defect control
The model is unconditional and does not directly control:
Defect type
Charge state
Formation energy
Target composition
Local structural failures
Some generated structures contain very short interatomic distances.
Such structures require additional filtering, structural relaxation, or energetic validation.
No stability guarantee
Generation by the diffusion model does not imply that a structure is:
- thermodynamically stable
- dynamically stable
- experimentally realizable
- energetically favorable
Additional materials-science validation is required.
Not intended for direct experimental prediction
This checkpoint should be treated as a research artifact and baseline model rather than a production materials-screening system.
Future Direction
The next stage of the project is conditional WS₂ defect generation.
The intended conditioning variables are:
Formation energy
Charge state
Defect type
The planned direction is to use classifier-free guidance (CFG) for controllable generation.
The unconditional model released here serves as the baseline against which the conditional model will be evaluated.
Reproducibility
The following parameters were kept consistent during the main held-out sampling experiment:
Checkpoint
500 diffusion steps
2.5 Å defect-region masking
Fixed WS₂ host templates
Z confinement
Random seeds were varied for independent generated samples.
Example Sampling
A typical AGeDi sampling command is:
agedi sample \
/path/to/AGeDi-WS2-Unconditional \
--n_samples 1 \
--n_atoms <NUMBER_OF_FREE_ATOMS> \
--template_path <FIXED_HOST_TEMPLATE.xyz> \
--confinement 4.87254786 12.67985539 \
--steps 500 \
--seed 42 \
--output outputs \
--name sample
<NUMBER_OF_FREE_ATOMS> depends on the defect-region mask associated with the input WS₂ template.
Example Workflow
The expected workflow is:
Fixed pristine WS₂ host
│
▼
Defect-region mask
│
▼
Fixed host + free atoms
│
▼
AGeDi diffusion model
│
▼
Generated WS₂ defect structure
The fixed host remains unchanged while the defect-region atoms are generated.
Model Card Summary
| Property | Value |
|---|---|
| Model | AGeDi-WS2-Unconditional |
| Task | Unconditional WS₂ defect generation |
| Dataset | 1,706 2D WS₂ defect structures |
| Training structures | 1,365 |
| Validation structures | 170 |
| Held-out test structures | 171 |
| Training epochs | 1,000 |
| Best validation loss | 0.028201 |
| Architecture | PaiNN |
| Atom basis | 64 |
| Interaction blocks | 4 |
| Cutoff | 6.0 Å |
| Sampling steps | 500 |
| Large benchmark | 1,000 generated structures |
Disclaimer
This checkpoint is released as a research artifact.
Generated structures should not be considered experimentally validated, thermodynamically stable, or suitable for direct materials screening without additional structural relaxation, energetic calculations, and domain-specific validation.
Citation
When citing this model, please cite the corresponding research work associated with this WS₂ defect-generation project once it is publicly available.
This repository contains the trained WS₂ unconditional checkpoint and associated configuration files.