| --- |
| 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** |
| |
