README.md

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