cosmufr-v09

CosmUFR v0.9 — Run 4

124M-parameter cosmological emulator using Unified Field Representation (UFR) belief-state dynamics. Trained on 84M simulations from 9 cosmological suites.

No attention, no transformers — pure energy-minimizing belief settling across 16 gradient descent steps.

What it does

Given log P(k) at z=0 and z=0.47 as input (400d total), CosmUFR infers all 8 cosmological parameters simultaneously with calibrated per-parameter uncertainties.

Performance (Run 4 · B200 · B=4096 · best.pt = Ep30)

Val metrics from 422K samples (live training monitor). Final 50K stratified test eval pending.

Architecture

• 124.6M parameters total
• d_b = 1024 structured belief state [s:512 | c:256 | u:128 | p:128]
• ObsEncoder (22M): log10 P(k) [400d] -> z [1024d], 4-block 1024-wide MLP
• BeliefProposal (21M): joint path (L_param+L_gen+L_nll) — DECOUPLED from seq path
• BeliefProposalSeq (21M): sequential path (L_seq only) — separate module, Run 4 fix
• SettlingCore (3M): 16-step energy gradient descent, learned preconditioner
• AttractorBank (4.2M): 4096 EMA-updated prototype belief vectors
• GenerativeHead (69M): k-continuous implicit field, 8-block 2048-wide MLP

Training

• Dataset: 84,105,515 samples · 9 simulation suites
• Run 4: B200 192GB VRAM · PyTorch 2.7 · CUDA 12.8 · BF16 · B=4096 · 40 epochs
• Best checkpoint: Ep30, composite score 0.6037
• Duration: ~15 hours
• GCS: gs://cosmufr-training/checkpoints/v10/b200-b4096-run4-20260414_031247/

Key fixes over v0.8

1. BeliefProposal decoupled: belief_proposal (joint) and belief_proposal_seq (seq) are separate modules — eliminates gradient interference that caused h oscillation
2. SettlingCore isolated: L_seq gradient cannot enter SettlingCore via b_prev or b_hat
3. Crash recovery: Optimizer state saved/restored on resume — no cold-restart Adam instability

Usage

import torch
from huggingface_hub import hf_hub_download
from model.cosmufr_arch import CosmUFRLite
from model.cosmufr_config import CosmUFRConfig

cfg   = CosmUFRConfig()
model = CosmUFRLite(cfg)
ckpt  = hf_hub_download("arajgor1/cosmufr-v09", "best.pt")
state = torch.load(ckpt, map_location="cpu", weights_only=False)
model.load_state_dict(state["model"])
model.eval()

# obs: [B, 400] — concat log10 P(k) at z=0 and z=0.47 (200 k-modes each)
obs    = torch.randn(1, 400)
b_prev = torch.zeros(1, cfg.d_b)
out    = model(obs, b_prev)
print("Params:", dict(zip(['Om','s8','h','ns','Ob','w0','mv','wa'],
                          out['params'][0].tolist())))
print("1-sigma:", out['variances'][0].sqrt().tolist())

Citation

@misc{cosmufr2026,
  title  = {CosmUFR v0.9: Belief-State Cosmological Emulator},
  author = {Rajgor, Aaditya},
  year   = {2026},
  url    = {https://huggingface.co/arajgor1/cosmufr-v09}
}