diff --git a/docs/superpowers/specs/2026-04-19-ofi-momentum-dense-reward-design.md b/docs/superpowers/specs/2026-04-19-ofi-momentum-dense-reward-design.md new file mode 100644 index 000000000..dfa892d3a --- /dev/null +++ b/docs/superpowers/specs/2026-04-19-ofi-momentum-dense-reward-design.md @@ -0,0 +1,164 @@ +# OFI Momentum Enrichment + Dense Micro-Reward Design Spec + +## Goal + +Provide dense per-bar gradient signal for the temporal pipeline (Mamba2 + attention) by: +1. Computing OFI momentum features (temporal derivatives of order flow) +2. Feeding them into Mamba2 history AND attention input via a learned embedding +3. Using flow momentum as a dense micro-reward for positioned bars + +Currently the model only receives reward at trade exit (sparse, delayed). Between open and close, reward=0 — the temporal pipeline has nothing to learn from. This spec adds dense per-bar feedback based on order flow dynamics and price confirmation. + +## Architecture + +### OFI Momentum Feature Pipeline + +**Input:** Raw OFI features at indices 42-49 of the 50-dim feature vector: +- 42: bid_ask_spread, 43: depth_imbalance_L1, 44: queue_pressure +- 45: spread_velocity, 46: depth_change_rate, 47: trade_arrival_rate +- 48: VPIN, 49: Kyle's_lambda + +**Delta computation:** Per bar, compute `delta_ofi[i] = ofi[i] - prev_ofi[i]` for all 8 features. Previous values stored in per-episode portfolio state `ps[30..37]`. Updated every bar. + +**Learned embedding MLP:** `[raw_ofi(8) ; delta_ofi(8)] = 16-dim → MLP(16→8) → ofi_embed(8-dim)` +- Implementation: cuBLAS SGEMM `[8, 16] × [16, B] → [8, B]` + elementwise bias+ReLU kernel +- Same pattern as `trade_plan_forward` (cuBLAS 2-layer MLP) +- Weights: 16×8 + 8 bias = 136 params. Xavier init. Trained jointly via C51/MSE backward. +- Kernel: `ofi_momentum_embed` — builds 16-dim input from features + ps, writes to scratch buf +- cuBLAS: W_ofi @ input → pre_h, then `add_bias_relu_f32` kernel → ofi_embed[B, 8] + +### Entry Point 1: Mamba2 History Enrichment + +**Current:** `h_history[B, K, SH2]` stores K=8 steps of SH2=256-dim trunk output. +**New:** `h_history[B, K, SH2+8]` — OFI embed appended per timestep. + +Changes: +- `mamba2_update_history` kernel: copies `[h_s2(SH2) ; ofi_embed(8)]` instead of just `h_s2` +- W_A, W_B projection GEMMs: k dimension changes from SH2 to SH2+8 + - cuBLAS GEMM: m=STATE_D(16), n=B*K, k=SH2+8(264) — trivial dimension change + - W_A, W_B parameter buffers grow from `[SH2, STATE_D]` to `[SH2+8, STATE_D]` +- W_C stays `[SH2, STATE_D]` — operates on the scan output, not the history input +- Mamba2 scan kernel: unchanged (operates on pre-projected values) +- `mamba2_h_enriched` output: stays `[B, SH2]` (temporal context added to trunk, not widened) + +Total param increase: 2 × 8 × 16 = 256 floats (1KB). Negligible. + +### Entry Point 2: Attention Input Enrichment + +**Current:** Attention input is `h_s2[B, SH2=256]`, projected to Q/K/V via `[D, D]` weight matrices. +**New:** Attention input is `[h_s2 ; ofi_embed] = [B, SH2+8=264]`. + +To keep all downstream attention dimensions unchanged (D=256, num_heads=4, head_dim=64): +- Add a projection GEMM: `h_attn_in[B, 256] = W_proj[256, 264] @ [h_s2; ofi_embed][264, B]` + - cuBLAS SGEMM: m=256, n=B, k=264. One extra GEMM per forward pass. + - Alternatively: expand W_Q/W_K/W_V from [256, 256] to [256, 264] directly — avoids the extra GEMM but changes 3 weight matrices. +- **Recommendation:** Expand W_Q/W_K/W_V input dimension from 256 to 264. This is simpler (no extra buffer, no extra GEMM) and the cuBLAS forward GEMMs just have k=264 instead of k=256. The existing GEMM descriptors need re-creation with the new k dimension. + +Total param increase: 3 × 256 × 8 = 6144 floats (24KB). Negligible. + +### Entry Point 3: Dense Micro-Reward + +**Location:** `experience_kernels.cu` env_step_batch kernel, replacing lines 1715-1721 (the `-0.0001 * |position|` holding cost). + +**Formula:** +```c +if (!segment_complete && fabsf(position) > 0.001f && micro_reward_scale > 0.0f) { + /* OFI momentum: mean of key delta features (depth_imb, VPIN, queue_press, trade_arrival) */ + float delta_depth = ofi[43] - ps[31]; /* depth_imbalance delta */ + float delta_vpin = ofi[48] - ps[36]; /* VPIN delta */ + float delta_queue = ofi[44] - ps[32]; /* queue_pressure delta */ + float delta_tarr = ofi[47] - ps[35]; /* trade_arrival_rate delta */ + + float flow_momentum = sign_pos * 0.25f * (delta_depth + delta_vpin + delta_queue + delta_tarr); + + /* Price confirmation: current bar's directional move (no lookahead) */ + float bar_move = (raw_close - open_price) / fmaxf(atr_abs, 1e-6f); + float price_confirm = sign_pos * bar_move; + + /* Informed flow intensity: VPIN × Kyle's Lambda */ + float vol_informed = ofi[48] * ofi[49]; + + /* Composite quality score */ + float quality = flow_momentum * (1.0f + vol_informed) + 0.5f * price_confirm; + + /* Bounded micro-reward — holding cost absorbed (negative quality = holding penalty) */ + reward = micro_reward_scale * tanhf(quality / 3.0f); +} +``` + +**Config:** `micro_reward_scale = 0.1` in `dqn-production.toml` (tunable via hyperopt). Should be small enough that sparse trade P&L (~0.5-5.0) dominates at exit, but large enough to provide gradient signal between bars. + +**Previous OFI storage:** After computing micro-reward, update `ps[30..37]` with current OFI values for next bar's delta computation. + +### Backward Pass + +The OFI momentum MLP (16→8) needs gradient flow. This follows the existing pattern: +- cuBLAS GEMM for dW (weight gradient): `dW = d_ofi_embed^T @ input` +- cuBLAS GEMM for dx (input gradient): `dx = W^T @ d_ofi_embed` (not needed — OFI features are frozen input) +- Bias grad: 2-phase shared-memory reduce (existing pattern) +- Adam update: existing `dqn_adam_update_kernel` + +The gradient flows through Mamba2 backward (already cuBLAS) and attention backward (already cuBLAS) into the OFI embed MLP via the chain rule. No new backward kernels needed — just wider GEMM dimensions. + +### Performance Impact + +| Operation | Current | New | Delta | +|-----------|---------|-----|-------| +| OFI delta kernel | — | 0.01ms | +0.01ms | +| OFI embed GEMM | — | 0.05ms | +0.05ms | +| Mamba2 proj GEMMs (k+8) | 0.45ms | ~0.47ms | +0.02ms | +| Attention Q/K/V GEMMs (k+8) | 0.77ms | ~0.79ms | +0.02ms | +| Micro-reward (in env_step) | 0 | ~0 (elementwise) | ~0 | +| **Total per-step** | **~45ms** | **~45.1ms** | **+0.1ms** | + +Negligible impact. The wider GEMM k-dimensions add <1% to cuBLAS time. + +### VRAM Impact + +| Buffer | Size | Note | +|--------|------|------| +| ofi_embed_buf [B, 8] | 0.5MB | scratch | +| ofi_input_buf [B, 16] | 1.0MB | scratch | +| W_ofi [8, 16] + bias [8] | 0.5KB | params | +| h_history wider [B, K, SH2+8] | +2MB | existing buf grows | +| W_A, W_B wider [SH2+8, STATE_D] | +1KB | existing params grow | +| W_Q/K/V wider [256, 264] | +24KB | existing params grow | +| ps[30..37] | 0 | reuses existing portfolio state slots | +| **Total** | **~3.5MB** | 0.01% of 38GB free | + +### Config Changes + +```toml +[reward] +micro_reward_scale = 0.1 # was 0.0 (disabled) +``` + +No new config fields — `micro_reward_scale` already exists. + +### Files Modified + +| File | Change | +|------|--------| +| `experience_kernels.cu` | OFI delta kernel + micro-reward body in env_step_batch | +| `gpu_dqn_trainer.rs` | OFI embed MLP buffers + GEMM + wider Mamba2 params + backward | +| `gpu_attention.rs` | Wider W_Q/W_K/W_V (k=264), re-create GEMM descs | +| `mamba2_temporal_kernel.cu` | Update history kernel to copy SH2+8 dims | +| `batched_forward.rs` | Wider attention input concatenation | +| `batched_backward.rs` | Wider attention backward dimensions | +| `dqn-production.toml` | `micro_reward_scale = 0.1` | + +### Success Criteria + +1. Smoke tests 20/20 pass +2. Per-step time stays <50ms (no regression from 45ms baseline) +3. Atom utilization remains 40-60% (PopArt + dense reward working together) +4. **Trade count drops below 1M** (model learns to hold via dense feedback) +5. **Val_Sharpe improves to > -100** within 50 epochs (dense signal accelerates learning) +6. Mamba2 temporal context shows non-trivial values (not all-zero — temporal pipeline is learning) + +### Non-Goals + +- No changes to C51 atom support or PopArt (already fixed this session) +- No changes to IQN/IQL (separate heads, unchanged) +- No new exploration mechanisms +- No hyperopt integration (use fixed `micro_reward_scale=0.1` initially)