spec: Cross-Pollinated Branch Intelligence — 5-component unified design
Combines TFT Variable Selection, GLU Gating, Cross-Branch Q Attention, Liquid adaptive tau, and Mamba-2 Selective Replay into unified branch architecture. Plus 11-item dead code cleanup. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -0,0 +1,321 @@
|
||||
# Cross-Pollinated Branch Intelligence — Design Spec
|
||||
|
||||
**Goal:** Combine Variable Selection (TFT), Gated Linear Units (TFT/Llama), Cross-Branch Q Attention, Liquid adaptive tau, and Mamba selective replay into a unified branch architecture that selects, gates, conditions, coordinates, and adapts.
|
||||
|
||||
**Problem:** The 4-branch DQN treats each branch as an independent FC layer reading the same h_s2 trunk. Every branch sees every feature equally (no selection), uses all-or-nothing ReLU gating (no soft suppression), produces Q-values in isolation (no coordination), and learns at a uniform rate (no per-branch adaptation). The replay buffer prioritizes transitions by TD-error magnitude regardless of current model state (no curriculum).
|
||||
|
||||
**Architecture:** Five additive layers on top of the existing 4-branch DQN + Layer 3 direction conditioning:
|
||||
|
||||
1. **Variable Selection** — per-branch softmax mask on h_s2 features
|
||||
2. **GLU Gating** — sigmoid gate replaces ReLU in branch FC
|
||||
3. **Cross-Branch Q Attention** — 2-head attention on the 12 Q-values
|
||||
4. **Liquid Tau** — per-branch adaptive learning rate via continuous-time ODE
|
||||
5. **Mamba-2 Selective Replay** — learned state-conditioned priority gate (from Mamba2's SelectiveStateSpace/StateImportance pattern in `crates/ml-supervised/src/mamba/selective_state.rs`)
|
||||
|
||||
Plus dead code cleanup (11 items across 6 files).
|
||||
|
||||
---
|
||||
|
||||
## Component 1: Variable Selection (from TFT)
|
||||
|
||||
### Problem
|
||||
All 4 branches read the same h_s2[B, SH2] with equal weight on all features. Direction should attend to trend features (MACD, RSI), magnitude to volatility (ATR), order to microstructure (OFI), urgency to regime (CUSUM, ADX). Currently the branch FC must learn feature relevance implicitly through its weight matrix — wasting capacity.
|
||||
|
||||
### Solution
|
||||
Each branch d learns a parameter vector `v_select_d[SH2]`. Softmax normalizes it into a feature importance distribution. Element-wise multiply produces `h_masked = h_s2 * softmax(v_select_d)`.
|
||||
|
||||
```
|
||||
mask_d[SH2] = softmax(v_select_d[SH2]) // learned, sum-to-1
|
||||
h_masked[B, SH2] = h_s2[B, SH2] * mask_d[SH2] // broadcast multiply
|
||||
```
|
||||
|
||||
### Properties
|
||||
- 4 × SH2 = 1,024 extra parameters (negligible)
|
||||
- No extra GEMM — one softmax (SH2=256, trivial) + one element-wise multiply per branch
|
||||
- Interpretable — softmax weights directly show which features each branch uses
|
||||
- Gradient flows through softmax to v_select: features that improve branch loss get higher weight
|
||||
- Initialized to uniform (1/SH2) so training starts unconstrained
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `variable_select_mask` in experience_kernels.cu: precomputes `softmax(v_select_d)` once per batch (not per sample — mask is shared across batch)
|
||||
- 4 parameter vectors stored in params_buf after existing weights (indices 26-29)
|
||||
- Element-wise multiply fused into the branch FC launch (multiply input before GEMM)
|
||||
|
||||
---
|
||||
|
||||
## Component 2: Gated Linear Units (from TFT GRN / Llama)
|
||||
|
||||
### Problem
|
||||
ReLU activation in branch FC is all-or-nothing: negative activations are zeroed, positive pass unchanged. This creates dead neurons (permanently zero for some input distributions) and can't suppress positive-but-irrelevant activations.
|
||||
|
||||
### Solution
|
||||
Replace ReLU with GLU (Gated Linear Unit):
|
||||
|
||||
```
|
||||
Current: h_bd = ReLU(W_bdf @ input + b_bdf)
|
||||
Proposed: gate = sigmoid(W_gate_d @ input + b_gate_d) [B, AH]
|
||||
value = W_bdf_d @ input + b_bdf_d [B, AH]
|
||||
h_bd = gate * value [B, AH]
|
||||
```
|
||||
|
||||
The sigmoid gate learns a smooth [0,1] suppression per activation. Unlike ReLU, it can suppress positive values ("this feature is active but irrelevant") and partially pass negative values ("this feature is slightly negative but informative").
|
||||
|
||||
### Properties
|
||||
- Doubles the branch FC parameters: 4 extra weight matrices [AH, K_d] + 4 bias vectors [AH]
|
||||
- Total: ~262,144 + 1,024 = 263,168 extra parameters
|
||||
- One extra GEMM per branch (4 total, same shape as existing branch FC)
|
||||
- Proven in TFT, Llama, PaLM — consistently outperforms ReLU in gated architectures
|
||||
- No dead neuron problem — sigmoid gate is always differentiable
|
||||
|
||||
### Implementation
|
||||
- Gate weights stored after variable selection vectors (indices 30-37: w_gate0, b_gate0, ..., w_gate3, b_gate3)
|
||||
- Forward: two parallel GEMMs per branch (gate path + value path), then element-wise multiply
|
||||
- Backward: standard GLU backward — dgate = dout * value * sigmoid'(gate_pre), dvalue = dout * gate
|
||||
- cuBLASLt RELU_BIAS epilogue is no longer used for branch FC — replaced by separate sigmoid kernel after gate GEMM
|
||||
- New CUDA kernel `glu_combine`: `output[i] = sigmoid(gate[i]) * value[i]`
|
||||
- New CUDA kernel `glu_backward`: computes dgate and dvalue from doutput
|
||||
|
||||
### Forward pass detail
|
||||
```
|
||||
// Variable selection (Component 1)
|
||||
h_masked = h_s2 * softmax(v_select_d) [B, SH2]
|
||||
|
||||
// For d==1: h_masked = [h_masked; Q_dir/dz] (Layer 3, already done)
|
||||
input_d = h_masked (or concat for d==1) [B, K_d]
|
||||
|
||||
// GLU: two parallel GEMMs
|
||||
gate_pre = W_gate_d @ input_d + b_gate_d [B, AH] — GEMM 1
|
||||
value = W_bdf_d @ input_d + b_bdf_d [B, AH] — GEMM 2 (existing)
|
||||
h_bd = sigmoid(gate_pre) * value [B, AH] — fused kernel
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Component 3: Cross-Branch Q Attention
|
||||
|
||||
### Problem
|
||||
After computing 12 Q-values (4 branches × 3 actions), action selection treats them independently. But action quality is interdependent: Long+Full is aggressive (direction × magnitude correlation), Long+Small is cautious (different risk profile). The Q-values should coordinate.
|
||||
|
||||
Layer 3 provides direction→magnitude conditioning at the FC level. Cross-branch attention provides ALL-to-ALL coordination at the Q-value level.
|
||||
|
||||
### Solution
|
||||
Apply 2-head self-attention on the 12 Q-values:
|
||||
|
||||
```
|
||||
Q_raw[B, 12] = compute_expected_q(v_logits, adv_logits)
|
||||
|
||||
// Attention on Q-values
|
||||
Q_proj = Q_raw @ W_QKV[3, 12, 12] + b_QKV[3, 12] // Q, K, V projections
|
||||
// 2 heads, dim 6 each
|
||||
For each head h:
|
||||
attn_h = softmax(Q_h @ K_h^T / sqrt(6)) // [B, 6, 6] (tiny)
|
||||
head_h = attn_h @ V_h // [B, 6]
|
||||
concat = [head_0; head_1] // [B, 12]
|
||||
Q_coord = concat @ W_O[12, 12] + b_O[12] // output projection
|
||||
Q_coord = layernorm(Q_coord + Q_raw) // residual + LN
|
||||
```
|
||||
|
||||
### Properties
|
||||
- 624 parameters total (3×144 QKV + 144 W_O + 24 biases + 24 LN)
|
||||
- 12-dim vectors — the attention computation is trivially cheap (no GEMM needed, just register-level dot products)
|
||||
- Single CUDA kernel, one thread per sample, all computation in registers
|
||||
- Residual connection ensures Q_coord ≈ Q_raw initially (W_O, W_QKV initialized near zero)
|
||||
- Generalizes Layer 3: direction↔magnitude interaction emerges naturally, plus order↔urgency and all other pairs
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `cross_branch_q_attention` in experience_kernels.cu
|
||||
- 624-param weight buffer, separate from DQN params (own Adam state)
|
||||
- Runs AFTER compute_expected_q, BEFORE action selection
|
||||
- Backward: standard attention backward on 12-dim vectors (trivial)
|
||||
- Also runs during training Q-stats computation (gpu_dqn_trainer.rs reduce_current_q_stats)
|
||||
|
||||
---
|
||||
|
||||
## Component 4: Liquid Tau — Per-Branch Adaptive Learning Rate
|
||||
|
||||
### Problem
|
||||
All 4 branches learn at the same rate via a shared Adam optimizer. But branches have different convergence speeds: direction (3 high-impact actions) converges faster than urgency (3 low-impact actions). The single Q-gap momentum (spread velocity) is a global modulator — it can't help a stuck branch while slowing a converging one.
|
||||
|
||||
### Solution
|
||||
Per-branch continuous-time dynamics inspired by Liquid Time-Constant (LTC) neurons:
|
||||
|
||||
```
|
||||
For each branch d:
|
||||
q_gap_d = max(Q_d) - mean(Q_d) // per-branch Q-gap
|
||||
velocity_d = q_gap_d - q_gap_ema_d // per-branch velocity
|
||||
|
||||
// Liquid ODE: tau controls adaptation speed
|
||||
tau_d = tau_min + (tau_max - tau_min) * sigmoid(velocity_d / delta_z)
|
||||
|
||||
// Branch scale evolves toward 1.0 at rate 1/tau_d
|
||||
branch_scale_d += (1.0/tau_d) * (1.0 - branch_scale_d) * dt
|
||||
```
|
||||
|
||||
- Fast learning (positive velocity) → large tau → slow adaptation → don't overshoot
|
||||
- Stuck (zero velocity) → small tau → fast adaptation → push harder
|
||||
- Shrinking Q-gap → very small tau → emergency push → prevent collapse
|
||||
|
||||
### Properties
|
||||
- 4 per-branch Q-gap EMAs (replace single q_gap_ema from spread momentum)
|
||||
- tau_min = 0.01, tau_max = 1.0 (from existing Liquid hyperparameters)
|
||||
- dt = 1.0 per training step (discrete approximation of continuous ODE)
|
||||
- Zero extra parameters — only CPU-side state (4 floats EMA + 4 floats branch_scale)
|
||||
- The existing `branch_scales[B, 4]` buffer from IQL carries the per-branch modulation — liquid tau writes to it
|
||||
- Subsumes the global spread_velocity modulator (spread_velocity becomes derived from branch 0's liquid tau)
|
||||
|
||||
### Implementation
|
||||
- In `reduce_current_q_stats`: compute per-branch Q-gaps from the 12 Q-values
|
||||
- In `update_eval_v_range` (or new method): run liquid ODE for each branch
|
||||
- Write updated branch_scales to the pinned device-mapped buffer
|
||||
- The C51 grad kernel already reads branch_scales — no kernel changes needed
|
||||
|
||||
---
|
||||
|
||||
## Component 5: Mamba-2 Selective Replay
|
||||
|
||||
### Problem
|
||||
PER (Prioritized Experience Replay) uses `|TD_error| + epsilon` as priority. This is state-agnostic: a transition with TD_error=0.5 gets the same priority whether the model is currently struggling with that state region or has already mastered it. The model wastes training on transitions it's already learned.
|
||||
|
||||
### Solution
|
||||
Learned selectivity gate conditioned on the current model state, inspired by Mamba2's selective state space mechanism:
|
||||
|
||||
```
|
||||
// During training, after forward pass produces save_h_s2:
|
||||
selectivity[B] = sigmoid(W_sel @ save_h_s2[B, SH2] + b_sel) // [B, 1]
|
||||
|
||||
// Update PER priorities with state-conditioned relevance
|
||||
priority[i] = |TD_error[i]| * (1 + selectivity[i]) + epsilon
|
||||
|
||||
// Train selectivity gate via BCE on gradient magnitude feedback
|
||||
grad_mag[i] = ||per-sample gradient from sample i|| // proxy for "usefulness"
|
||||
target[i] = grad_mag[i] / mean(grad_mag) // normalize to ~1.0
|
||||
L_sel = BCE(selectivity, clamp(target, 0, 1)) // binary cross-entropy
|
||||
// Separate Adam update for W_sel, b_sel
|
||||
```
|
||||
|
||||
### Properties
|
||||
- 257 extra parameters: W_sel[1, SH2] + b_sel[1]
|
||||
- Separate Adam state: 2 × 257 = 514 floats (negligible)
|
||||
- Forward: one GEMM [1, SH2] per training batch → sigmoid → scale priorities
|
||||
- The gate learns which trunk activations predict useful training signal
|
||||
- Early training: selectivity ≈ 0.5 (uniform) — behaves like standard PER
|
||||
- Late training: selectivity differentiates — focuses on frontier states where model is uncertain
|
||||
- Per-sample gradient norm computed from the existing grad_buf (sum of squared per-sample gradients)
|
||||
|
||||
### Implementation
|
||||
- W_sel and b_sel stored in a separate small buffer (not in main params_buf)
|
||||
- New CUDA kernel `selectivity_forward`: `sel[i] = sigmoid(dot(W_sel, h_s2[i]) + b_sel)` — one thread per sample
|
||||
- New CUDA kernel `selectivity_backward`: BCE gradient → dW_sel, db_sel
|
||||
- Per-sample gradient norm: reduce grad_buf per-sample (need per-sample loss tracking, already available via is_weights)
|
||||
- Priority update in PER: multiply existing priority by `(1 + sel)` after each training step
|
||||
- Separate Adam update for W_sel/b_sel (2 momentum buffers, tiny)
|
||||
|
||||
### Training loop integration
|
||||
```
|
||||
1. Sample batch from PER (existing)
|
||||
2. Forward pass → save_h_s2
|
||||
3. Selectivity forward: sel = sigmoid(W_sel @ save_h_s2 + b_sel)
|
||||
4. C51/MSE loss + backward → grad_buf
|
||||
5. Compute per-sample gradient magnitude from grad_buf
|
||||
6. Selectivity backward: BCE(sel, normalized_grad_mag) → update W_sel
|
||||
7. Adam step (main params)
|
||||
8. Update PER priorities: priority *= (1 + sel)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Dead Code Cleanup
|
||||
|
||||
Remove before building new system:
|
||||
|
||||
| Item | File | Action |
|
||||
|------|------|--------|
|
||||
| `_reward_norm_kernel` field + load | gpu_experience_collector.rs | Delete field, delete load_function call |
|
||||
| `reward_normalize_kernel` function | nstep_kernel.cu | Keep in .cu (other kernels in same file), delete Rust load |
|
||||
| `cea_weight` parameter | experience_kernels.cu env_step | Remove param, remove from all Rust launch sites |
|
||||
| `mixup_alpha`, `mixup_seed`, `mixup_barrier` params | c51_loss_kernel.cu | Remove params, remove from all Rust launch sites |
|
||||
| `clipped_saxpy_kernel` field | gpu_dqn_trainer.rs | Delete field, delete load |
|
||||
| `apply_cql_clipped_saxpy` method | gpu_dqn_trainer.rs | Delete entire method |
|
||||
| `clip_grad_buf_inplace` method | gpu_dqn_trainer.rs | Delete entire method |
|
||||
| `debug_forward_no_graph` method | gpu_dqn_trainer.rs | Delete entire method |
|
||||
| `GpuTrainingGuard::new()` alias | gpu_training_guard.rs | Replace callers with direct constructor |
|
||||
| Decision transformer `#[allow(dead_code)]` block | decision_transformer.rs | Delete dead block |
|
||||
| GPU curiosity `#[allow(dead_code)]` block | gpu_curiosity_trainer.rs | Delete dead block |
|
||||
|
||||
---
|
||||
|
||||
## Weight Layout
|
||||
|
||||
Existing 26 tensors (indices 0-25) unchanged. New tensors appended:
|
||||
|
||||
| Index | Tensor | Size | Purpose |
|
||||
|-------|--------|------|---------|
|
||||
| 26 | v_select_0 | SH2 | Direction variable selection |
|
||||
| 27 | v_select_1 | SH2 | Magnitude variable selection |
|
||||
| 28 | v_select_2 | SH2 | Order variable selection |
|
||||
| 29 | v_select_3 | SH2 | Urgency variable selection |
|
||||
| 30 | w_gate_0 | AH × SH2 | Direction GLU gate weight |
|
||||
| 31 | b_gate_0 | AH | Direction GLU gate bias |
|
||||
| 32 | w_gate_1 | AH × (SH2+3) | Magnitude GLU gate weight (direction-conditioned) |
|
||||
| 33 | b_gate_1 | AH | Magnitude GLU gate bias |
|
||||
| 34 | w_gate_2 | AH × SH2 | Order GLU gate weight |
|
||||
| 35 | b_gate_2 | AH | Order GLU gate bias |
|
||||
| 36 | w_gate_3 | AH × SH2 | Urgency GLU gate weight |
|
||||
| 37 | b_gate_3 | AH | Urgency GLU gate bias |
|
||||
|
||||
Separate buffers (not in params_buf, own Adam state):
|
||||
- Q attention: 624 floats
|
||||
- Selectivity: 257 floats
|
||||
|
||||
NUM_WEIGHT_TENSORS: 26 → 38
|
||||
|
||||
---
|
||||
|
||||
## CUDA Graph Impact
|
||||
|
||||
- Variable selection softmax: runs BEFORE graph (once per batch, not per-step)
|
||||
- GLU gate GEMM: inside graph (replaces RELU_BIAS fused epilogue)
|
||||
- Cross-branch Q attention: AFTER graph (runs on Q-values, not on the training forward)
|
||||
- Liquid tau: CPU-side only (writes to pinned device-mapped branch_scales)
|
||||
- Selectivity gate: AFTER graph (runs on save_h_s2 from graph forward)
|
||||
|
||||
Only the GLU gate GEMM is inside the CUDA graph. This requires graph recapture (one-time cost). The existing RELU_BIAS epilogue is removed for branch FC — replaced by separate gate GEMM + sigmoid_multiply kernel.
|
||||
|
||||
---
|
||||
|
||||
## Success Criteria
|
||||
|
||||
| Metric | Current (train-w6qfd) | Target | Measurement |
|
||||
|--------|----------------------|--------|-------------|
|
||||
| val_Sharpe sustained | 22.18 frozen at epoch 22 | >25 through epoch 50+ | No freeze |
|
||||
| Q-gap growth | 0.133 locked | >0.2 by epoch 30, still growing | Per-branch Q-gaps |
|
||||
| OOS Sharpe | ~0 oscillating | >0 sustained for 10+ epochs | OOS from trade stats |
|
||||
| Feature selectivity | N/A | Branches show different feature masks | softmax(v_select) divergence |
|
||||
| Branch learning rates | Uniform | Direction converges 2-3× faster than urgency | Per-branch tau values |
|
||||
| Replay selectivity | Uniform PER | Selectivity varies 0.2-0.8 across states | selectivity distribution |
|
||||
|
||||
---
|
||||
|
||||
## Risks and Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|-----------|
|
||||
| GLU doubles branch FC parameters → overfitting | Q-gap momentum + Liquid tau prevent overfitting by modulating gradient intensity |
|
||||
| Cross-branch attention creates circular dependency in action selection | Residual connection + near-zero init ensures Q_coord ≈ Q_raw early → gradual coordination |
|
||||
| Selectivity gate collapse (all 0 or all 1) | BCE loss normalizes targets to ~1.0; sigmoid saturates slowly at ±5 |
|
||||
| Variable selection mask collapse (one feature dominates) | Softmax temperature: mask = softmax(v_select / temperature), temperature=1.0 initially |
|
||||
| Graph recapture from GLU change | One-time cost at epoch 0, well-understood pattern |
|
||||
| Checkpoint incompatibility (38 vs 26 tensors) | New training runs only, no migration from old checkpoints |
|
||||
| Selectivity gate training instability | Separate Adam with low LR (1e-4), gradient clipping to 1.0 |
|
||||
|
||||
---
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. **Dead code cleanup** — remove 11 items, clean build
|
||||
2. **Variable selection + GLU** — branch FC upgrade (Components 1+2, tightly coupled)
|
||||
3. **Cross-branch Q attention** — post-Q coordination (Component 3, independent)
|
||||
4. **Liquid tau** — per-branch dynamics (Component 4, independent)
|
||||
5. **Mamba selective replay** — learned curriculum (Component 5, independent)
|
||||
|
||||
Components 3, 4, 5 are independent of each other and can be implemented in any order after 1+2.
|
||||
Reference in New Issue
Block a user