spec: OFI momentum enrichment + dense micro-reward design
Dense per-bar reward using order flow momentum × price confirmation. OFI embedding MLP (16→8 via cuBLAS) feeds into Mamba2 history AND attention input. Replaces sparse exit-only reward with continuous temporal signal for the SSM and attention heads to learn from. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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)
|
||||
Reference in New Issue
Block a user