fix(chrombpnet): BPNet CHIP — joint strand softmax, exp(C)-n_tracks, bias shapes

[lucapinello]

Item C. You asked me to think harder before spending the 4 GPU-hours — that was the right call. The joint-softmax convention you picked is confirmed, but I found two further defects, one of them larger, that would have been baked into the rebuild. All three had to be settled together because each changes what the 744 CHIP CDF rows must be built from.

The convention question is now answered from the vendored training code, not by analogy: chorus/oracles/chrombpnet_source/templates/BPNet/ is upstream kundajelab/bpnet-refactor (losses.py byte-identical; arch.py/custommodel.py differ only by import shims).

1. Joint strand softmax — confirmed, and now proven

custommodel.py:93-104 (the orig_multi_loss=False branch) reshapes the (B,1000,2) profile to (B,2000) and applies one multinomial_nll, which normalizes over the last axis. The trained parameterization is a single softmax spanning both strands.

Which branch trained these weights is provable, not assumed: orig_multi_loss=False yields a (2,1) logcounts kernel; True yields (3,2) and refuses to load the real h5. All 744/744 cached models have (2,1).

So per-strand isn't merely disfavoured — with a single count head there is no per-strand total to scale by, making it inexpressible.

2. The count inverse is exp(C) − n_tracks, not expm1(C) ← the one I nearly baked in

bpnet-refactor's generator builds the count target per track as log1p: np.log(np.sum(profile_predictions, axis=1) + 1) over an array shaped (batch, positions, tracks) — so the sum is over positions — and custommodel.py:57 pools a task's tracks with reduce_logsumexp. The trained target is therefore

C = log( Σ_t (1 + c_t) ) = log(n_tracks + total)   ⟹   total = exp(C) − n_tracks

n_tracks	true total	expm1(C)	exp(C) − n
1 (ATAC/DNASE)	100	100.0000 ✓	100.0000 ✓
2 (CHIP)	100	101.0000 ✗	100.0000 ✓

For the 42 single-track models this reduces to expm1 exactly — verified bit-identical, max diff 0.000e+00. So #113 stands and those 42 CDF rows stay valid. For the 744 CHIP models expm1 left exactly one read of inflation: measured 1.12×–1.95× on real loci, worse at quiet sites (median 1.78× on background 501 bp window sums).

Upstream's own bpnet/cli/predict.py uses a bare exp(C), off by +n_tracks the other way. The target construction is the authority, not that CLI.

3. The count-bias input shape was hardcoded (N, 1) for a (None, 2) input

BPNet logsumexp-reduces that bias before the final Dense (main_logsumexp_counts_bias_0). The wrong width — silently broadcast by Keras, not rejected — made the term log(1)=0 instead of log(2)=0.6931, shifting every predicted log-count down by a constant w·log 2 = 0.588506 (w = 0.849035, the Dense weight — matching the measured shift to all printed digits).

log-count	chorus counts	correct/chorus
0.50	0.649	3.04×
1.88	5.524	1.95×
8.00	2979.96	1.80×

Profile logits were unaffected. The builder derived the shape correctly and was not affected, so the oracle silently disagreed with its own CDFs. Fixed at all four call sites via _zero_bias_inputs (shapes derived from model.inputs): _predict_direct, predict_sliding, the env-path predict_template.py, and the public one-shot helper oracles/bpnet.py — whose docstring also prescribed the superseded per-strand exp(counts) recipe.

Also fixed

predict_sliding used SUM-LOGITS for CHIP while _predict used per-strand, reporting it under a ':+' id. It now takes the joint softmax and emits the plus strand, matching the id it reports and the pooled per-strand CDF row it's looked up against. A both-strand total there would be 1.53–2.93× the _predict value behind the same id.

Verification

• builder ↔ oracle agree to 2.0e-07 on every strand/locus combination
• the two strands conserve exp(C) − 2 exactly
• DNASE:K562 bit-identical to before (0.000e+00)
• 23 tests in tests/test_chrombpnet_counts.py, including a helper that builds the count target the way training does so the expectations can't drift from the loss, a bit-identity test for the single-track inverse, and a check that expm1 overstates a two-track model by exactly one read
• Fast suite: 491 passed, 4 skipped, 0 errors

Rebuild — not done yet, deliberately

The 744 CHIP rows must be rebuilt; the 42 ATAC/DNASE rows must not (unchanged, and merge_shards appends and de-dups by track_id, so the base must contain only those 42 at merge time). Expected CHIP counts double, since both strands now feed the same row: effect 37,344 / summary 68,008 / perbin 2,176,256.

Two caveats worth your judgement before I start it:

• Zero count-bias may be out of distribution. (B,2) zeros is the right shape, but bpnet-refactor trains with real control bigwigs unless set_bias_as_zero was used, and which applies to the ENCODE-published BPNet models can't be determined from the weights. The 0.5885 uplift correctly encodes "no control" — it isn't necessarily the regime these models were trained in. Fine for percentiles (self-consistent), worth knowing before quoting absolute CHIP counts as calibrated.
• _MIN_MAGNITUDE['tf_binding'] = 0.1 (variant_report.py:594) becomes more load-bearing as CHIP values shrink, and a CDF rebuild doesn't address it.

Say go and I'll run the CHIP-only rebuild and open the HF dataset PR with the same discipline as PR #2.

🤖 Generated with Claude Code