RefineRx

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):

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