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/.
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)
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).
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):
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").
