File size: 35,171 Bytes
b415c47
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
# HEART β€” Efficient Attention with Rank-factorized bias Transformer

**TL;DR:** window-attention SR transformer. HAT-iLN quality, FlashAttention
speed, none of the complexity. **One architecture, one config. Trains in bf16.
Exports to dynamic ONNX with a single command** β€” no fused/unfused checkpoint
pairs, no custom CUDA kernels, no relative-position tables.

## What is HEART?

HEART is what HAT-iLN becomes when you strip out the parts nobody actually
deploys:

- **RIB** replaces the relative-position-bias *table* β€” learned low-rank
  position features concatenated onto Q/K, so attention runs on standard
  FlashAttention/SDPA kernels. No table, no index gather, no mask.
- **OCAB removed** β€” HAT's most complex, most memory-hungry block. Our
  ablations showed it isn't needed at window 32.
- **Keeps the proven bits:** i-LN (input-adaptive normalization, bf16-stable),
  window attention, a convolutional branch (CAB).

The result: a ~16.7M single-image SR architecture that **ties HAT on PSNR,
trains in stable bf16 AMP, and exports to ONNX that runs in TensorRT / ONNX
Runtime / DirectML.**

## Why HEART exists

Most SR papers chase benchmark numbers by adding modules that are annoying to
train and painful to deploy β€” dynamic routing, deformable attention, blocks
that only work with a specific resolution, or a "fused" checkpoint for
inference and an "unfused" one for fine-tuning. The result looks good on a
leaderboard and gathers dust in practice.

HEART has a different goal: **real-world simplicity.** Train it anywhere,
convert it to dynamic-shape ONNX, run it everywhere. No fusion step. No second
model file. No custom kernels.

## When to use HEART

- You want HAT-class quality on real images (denoise, deblur, upscale) without
  babysitting the training.
- You need weights you can convert to ONNX and ship β€” no fusion step, no "use
  variant B for inference."
- You are tired of benchmark-chasing networks that add modules no one can
  actually deploy.
- You need desktop quality at reasonable inference cost. For mobile/edge, use
  NERVE (coming in the BODY suite).

## Quickstart

1. Copy `traiNNer/archs/heart_arch.py` into your traiNNer-redux clone
   (`traiNNer/archs/`).
2. Copy the release config from `options/train/RELEASE/4x_HEART_release.yml`
   and point `dataroot_gt` / `dataroot_lq` at your data.
3. Train:
   ```bash
   python train.py -opt options/train/RELEASE/4x_HEART_release.yml --auto_resume
   ```
   (`--auto_resume` is **mandatory** β€” without it an interrupted run restarts
   from iter 0 and archives the experiment dir.)
4. Convert to dynamic ONNX (each model, one at a time):
   ```bash
   ulimit -v 14000000
   ./venv/bin/python scripts/heart/export_minimal.py \
       experiments/4x_HEART_pretrain600k/models/net_g_ema_615000.safetensors \
       onnx/4x_HEART_1x3xHxW_fp32_op17.onnx 4
   ```
5. Run inference anywhere ONNX Runtime or TensorRT runs.

Full architecture docs, ablation results, and training recipes are below.

> **Training crop size: use `lq_size: 96` (a multiple of the 32 px window).**
> HEART partitions the image into 32Γ—32 windows. Training at a crop size that
> is *not* a multiple of 32 (e.g. 80) makes the model reflect-pad its own
> bottom/right edge every batch, so it learns a false "edge = mirrored
> content" response β€” visible as a bright bottom band and a tile-grid pattern,
> and strongly amplified by a GAN loss. All HEART ablations and release runs
> used `lq_size: 96` (= 3Γ—32); a larger multiple (128) also works. This
> matters most for GAN/OTF finetunes, which are most sensitive to boundary
> statistics.

## BODY suite status

| Model | Role | Params | Status |
|---|---|---|---|
| HEART | desktop quality tier | 16.7M | trained + ONNX |
| NERVE | mobile/edge tier | ~270K | trained |

## Design

```
conv_first β†’ [6Γ— RHAG-RIB] β†’ conv_after_body (+ global residual)
           β†’ pixelshuffle upsampler

RHAG-RIB (residual group):
  [6Γ— HAB-RIB blocks] β†’ 3x3 conv β†’ group residual

HAB-RIB block:
  i-LN β†’ RIB window attention (window 32, FlashAttention/SDPA)
       + CAB conv-attention branch (scale 0.01)
  i-LN β†’ MLP (mlp_ratio 2)
```

- **i-LN** (Image Restoration tailored LayerNorm, arXiv:2504.06629) β€” stable
  training in bf16/AMP, no divergence.
- **RIB window attention** β€” FlashAttention-compatible, window 32, shifted
  windows via non-wrapping pad-and-partition (no masks, no position tables).
- **CAB conv branch** β€” HAT's proven local inductive bias, kept.
- **OCAB removed** β€” the most complex and memory-hungry part of HAT is gone.

## Goals

| Goal | How |
|---|---|
| Great results | HAT's proven hybrid block (window + conv) + large-window global context |
| Fast train & inference | FlashAttention via RIB (SST measured 2.1x train / 2.9x inference vs RPB) |
| Stable training | i-LN + EMA std clamping, bf16/AMP safe |
| Simple code | ~640 lines, reuses existing blocks, no masks/tables/OCAB |
| Easy deployment | Pure PyTorch, ONNX-exportable, no custom CUDA kernels |

Note on speed: the FlashAttention win only exists in bf16/fp16 (measured
~1.6x faster than HAT-iLN-M at bf16 inference). In fp32, HEART is slightly
slower than HAT-iLN β€” fp32 has no flash kernel. Use reduced precision.

## Status

Implemented, verified, and releasing. Architecture is frozen; all weights and
ONNX files below are measured, not extrapolated.

Final architecture (post-ablation, 2026-08-31): RIB window-32 attention + CAB,
`attention_freq=2`, rank 8, linear MLP β€” the 20k sweep winner. RoPE,
no-positional, GDFN, SwiGLU, LayerScale, window 48/64, and rank 18/32 were all
measured and rejected; the losing code paths were removed from `heart_arch.py`
(646 lines, 16/16 tests, lint + typecheck clean).

Release artifacts (all exist, paths verified):

| Artifact | Checkpoint | ONNX |
|---|---|---|
| 4x bicubic | `experiments/4x_HEART_pretrain600k/models/net_g_ema_615000.safetensors` | `onnx/4x_HEART_1x3xHxW_fp32_op17.onnx` |
| 2x | `experiments/2x_HEART/models/net_g_ema_200000.safetensors` | `onnx/2x_HEART_1x3xHxW_fp32_op17.onnx` |
| 4x OTF fidelity (v1.1) | `experiments/backup_otf_v2/net_g_ema_150000.safetensors` | `onnx/4x_HEART_OTF_1x3xHxW_fp32_op17.onnx` |
| 4x OTF perceptual (GAN) | training | β€” |

All checkpoints backed up to a private HF repo
(`Phips/heart-sisr-pretrain`). ONNX files pass `onnx.checker` and match
PyTorch numerically (ORT diff < 1e-2).

In flight: **4x OTF perceptual (GAN) finetune** (~3 days remaining). When it
finishes, its weights + ONNX join the release set.

Known limitation: a faint **window-32 grid artifact** (128 px period in 4x
output) on smooth gradients and fine lines in OTF outputs. Diagnosed via FFT;
root cause is window boundaries (not JPEG blocks or PixelShuffle). Overlap
attention contradicts the frozen arch, so it is documented here, not fixed in
1.0.

Release-gate checks passed: `--auto_resume` (CLI flag) resume verified,
tiled inference seam-free, ONNX legacy-tracer opset 17 + onnxruntime,
fp16/bf16 inference, bf16 AMP training stability, torch.compile validated.
  the bare model, so checkpoints and val are unaffected. Kernel profile of
  one step: ~23% flash attention, ~30% small elementwise ops (what compile
  fuses), ~12% convs β€” no architectural hotspot left for custom kernels.

Dev notes for anyone touching the attention code:

- Shifted windows are non-wrapping: the feature map is padded top/left by half
  a window (reflect) and windows are partitioned from the padded origin. Border
  windows see reflected padding, never content from the opposite image border.
  Do NOT reintroduce `torch.roll` cyclic shifting β€” without an attention mask
  it couples opposite borders, and adding a mask would force SDPA off the flash
  kernel (the whole point of RIB).
- Logits are pre-scaled per SST eq. 5 (`q/√D`, `pos/√R`) and SDPA is called
  with `scale=1.0`. Re-adding default SDPA scaling silently divides logits
  by √48 β‰ˆ 7 and shifts attention toward uniform β€” shape-only tests will
  not catch it. Without this fix the model trains around it and loses
  quality for no reason.
- RIB features are computed in fp32 then cast to the input dtype. Skipping
  the cast breaks pure fp16 inference (buffer/param dtype mismatch).

Next step is a short ablation against the baselines before the full training
schedule:

- `heart` (this network, 19.2M params, RIB window 32)
- `atd` (20.3M, the efficiency champion)
- `hat_m` (20.8M, classic HAT)

The ablation trains each for ~20-50k iterations on DIV2K and compares
Urban100 PSNR and wall-clock speed. HEART proceeds to full DF2K training only
if it wins.

## Ablation results (2026-08-15, post shifted-window fix)

30k iterations each, identical settings (post-fix HEART code). RTX 3060 12 GB,
DIV2K β†’ Urban100:

| arch  | best Urban100 PSNR | @iter | final SSIM | train it/s | peak VRAM |
|-------|--------------------|-------|------------|------------|-----------|
| HEART | 25.0235            | 30000 | 0.7497     | 0.86       | 1.44 GB   |
| HAT_M | 25.0524            | 30000 | 0.7517     | 0.95       | 2.84 GB   |
| ATD   | OOM at iter 1 β€” did not run on the same config |

Reading the table honestly:

- **Quality:** HAT_M leads by +0.03 dB / +0.002 SSIM at the 30k cutoff β€”
  within noise for a short schedule, and both were still improving on their
  final validation. Effectively a tie; only a full-length run separates them.
- **VRAM:** HEART uses ~half the peak memory of HAT_M (1.44 vs 2.84 GB).
- **Train speed:** HAT_M trains ~10% faster (0.95 vs 0.86 it/s). Window 32
  attention is 4x the per-pixel attention work of window 16; the flash kernel
  makes it feasible, not free. HEART's speed advantage is at *inference*
  (~1.6x vs HAT-iLN-M bf16), not training throughput.
- **ATD OOM caveat:** HEART/HAT_M ran with `use_checkpoint: true`; ATD's
  factory has no such option. The OOM is a real deployment-relevant result
  (ATD does not train on a 12 GB card at lq_size 64 / batch 4 as configured),
  but the memory comparison is not fully symmetric because of that.

### Inference benchmark (same trained EMA checkpoints, RTX 3060)

Median of 5 runs after warmup, `cudnn.benchmark=False`,
`/tmp/kilo/inference_bench.py`:

| arch  | dtype | input    | time    | peak VRAM |
|-------|-------|----------|---------|-----------|
| HEART | bf16  | 320x180  | 0.884s  | 0.45 GB   |
| HAT_M | bf16  | 320x180  | 1.092s  | 1.20 GB   |
| HEART | bf16  | 480x270  | 2.024s  | 0.90 GB   |
| HAT_M | bf16  | 480x270  | 2.351s  | 2.47 GB   |
| HEART | fp32  | 320x180  | 2.350s  | 0.89 GB   |
| HAT_M | fp32  | 320x180  | 1.816s  | 2.37 GB   |
| HEART | fp32  | 480x270  | 5.347s  | 1.78 GB   |
| HAT_M | fp32  | 480x270  | 3.973s  | 4.94 GB   |

In bf16 (the deployment-relevant precision) HEART wins both speed (1.16-1.24x)
and VRAM (~2.7x less). In fp32 HAT_M is faster (no flash kernel for fp32), but
HEART still uses 2.7-2.8x less VRAM.

### Training speed benchmarks (RTX 3060, bf16 AMP, channels-last)

torch.compile (`use_compile: true`) validated in the training harness:
~1.5x training speed, quality-neutral (300-iter smoke: 21.88 dB vs 21.80 dB
uncompiled).

| config | it/s | 800k iters | peak VRAM |
|---|---|---|---|
| lq64 bs4, no compile (ablation) | 0.86 | ~10.7 days | 1.44 GB |
| lq64 bs4 + compile | 1.26 | ~7.4 days | 1.44 GB |
| lq96 bs2 + compile | 1.19 | ~7.8 days | 0.98 GB |
| lq96 bs4 + compile | ~0.63 | ~14.7 days | 1.70 GB |
| lq96 bs4 / lq128, no checkpoint | OOM | β€” | β€” |

compile_mode comparison (@ lq96 bs2, 3060): `reduce-overhead` 1.23 it/s >
`default` 1.18 > `dynamic=False` 1.19 β‰ˆ `default+TF32` 1.19 > no compile 0.77.
`max-autotune` is not worth it on small GPUs (inductor warns: not enough SMs).
At the ablation config (lq64 bs4), `max-autotune-no-cudagraphs` was measured
equal to `reduce-overhead` (both 1.28 it/s) with 2x the compile time and
slightly more VRAM β€” it gains nothing on a 28-SM card because the tuned GEMM
search falls back. `fast_matmul` (TF32) is a no-op under bf16 AMP β€” measured.

Kernel profile of one training step (uncompiled): flash attention ~23%,
small elementwise ops ~30% (what compile fuses β€” the source of the 1.5x),
convs ~12%, MLP gemms ~4%. No hotspot remains that a custom Triton kernel
could profitably target.

### Scoreboard vs HAT_M (30k ablation)

| Metric | Winner |
|---|---|
| Quality | tie (HAT_M +0.03 dB, within noise) |
| Inference speed bf16 | HEART (1.16-1.24x) |
| Inference VRAM | HEART (~2.7x less) |
| Training VRAM | HEART (1.44 vs 2.84 GB) |
| Training speed (no compile) | HAT_M (+10%) |
| ATD on 12 GB | does not run (OOM) |

### Rank ablation (done, 2026-08-17)

RIB's per-head positional-bias matrix has rank <= `rank`. Three 20k-iteration
runs compared rank 8 / 18 / 32 on the identical ablation config **with
torch.compile reduce-overhead** (`HEART_rank{8,18,32}_ablation.yml`, lr decay
at 16k). All three used compile so the comparison is not confounded (compile
shifts short-run PSNR ~0.08 dB, larger than the 0.03 dB decision threshold).

| rank | 20k PSNR |
|---|---|
| **8** | **24.7820** |
| 32 | 24.7586 |
| 18 | 24.7346 |

**Winner: rank 8** (smallest within 0.03 dB of the best). Rank 8 beating
rank 18 means positional capacity is *not* the bottleneck at rank 8 β€” a
lower-rank positional mechanism is enough, and it is leaner (concat dim 40
vs 48) and slightly faster. The canonical default in `heart_arch.py` is now
`rank=8`. The rank-18 uncompiled 30k reference (25.0235) was historical only.
The release configs are patched to rank 8.

### Architecture ablation: window size & attention frequency (2026-08-19)

Ran because an 8-day release was at stake: if window 48/64 or sparser
attention changed the quality-per-second picture, it had to be found before
committing. Setup: lq96 bs2 (the release crop), 20k iters, compile
reduce-overhead, rank 8, DIV2K, seed 1024 β€” identical settings except the
variable (`HEART_{w32,w48,w64,attn2}_ablation.yml`). The release run was
halted at 75k iters (checkpointed, resumable) while this ran.

| config | 20k PSNR | it/s (compiled) | params |
|---|---|---|---|
| **attn2** (`attention_freq=2`, RIB) | **24.6282** | **2.02** | 16.68M |
| attn3 (`attention_freq=3`) | 24.5606 | ~2.3 | 15.87M |
| w32 (baseline, RIB rank 8) | 24.5063 | 1.29 | 19.10M |
| none (no positional, freq=1) | 24.4965 | ~1.4 | 18.94M |
| w48 | 24.4347 | 0.80 | 19.10M |
| rope (RoPE, freq=1) | 24.3615 | ~1.4 | 18.94M |
| gdfn (Gated-Dconv FFN, ratio 2.0, freq=2) | 24.2546 | ~1.9 | 19.28M |
| swiglu (SwiGLU FFN, ratio 4/3 param-matched, freq=2) | 23.7403 (best @12k; collapsed to 23.21 @20k) | ~1.9 | ~16.7M |
| qknorm (per-head QK LayerNorm, freq=2) | 24.6334 (tie with attn2, +0.005) | ~2.0 | ~16.7M |
| swiglu + qknorm | 23.9576 | ~1.9 | ~16.7M |
| layerscale (per-block residual scaling, init 1e-6) | 24.0288 | ~2.0 | ~16.7M |
| w64 | skipped (w48 already lost; 4.2x slower) | 0.31 | 19.10M |

Reading:

- **attention_freq=2 wins on both axes**: +0.12 dB over the baseline AND
  1.56x faster training, with 2.4M fewer params. HEART has more attention
  than it needs at 20M scale.
- **GDFN rejected** (-0.37 dB vs attn2): doubling the per-block 3x3 conv
  density (CAB already provides spatial convs) hurt. The FFN stays a linear
  MLP.
- **Positional barely matters**: no-positional ties RIB at freq=1 (-0.01 dB,
  noise), while RoPE's fixed rotation actively hurt (-0.15 dB). RIB stays as
  cheap insurance; it is not a quality driver.
- **Window 48/64 closed**: window size is not a quality bottleneck at this
  crop size (w48 -0.07 dB, 38% slower). SST's large-window gains needed
  scaled crops+datasets this regime doesn't have.
- **Phase 2 confirmation** (in progress): w32 vs attn2 head-to-head at 100k
  iters (`HEART_{w32,attn2}_confirm.yml`, lr decay at 50k). The 20k ordering
  must survive to 100k before the release config changes.

### Release training (actual pipeline, 2026-08-25..28)

The release trained in five chained phases, all unattended with
`--auto_resume`. This is the pipeline that actually ran (not the earlier
3-phase design, which was superseded).

| Phase | Config | Iters | Result |
|---|---|---|---|
| 1. 4x pretrain (bicubic) | `options/train/RELEASE/4x_HEART_release.yml` | 615k (early-stopped) | 26.2139 dB |
| 2. 2x finetune | `options/train/RELEASE/2x_HEART_release.yml` | 200k | 33.9959 dB |
| 3. 4x OTF v1.0 (real-world) | `options/train/RELEASE/4x_HEART_OTF.yml` | 300k | 22.8473 dB |
| 4. 4x OTF v1.1 (anti-alias) | `options/train/RELEASE/4x_HEART_OTF_v11_finetune.yml` | 150k | 22.9529 dB |
| 5. 4x OTF perceptual (GAN) | `options/train/RELEASE/4x_HEART_OTF_gan.yml` | 150k | in progress |

Key operational lessons learned the hard way:

- **`--auto_resume` is mandatory** (CLI flag, not YAML). Without it an
  interrupted run restarts from iter 0 and archives the experiment dir.
- **Dataset:** LUCID CC0 v2 HC-512 (100,866 tiles, 512x512, CC0). LR generated
  with chainner CubicCatrom bicubic (the chaiNNer ecosystem standard).
- **OTF degradation:** Real-ESRGAN-style two-pass blur + noise + JPEG, applied
  on-the-fly by `realesrgandataset`.
- **ONNX export froze the desktop twice** (see Part 1 of this plan) β€” the fix
  is a per-process memory cap, not `nice`/`ionice`.

Pipeline scripts (durable, in `scripts/heart/`): `download_lucid.py`,
`prepare_lucid.py`, `run_release_sequence.sh`, `watch_download_then_prepare.sh`,
`export_one.py`, `export_minimal.py`, `infer_otf_comparison.py`,
`grid_diagnosis.py`. Logs in `/tmp/kilo/`.

## Variants

HEART 1.0 ships as a single architecture (`heart`, 19.2M params, embed 180,
6Γ—6 blocks, window 32, CAB compress/squeeze 3/30). A principled S/M/L scaling
family (width, depth, heads, rank, CAB) is a HEART 2.0 item β€” deliberately
not added before the release model proves the architecture.

## Positioning

HEART deliberately does **not** compete on benchmark PSNR. Research
architectures (HAT, DAT, RGT, Restormer, DRCT, Mamba-based SSMs) exist to
move leaderboard numbers, and they will out-score HEART on some val sets β€”
that is expected and accepted.

HEART competes on everything that happens *after* the paper: training
without babysitting, exporting without a fuse step, running in chaiNNer
tomorrow, and being understood by a new maintainer in 2028.

The rules that define the position:

- **One configuration.** No S/M/L variants, no fused/unfused checkpoint
  pairs, no "lite" forks. One architecture, one way to run it.
- **Zero custom kernels.** Everything a user needs ships in vanilla
  PyTorch; ONNX opset-17 export is verified, not promised.
- **Anything that loses the ablation protocol is removed from the code.**
  14 variants were measured; the winners stayed, the losers were deleted β€”
  not kept behind flags.
- **The architecture is boring on purpose.** The block is i-LN -> [RIB
  window attention | CAB] -> MLP. No dynamic routing, no deformable
  anything, no state-space branches.

If a technique requires a fuse step, a custom kernel, a shape assumption,
or a "use variant B only for inference" note, it does not belong in HEART.

## Deployment notes

- **Arbitrary input sizes:** `check_image_size` pads the input to a window
  multiple (reflect, replicate fallback for tiny images) and crops the
  output back. Verified at 32x32, 64x64, 128x128, 323x711, 777x65,
  720x1280. Users never handle padding themselves.
- **ONNX (verified, 2026-08-21, trained 20k checkpoint):** legacy tracer
  (`dynamo=False`), opset 17, dynamic H/W + batch. `onnx.checker` passes;
  onnxruntime output matches PyTorch (max abs diff 0.00026 at 128x128).
  Export pipeline for releases: export -> onnxsim -> ORT optimizer ->
  validate at 64x64 / 256x256 / 720x1280. Note: ORT *CPU* inference of
  large images needs several GB of RAM (a 16.7M transformer is not a CPU
  workload) β€” deploy on GPU providers (CUDA/TensorRT/DirectML).
- **NCNN: not supported, by design.** The fused SDPA node does not convert
  to NCNN, and window attention is not an NCNN-style workload anyway.
  HEART targets ONNX Runtime / TensorRT / DirectML; a mobile-conv sibling
  (NERVE) is the future NCNN member of the suite.
- **Large images β€” single pass vs tiling:** single pass fits up to ~720x1280
  LR on a 12 GB card (5.55 GB measured, bf16). 1080x1920 LR single-pass
  OOMs (~10.6 GB); tiled inference (tile 256, overlap 32, weighted blend)
  runs it at 0.86 GB and is seam-free (measured). Use chaiNNer seamless
  tiling / the repo's `tile_size` option for anything bigger than 720x1280.
- **Real-world degradations:** the backbone is degradation-agnostic. For
  denoising/JPEG cleanup train at `scale=1` (verified working) with a
  Real-ESRGAN-style degradation pipeline, same as HAT-R relates to HAT.
- **spandrel:** registered in-repo via `SPANDREL_REGISTRY` +
  `store_hyperparameters`, so traiNNer-redux/chainer forks that carry this
  repo load it directly. **Upstream-PR decision (made):** at PR time,
  `heart_arch.py`'s imports from `arch_util` (iLN) and `hat_iln_arch`
  (CAB, AffineTransform, DropPath, Mlp, PatchEmbed, PatchUnEmbed, Upsample)
  get vendored into the self-contained spandrel copy (~200 lines) rather
  than landing HAT-iLN in spandrel first. No code change in this repo until
  then.

## Experiment log (everything that shaped HEART)

Chronological record of every experiment, success and failure. The failures
matter as much as the wins β€” each one closed a direction.

| When | Experiment | Result | Decision / lesson |
|---|---|---|---|
| design | HAT-iLN base: RPB table -> RIB, OCAB removed, window 32 | β€” | Attention runs on FlashAttention/SDPA; simpler than HAT |
| 08-14 | Shifted windows via `torch.roll` | BUG: cyclic wrap coupled opposite borders (no mask) | Fixed: non-wrapping pad-and-partition |
| 08-14 | `F.pad` for V concat | minor | cleanup |
| 08-15 | 30k arch comparison HEART vs HAT_M vs ATD | 25.02 vs 25.05 vs OOM | tie with HAT_M at equal settings -> proceeded |
| 08-15 | Inference bench vs HAT_M | bf16: 1.16-1.24x faster, ~2.7x less VRAM; fp32 slower | deployment story is bf16 |
| 08-15 | torch.compile (smoke + harness) | +46-54% training, quality-neutral | adopted (`use_compile: true`) |
| 08-15 | compile_mode sweep | reduce-overhead 1.23 > default 1.18 > dynamic=False/TF32 1.19; max-autotune-no-cudagraphs == reduce-overhead with 2x compile time | reduce-overhead adopted; max-autotune rejected (28 SMs) |
| 08-15 | `fast_matmul` (TF32) | no-op under bf16 AMP | rejected |
| 08-15 | Crop/batch probes | lq96 bs4 OOM without checkpoint; lq96 bs2 fits, 9 windows | release crop = lq96 bs2 |
| 08-15 | `heart_light` (8.7M variant) | created, then removed | S/M/L family deferred to v2; 1.0 = one architecture |
| 08-17 | Rank ablation 8/18/32 @20k | 8 wins (24.782) > 32 (24.759) > 18 (24.735) | rank 8 default β€” positional capacity is NOT the bottleneck |
| 08-18 | Window ablation 48/64 @20k | w48 loses 0.07 dB, 38% slower; w64 skipped | spatial context is NOT the bottleneck at lq96 |
| 08-18 | Attention frequency 1/2/3 @20k | freq=2 wins: 24.628 (+0.12 dB) at 1.56x speed; freq=3 drops | attention density sweet spot = every 2nd block |
| 08-19 | RoPE instead of RIB @20k | loses 0.15 dB (24.36 vs 24.51) | learned low-rank positional beats fixed rotation |
| 08-19 | No positional at all @20k | ties RIB at freq=1 (24.497 vs 24.506) | RIB stays as cheap insurance; not a quality driver |
| 08-19 | GDFN (gated depthwise FFN) @20k | loses 0.37 dB (24.25 vs 24.63) | more per-block conv density hurts (CAB already convs) |
| 08-20 | SwiGLU FFN @20k | loses decisively: best 23.74 @12k, declined to 23.21 @20k | gated FFN destabilizes the i-LN-rescaled block β€” MLP stays |
| 08-20 | QK-Norm @20k | 24.6334 β€” tie with attn2 (+0.005) | tie-breaker: plain (simplest) β€” redundant with i-LN + RIB |
| 08-20 | SwiGLU + QK-Norm @20k | 23.9576 | gate instability again |
| 08-20 | LayerScale @20k | 24.0288 (βˆ’0.60) | 1e-6 init over-damps the branches; existing i-LN+conv_scale damping is right |
| 08-20 | Final cleanup | all test switches (ffn_type/qk_norm/layer_scale) stripped; canonical default attention_freq=2 | the file is the single-path measured-best architecture |
| 08-16 | Resume test | `--auto_resume` is a CLI flag, not YAML | all training commands pass it |
| 08-16 | Tiled inference | seam-free (seam diff == interior diff) | large-image tiling safe |
| 08-14..20 | ONNX export | legacy tracer opset 17 + checker + onnxruntime pass | deployment path verified |

Process lessons (drivers/pipeline): /tmp gets wiped on reboot β€” logs restored
from durable experiment copies; `ls dir/*.png | wc -l` breaks at ~100k files
(ARG_MAX) β€” use `find`; drivers are idempotent/self-healing and training runs
deprioritized (nice/ionice) so the machine stays usable.

## Considered-and-decided register

Every idea that was considered, with the decision and reason. If someone asks
"why didn't you use X" β€” the answer is here. Status: βœ… integrated (and
measured), ❌ rejected by measurement, 🚫 discarded by analysis (never run),
⏳ pending, πŸ“Œ deferred (not forgotten).

### Integrated

| Idea | Why |
|---|---|
| βœ… RIB instead of RPB table | Concatenated position features make attention FlashAttention/SDPA-compatible; measured (ties HAT_M, rank 8 suffices) |
| βœ… i-LN normalization | Input-adaptive holistic norm; bf16 stability is the paper's core claim, verified in training |
| βœ… CAB conv branch | Local inductive bias; the measured carrier of quality at this scale |
| βœ… Non-wrapping shifted windows | Fixes the cyclic-wrap border bug without the Swin mask (which would kill the flash kernel) |
| βœ… `attention_freq=2` | Measured winner: +0.12 dB at 1.56x speed vs attention-every-block |
| βœ… rank 8 RIB | Measured winner of 8/18/32 sweep |
| βœ… Window 32 | w48/64 measured losers at lq96 |
| βœ… torch.compile reduce-overhead | +46% training speed, quality-neutral; best of the mode sweep |
| βœ… bf16 AMP + gradient checkpointing | Verified stable; 1.4-1.7 GB training VRAM |
| βœ… Chainer CubicCatrom LR (release data) | The chaiNNer ecosystem standard β€” models behave as users generate LR |

### Rejected by measurement (ran, lost)

| Idea | Result | Why rejected |
|---|---|---|
| ❌ RoPE positional | βˆ’0.15 dB @20k | fixed rotation worse than learned low-rank positional |
| ❌ No positional | ties RIB at freq=1, but attn2+RIB wins overall | RIB stays as cheap insurance |
| ❌ GDFN gated-FFN | βˆ’0.37 dB @20k | more per-block conv density hurts (CAB already provides spatial convs) |
| ❌ SwiGLU FFN (param-matched 4/3) | best 23.74 @12k, declined to 23.21 @20k (βˆ’0.9 to βˆ’1.4 dB vs attn2) | the multiplicative gate destabilizes the i-LN-rescaled block (loss degraded after 12k) |
| ❌ SwiGLU + QK-Norm | 23.9576 @20k | same gate instability |
| ❌ LayerScale (init 1e-6) | 24.0288 @20k (βˆ’0.60 dB) | the near-zero init suppresses the branches through warmup; i-LN rescale can't compensate β€” the block's existing damping (i-LN + conv_scale 0.01) is the right amount |
| βš–οΈ QK-Norm | 24.6334 β€” tie with attn2 (24.6282, +0.005 dB) | tie-breaker: plain wins (no component for zero measured gain) |
| ❌ Window 48 / 64 | w48 βˆ’0.07 dB + 38% slower; w64 skipped | spatial context is not the bottleneck at lq96 |
| ❌ `attention_freq=3` | βˆ’0.07 dB vs freq=2 | density sweet spot is every 2nd block |
| ❌ rank 18 / 32 | rank 8 wins | positional capacity is not the bottleneck |
| ❌ max-autotune compile modes | == reduce-overhead, 2x compile time | tuned GEMM search falls back on 28-SM GPUs |
| ❌ `fast_matmul` (TF32) | no-op | everything already runs bf16 on tensor cores |
| ❌ `compile(dynamic=False)` | no gain | shapes were effectively static anyway |
| ❌ `heart_light` 8.7M variant | created, then removed | hand-picked config, not a principled family; S/M/L is a v2 item |

### Discarded by analysis (never run β€” the reason is the evidence)

| Idea | Why discarded |
|---|---|
| 🚫 OCAB (HAT's overlapping cross-attention) | RPB tables + gather construction + memory are exactly what RIB/HEART removed; the 30k ablation already ties HAT_M *with* OCAB β€” the missing piece isn't needed |
| 🚫 Swin-style attention mask (shift fix alternative) | an additive mask forces SDPA off the flash kernel β€” the whole point of RIB |
| 🚫 RepCAB (reparameterized CAB) | `fuse()`-before-export creates two weight formats and breaks the drop-in property, for ~2-3% inference; quality hypothesis contradicted by the GDFN result |
| 🚫 MDTA transposed channel attention | global channel-covariance context β€” falsified direction (see w48) |
| 🚫 Stripe/multi-axis attention (HMA) | full-row/column context β€” same falsified hypothesis, plus partitioning complexity |
| 🚫 Dense residuals (DRCT) | groups and blocks already have residuals; dense skips grow activation memory |
| 🚫 Mamba/SSM backbones | custom CUDA kernels (`selective_scan`) violate the zero-custom-kernel constraint |
| 🚫 IET adaptive token selection | dynamic gather/indexing breaks dynamic ONNX tracing |
| 🚫 UCAN / SAT / FPLIA | sub-1M-parameter tier techniques; irrelevant at ~16.7M |
| 🚫 Muon optimizer | its gains are LLM-pretraining-scale; the RIB net is ~4k params |
| 🚫 FlashBias Triton kernel | HEART already *is* the SDPA variant of FlashBias; no bias-add attention exists to accelerate |
| 🚫 LAformer swap-in | different architecture, not an optimization; the deployability premise is already satisfied by HEART |
| 🚫 Triton-fused RIB concat | ~2-4% of a step at the cost of custom kernels; torch.compile already fuses the elementwise ops |
| 🚫 4D-native block refactor | single-digit % after compile's 46% landed; complexity not justified |
| 🚫 Hybrid window schedule (16/32 mix) | smaller windows contradict the measured landscape (w48 lost; context isn't the lever) |
| 🚫 lq128 crops | OOM without checkpointing and pathological compile on this GPU |
| 🚫 Progressive patch training (LQ64 -> LQ96) | measured apples-to-apples (bs2, compiled): lq64 3.86 it/s vs lq96 1.52 it/s (2.5x per iteration), but pixel throughput only 31.6 vs 28.0 kpx/s (+13%) β€” and lq64 crops see 4 windows vs lq96's 9, weakening cross-window mixing. ~8% total wall-time for real pipeline complexity and quality risk |

### Pending

None β€” the ablation program is complete. 14 variants measured; the core
(RIB rank-8 + attention_freq=2 + window 32 + CAB + plain MLP + i-LN) beat or
tied all of them. The architecture is at a measured local optimum.

### Deferred (not forgotten β€” see the 2.0 backlog)

πŸ“Œ INT8 quantization, TensorRT export CI, RIB positional caching in eval, RIB
dtype sweep, ENAF-style early exits for 4K+ images, S/M/L family, larger
training crops (bigger GPU), video adaptation, rank 2/4 micro-sweep,
attention-freq bs4 batch check, MLP ratio sweep.

## HEART 2.0 backlog

Frozen for 1.0 (do not touch the architecture while the release model trains).
Ideas that need a fully trained 1.0 baseline before they can be judged
properly β€” the rule is: only adopt what wins on *quality per second* against
the 1.0 model, not what moves a benchmark number.

| Idea | Why it's on the list | Effort |
|---|---|---|
| Gated-Dconv FFN (GDFN, Restormer) | MEASURED (20k, ratio 2.0, freq=2): lost 0.37 dB (24.25 vs 24.63) β€” doubling per-block 3x3 conv density (CAB already provides spatial convs) hurt. Code removed. | β€” |
| SwiGLU FFN + QK-Norm | External review (2026-08-20): modernize the one remaining vanilla component (the MLP) + per-head QK LayerNorm for logit stability. TESTING: three 20k runs on the attn2 base (swiglu-only, qknorm-only, both). SwiGLU param-matched at ratio 4/3 (their 8/3 would have doubled params). | Low-Med |
| RepCAB (reparameterized CAB) | Train-time multi-branch conv fused at export. REJECTED for 1.0: the fuse()-before-export step conflicts with the trivial-export property that is part of the maintainability goal. | β€” |
| Transposed channel attention (MDTA) | Global channel-covariance attention β€” chases global context, which the w48 result already falsified as the lever. Backlog only if GDFN-style spatial enrichment proves the direction. | Medium |
| Multi-axis stripe attention (HMA) | Full-row/column context β€” same global-context rationale, contradicted by w48 losing at 20k. Skip unless a crop-size scaling study revives the context hypothesis. | Medium |
| Dense residual connections (DRCT) | Zero-param dense skips; NTIRE-2024 proven, but our blocks/groups already have residuals, and dense skips grow activation memory. Low expected value; cheap to test one day. | Low |
| ENAF early-exit routing | Inference-time wrapper for huge images, not an architecture change; post-release option only. | High |
| IET adaptive token selection | Dynamic gather/indexing risks breaking dynamic ONNX tracing β€” rejected on deployability grounds. | β€” |
| Mamba/SSM backbones | Custom CUDA kernels (selective_scan) violate the zero-custom-kernel constraint β€” rejected. | β€” |
| Larger windows (48 / 64) | MEASURED (20k, lq96): window 48 lost to 32 (-0.07 dB) and cost 38% more time. Closed for this crop/GPU regime unless crops scale too. | β€” |
| Attention frequency (H C H C) | MEASURED (20k, lq96): `attention_freq=2` beat the baseline (+0.12 dB) at 1.56x speed. Live candidate for the release config; longer-run confirmation pending. | Low-Med |
| RoPE instead of RIB | MEASURED (20k, lq96): RoPE lost 0.15 dB to RIB (24.36 vs 24.51) β€” learned low-rank positional wins; fixed rotations can't substitute. Closed. | β€” |
| Rank beyond release range (4/12/24/…) | The release gate only tests 8/18/32. Fine-grained rank scaling could shave the positional capacity further. | Low |
| RIB positional caching in eval | q_pos/k_pos recomputed every forward in every block; tiny but free to cache for inference. Micro-opt, measure before doing. | Low |
| RIB dtype sweep (bf16/fp16 vs fp32) | fp32 cast is for fp16-inference stability; whether bf16 RIB changes quality/speed is unmeasured. | Low |
| CAB variants / S/M/L family | An earlier hand-picked `heart_light` (8.7M) was removed for 1.0; a principled S/M/L scaling rule (width, depth, heads, rank, CAB) is the v2 way to extend the family. | Low-Med |
| TensorRT + export CI | ONNX verified; TensorRT untested. A CI export test (PyTorch→ONNX→ORT→TensorRT, compare max-abs-error/PSNR/runtime) would harden the deployment claim. | Medium |
| Quantization (INT8) | Untested; a 3060-class GPU model could get much faster. | Medium |
| Larger training crops | lq 128+ needs checkpointing + slower it/s; only test with a bigger GPU. | Medium |
| Video adaptation | temporal stability (like TSPAN) β€” separate scope. | High |
| Cyclic vs reflect shift | Deliberately closed for 1.0 (cyclic + mask kills the flash kernel). Revisit only if a flash-compatible exact-cyclic scheme emerges. | β€” |

Priority order after 1.0: confirm attention_freq=2 on a longer run (if the
release adopts it, this is moot) -> export hardening (TensorRT/CI) ->
RoPE swap-in benchmark -> rank/capacity scaling. Everything is measured
against the frozen 1.0 model, not against papers.