File size: 6,467 Bytes
c2dfde3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c78e43f
 
 
 
 
 
 
 
 
 
 
 
 
 
c2dfde3
 
 
 
fff0d09
 
c2dfde3
 
 
 
 
 
 
 
fff0d09
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c2dfde3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
---
license: mit
tags: [crystal-structure-prediction, materials-science, transformer, coordination-chemistry]
library_name: pytorch
---

# TopG — a coordination-conditioned generator of periodic cation frameworks

**Status: research artifact, mid-project. Read the limitations before using any number.**

Given a *coordination query* — for each inequivalent metal site, its site symmetry, coordination
sequence and shell-distance ratios — the model emits a **complete crystal structure**: space group,
Wyckoff positions, free coordinates and free cell-shape parameters. Nothing downstream re-derives
any of it.

## Model

| | |
|---|---|
| architecture | causal transformer decoder, 4 layers, d=256, 8 heads, FFN 1024 |
| parameters | **4.5 M** |
| conditioning | one prefix token per site + one global token (so k is a sequence length, not a weight shape) |
| vocabulary | 2,317 tokens — per-space-group Wyckoff labels, coarse/fine coordinate and cell grids |
| training | 69,190 enumerated records, 100 epochs, ~100 min on a Quadro RTX 4000 |

## ⚠ CHECKPOINT / CODE COMPATIBILITY — READ BEFORE LOADING

The conditioning vector changed on 2026-08-08 (WP-018 added a per-shell **SPECIFIED** bit so partial
queries could be expressed). That is a **breaking representation change**:

| arms | conditioning | load with |
|---|---|---|
| `H_pq*` | `PER_SITE_DIM=68`, `COND_DIM=1020` | **current `main`** (commit `74dae8f` or later) |
| `G_100*` | `PER_SITE_DIM=62`, `COND_DIM=948` | **commit `203c036`** (the parent of the change) |

Loading a `G` checkpoint with current code raises
`size mismatch for site_proj.0.weight: [256, 62] vs [256, 68]`. That is the expected, loud failure —
not a corrupted download.

Two families, three seeds each:
- **`G_100*`** — full-depth queries (all six shells specified).
- **`H_pq*`** — partial queries: trained with the query depth truncated to 2–6 shells from a declared,
  seeded distribution, so a user can ask for less than six shells.
- **`I_full*`** — the control for the above: identical to `H_pq*` in every respect **except** the
  truncation. Included so the partial-query claim can be checked rather than taken on trust.

Two checkpoints per arm, because **the two halves of the task peak at different times**:
- `ckpt_00020.pt` — best validation loss, best **skeleton** accuracy.
- `ckpt_00100.pt` — best **geometry** (r3). Measured, paired, 3 seeds: r3 1.50%→2.50% (p=0.017) and
  1.02%→2.86% (p=0.0005) between these two points, while skeleton-exact falls.

**There is no single best checkpoint. Select per metric and say which you used.**

## ⭐ Partial queries: ask for less than six shells

A user rarely knows six coordination shells. They know the first one or two. The `H_pq*` arms are
trained with the query depth truncated to 2–6 shells, so a **partial** query is both expressible and
answerable — and the control shows the *training*, not just the encoding, is what makes it work.

Both families trained identically under the same conditioning; only the truncation differs
[3 seeds, 400 queries × 3, K=24]:

| query depth | applicable rung | `I_full` (full-query trained) | `H_pq` (partial-query trained) |
|---|---|---|---|
| **2 shells** | r0 | 19.4% · skeleton 10.2% | **61.5% · skeleton 49.6%** |
| **3 shells** | r2 (3sh + 4%) | 1.0% · skeleton 15.8% | **4.2% · skeleton 50.2%** |
| 6 shells | r3 (6sh + 4%) | 2.4% · skeleton 54.2% | 2.4% · skeleton 53.6% |

**At full depth the families are identical** — partial support is free. **At shallow depth a
full-query model collapses**: shown an input pattern it never saw (`specified=[1,1,0,0,0,0]`), its
skeleton accuracy falls to 10.2%, worse than anywhere on its own training curve.

⚠ **r0 is a weak criterion** (two shells, no ratio test) and **10.8% of 2-shell queries admit ≥2
distinct structures** — so the absolute 61.5% is against a weak bar. What carries weight is the
*paired contrast*: same bar, same queries, 19.4% vs 61.5%.

**Use `H_pq*` if you want to send partial queries. Use either family at full depth.**

## What it does well

| capability | measured [per-QUERY unless stated] |
|---|---|
| emits a valid parseable structure | **99%** of beams; every death diagnosed as a model choice |
| predicts N without being told | **92% per candidate**, 98.75% per query |
| predicts k without being told | **97% per candidate** |
| names the exact skeleton (sg + Wyckoff labels + N) | **58.8%** |
| satisfies a **2-shell** query (r0) | **61.5%** |

Skeleton accuracy is roughly flat from N=2 to N=192 and k=2 to k=11.

## What it does not do

**Metric precision.** r3 (six shells + 4% distance tolerance — the project's historical criterion) is
**2.5–3.1%**. Loosening the tolerance to 15% reaches only ~5.5%: the distances are wrong, not nearly
right. Structure:

- **Below ndim ≈ 9 the CELL is the binding constraint.** In cubic, where symmetry fixes the cell
  entirely, r3 is **18.75% vs 0.25%** — matched on N and k, and on ndim.
- **Above ndim ≈ 9 the COORDINATES bind.** Cubic scores zero there too.
- Every r3 success in the evaluation sits at **ndim ≤ 9 and k ≤ 4**.

## Honest instrument facts

- **Grid ceiling 95.67%** — a *perfect* model scores 95.67%, not 100%, on this evaluation set at this
  emission grid. Ground truth round-tripped through encode→decode→rebuild.
- **Model-free control: 0/400.** Oracle skeleton + the training split's median geometry scores zero.
- **Train-seed noise floor: 2–4 pp** on skeleton, **0.5 pp** on r3, from 3 cold starts differing only
  in the seed. **No discrete claim below ~9 pp is interpretable from a single run at a long epoch
  budget.**

## Limitations

- Trained on **enumerated (synthetic) structures only**. Materials Project data is held out for
  evaluation and never trained on (enforced in code).
- Space groups 3–230; **triclinic is thin** (348 corpus records) and cubic external validation is
  effectively absent (7 MP queries survive rebuild).
- Queries are **scale-free** — the model emits cell *shape*, never absolute length.
- No refinement step exists yet, so the "model output as a starting point" question is open.

## Provenance

Code, every preregistration, every result and every correction:
**https://github.com/rgpalgrave/TopG**. `SHA256SUMS.json` here matches the checkpoints; each
`arm_meta.json` records the corpus hash, split hashes, seed, vocab hash, tokenizer version and git
commit.