--- license: mit library_name: pytorch pipeline_tag: other tags: - perturbation-biology - single-cell - t-cells - crispr - adaptive-computation - halting - state-model --- # RefineRx — adaptive-depth halting as a per-perturbation signature RefineRx adds a **learned, per-perturbation adaptive-depth halting** mechanism to the ARC **STATE** perturbation-response architecture, and asks whether the resulting halting depth E[N] is a reproducible, effect-size-independent signature. This repository ships both the **positive** and the **negative** result of that question: - **✅ Positive (the working artifacts): four `oracle-ACT` halt heads** on **frozen** cell-line STATE-embedding (ST-SE) backbones (K562, HepG2, Jurkat, RPE1). Read off a fixed, gradually-refining backbone, these recover a usable per-perturbation depth signature (reproducible within-line, effect-independent, non-redundant with network topology). - **⚠️ Negative (documented, for the record): CD4-native fused-halting checkpoints** trained from scratch on CD4+ T-cell **pseudobulk**. The response head fits, but the halting depth **collapses to a constant** on this pseudobulk substrate — see the key result below. These checkpoints are shipped as the reproducible negative, not as a recommended depth model. > **Naming note.** "CD4-native" refers only to the from-scratch pseudobulk > checkpoints in `checkpoints/`; it is the *collapsed* model. The usable > depth signal in this repo comes from the frozen-backbone halt heads in > `halt_heads/`. **Paper:** *When Does a Perturbation Model Know Enough?* (Yash Raj). **Code:** https://github.com/yashraj59/RefineRx --- ## ⚠️ Key result — read this first **On the CD4 pseudobulk backbone, the fused-halting expected depth E[N] collapses to a constant** (E[N] → 6.0, the round budget; across-perturbation std ≈ 1e-4), **even though the model fits the perturbation response well.** The response head works; the *depth* signal does not survive on this substrate. Why it collapses — a head-free diagnostic locates the cause **upstream of the halt head**: - The oracle stopping round **r\* is degenerate**: r\* = max (the final refinement round) for **100% of held-out perturbations**, with zero spread. - The per-round magnitude-free distances form a **step function** (1.04 → 1.23 → 1.20 → 1.24 → 0.99 → 0.03): refinement rounds 2–5 are no better than round 1, and only the *final* round converges. There is no smooth accuracy-vs-depth curve for a halt head to calibrate on. - Increasing the ponder weight 10× (γ: 0.01 → 0.1) does **not** restore spread. ### What we do — and do not — claim We attribute the collapse **primarily to pseudobulk aggregation, and we are careful not to over-read it:** - **The signal exists at single-cell resolution.** The *frozen*-backbone single-cell CD4 analysis (two donors, D1/D4) on this same biology recovers a **reproducible within-donor depth signature** (split-half ρ = 0.62–0.75, rising with stimulation), and at the stimulated 48 h endpoint separates resting-sparing inflammatory-program suppressors from damaging/generic perturbations. So depth **is** recoverable from CD4 data — just not from aggregated pseudobulk. - **Pseudobulk is the leading suspect,** because averaging cells removes the within-population variation the single-cell signature draws on (and raw CD4 pseudobulk response is donor-unstable by construction). - **We could not run the clean control.** Separating "pseudobulk hides it" from "the from-scratch x0-recurrence mechanism kills it" requires training a CD4-native backbone from scratch **at single-cell resolution**. The full CD4+ single-cell screen is ~22M cells; our frozen analysis used a two-donor subsample capped at 50 cells/perturbation (3.9M cells). From-scratch single-cell training did not reach stable dataloader throughput in the compute available, so this control is left as future work. **The honest, narrow statement:** *on aggregated pseudobulk the from-scratch fused model fits the response but yields no usable depth signature, most plausibly because pseudobulk hides it.* This does not, on its own, refute end-to-end halting. **Interpretation.** Halting is a **computational proxy for response complexity, not biological time.** Depth is a property of a fitted model and its refinement axis — it should not be read as a biological trajectory, kinetic ordering, or a claim about how long a real cell "computes" a response. --- ## Contents ``` checkpoints/ best.ckpt # CD4-native STATE transition + fused halting (~449 MB) last.ckpt # last-step checkpoint of the same run (~449 MB) halt_heads/ halthead_k562.pt + halthead_k562_meta.json # oracle-ACT head, frozen K562 ST-SE backbone halthead_hepg2.pt + halthead_hepg2_meta.json # oracle-ACT head, frozen HepG2 ST-SE backbone halthead_jurkat.pt + halthead_jurkat_meta.json # oracle-ACT head, frozen Jurkat ST-SE backbone halthead_rpe1.pt + halthead_rpe1_meta.json # oracle-ACT head, frozen RPE1 ST-SE backbone config/ config.yaml # full run config for the CD4-native checkpoints hparams.yaml # Lightning hparams snapshot (version_0) state_halt.yaml # state_halt model config template (repo default) ``` --- ## 1. CD4-native transition + fused-halting checkpoints `checkpoints/best.ckpt` and `checkpoints/last.ckpt`. **Architecture** — STATE transition model (`state_halt`) with a fused adaptive-depth halting mechanism combining: - **oracle-r\*** supervision (project §4.4), and - **joint calibration** of the halt head (project §4.5), using a **magnitude-free per-round target**, with **ponder loss gated on after warmup** and **KL = 0**. Backbone / halting hyperparameters (from `config/config.yaml`): | field | value | |---|---| | model | `state_halt` | | transformer backbone | `llama`, bidirectional, 8 layers, 12 heads, head_dim 28 | | hidden_dim | 336 | | intermediate_size | 3072 | | cell_set_len | 64 | | n_refine_rounds | 6 | | predict_residual / softplus | true / true | | distributional loss | energy (blur 0.05) | | confidence_token | true (weight 0.01, target_scale 10.0) | | halt_tau / halt_alpha / halt_beta / halt_gamma | 0.05 / 0.5 / 0.1 / 0.1 | | halt_warmup_steps | 2000 | | halt_magnitude_free | true | | optimizer | AdamW, lr 1e-3, weight_decay 5e-4, grad_clip 10 | | batch_size / max_steps | 64 / 15000 | | train_seed | 42 | **Training data** — CD4+ T-cell **pseudobulk**: - **278,684 pseudobulk profiles × 2,001 HVG** (HVG space shared with the frozen-backbone CD4 analyses; includes PuroR). - Conditions: **Rest / Stim8hr / Stim48hr**. - **4 donors** (used as the batch covariate). - Input representation: `X_hvg`, `log1p(counts)` only (no library-size / 1e4 rescaling), following the cell-load contract. - Data schema: `pert_col=gene`, `cell_type_key=condition`, `batch_col=donor`, `control_pert=NTC`, `output_space=gene`. The model **fits the perturbation response**; see the key result above for what fails (depth identifiability on pseudobulk), not what works (response prediction). --- ## 2. Cell-line oracle-ACT halt heads `halt_heads/halthead_{k562,hepg2,jurkat,rpe1}.pt` (+ `_meta.json`). Each is an **oracle-ACT halt head** trained on a **frozen ST-SE cell-line llama backbone** (no gradient to the backbone). They reproduce the expected-depth E[N] signal at reproducibility ρ = 1.0 on their own line. **Head module** (`AdaptiveStateRefine`, from the `_meta.json`): - Frozen backbone, **8 layers → 8 exit rounds** (exit_layers 1..8), hidden H = 328. - **Halt head:** `Sequential(LayerNorm(H), Linear(H,64), SiLU, Linear(64,1))`, last-bias init −2.0 (sequential-hazard formulation). - **Error head:** `Sequential(LayerNorm(H), Linear(H,64), SiLU, Linear(64,1))` + softplus. - **Refinement token:** `Parameter(H)`, init `randn*0.02`, appended at sequence index S. - Feature dims: pert_dim 2024, input_dim 2058, n_genes 2000. **Training** (per `_meta.json`): 50 epochs (15 warmup), 4 seeds, cell_set_len S = 64, AdamW lr 3e-3, grad_clip 1.0, τ 0.05 / α 0.5 / β 1.0 / γ 0.1 / δ 0.1, min_cells 20, **ponder gated on after warmup**, **KL = 0**. Per-line perturbation counts are in each meta file (e.g. K562 = 968). **These are the positive-control counterpart to the CD4 pseudobulk collapse.** Read off a *frozen, gradually-refining* backbone, the halt heads DO admit a usable per-perturbation depth signal: within-line split-half ρ ≈ 0.76–0.85, effect-size independent once #DE and cell count are controlled (partial ρ ≈ 0), and **non-redundant with network topology** — no GRN/PPI graph statistic reproduces the per-perturbation ordering of E[N] (best |ρ| = 0.23, below a 0.3 novelty ceiling). The signature is reproducible *within* a cell type but does **not** port across cell lines (cross-line ρ = 0.14) — it is a cell-type-specific property. Two claims must be held apart. As a **descriptor**, E[N] is *non-redundant* with the network (above): it captures per-perturbation structure the graph statistics miss. For one downstream **target-class ranking task**, however, it is *non-additive* with a STRING baseline built from the same functional modules — adding |ΔE[N]| gives no AUC lift over response + STRING in three of four lines. Non-redundancy as a descriptor and non-additivity for that one classifier are different statements; only the latter is negative, and it does not diminish the signature's novelty. --- ## Provenance & caveats - Checkpoints were **stage-copied from a live training job** (consistent snapshot; the training run was not interrupted). - The three config files describe the CD4 run; `state_halt.yaml` is the repo default template (hidden_dim 768 there vs 336 in the actual CD4 run — use `config.yaml` for the shipped checkpoints). - **Depth is model-specific and is not biological time.** Do not interpret E[N] or r\* as a kinetic / temporal ordering of the biological response. - The collapse on pseudobulk is a substrate result, not a proof that CD4 halting is intrinsically unidentifiable (see "What we do — and do not — claim"). ## Citation / code Source, training scripts, and analysis: **https://github.com/yashraj59/RefineRx**