empirical-machines/atomformer-p2-v2

AtomFormer P2 v2 (l, 59M)

Checkpoint of the AtomFormer trunk after phase P2 ("instinct": behavior cloning of recorded DFT-relaxation optimizer steps) of a search-native atomistic-ML program. Warm-started from the P1 verifier checkpoint and trained compound (E/F + verifier + policy).

Architecture

Bone-stock pre-norm transformer over atoms-as-tokens:

• l size: d_model 640, 12 layers (dense), 10 heads — 59.09M params
• Geometry channel: pairwise-distance attention bias B_ij = W·RBF(d_ij) (min-image under PBC) + radial-density fingerprint input + two-body pair-energy readout
• 8 global register tokens (readout site for value/meta heads)
• Conservative forces by autograd: F = −dE/dx (no direct force head)
• No equivariance constraints in weights — E invariant by construction (distances only), F exactly equivariant by differentiation, permutation by set attention
• Additive task/fidelity + charge + spin conditioning; frozen per-task composition references (refs_p4)

Heads:

• E energy (per-task normalized), F = −dE/dx forces
• P1 verifier V (from registers): E_end (basin endpoint), dE_togo, convergence logit
• P2 policy pi: per-atom displacement field (pairwise projection, non-antisymmetrized), trained by BC on recorded BFGS-on-DFT relaxation steps (dr to next frame)

Training

• Compound E/F + P1 + P2 fine-tune over 2.00B tokens (63,919 steps)
• Mix: MPtrj + T1x + OC20 (E_ads-referenced) + SPICE2 + OMat24 (300M cap); 1.11M action-labeled consecutive frames (mostly MPtrj) for policy BC
• 8×A100 DDP, lr 2e-4, AdamW, lam_f 2.0, lam_p2 10 (rel_l1, scale floor 0.03 Å, hygiene window 0.002–0.5 Å), fp32 evals

Final held-out eval

Domain	Task	E MAE (meV/atom)	F MAE (meV/Å)
MPtrj (PBE+U crystals)	1	30.4	91.1
OMat24	5	40.2	203.6
OC20 (E_ads)	2	4.5	64.9
T1x (wB97X)	3	4.2	57.1
SPICE2 (wB97M-D3BJ)	0	4.1	75.1

P1 verifier: convergence accuracy 0.823. P2 policy: action L1 0.0359 Å (predict-zero baseline 0.0363 Å; rel_l1 0.53).

Files

• ckpt_final.pt — PyTorch checkpoint: model (135 tensors, 59.09M params), cfg, tokens
• config.json — full training config
• results.jsonl — complete train/eval event log

Usage

import torch
ck = torch.load("ckpt_final.pt", map_location="cpu", weights_only=False)
state_dict = ck["model"]

Instantiating the model requires the AtomFormer code, which is not yet public.

Status

Research preview from an active experiment series (P0 perception → P1 judgment → P2 instinct, this checkpoint → P3 search-in-the-loop expert iteration → P4 scope).