spec: Supervised Architecture Transfer — 7 concepts from KAN/xLSTM/Diffusion/TGNN/Liquid/TFT/TLOB
1. KAN spline activations → adaptive branch gates (learnable activation shape) 2. xLSTM matrix memory → temporal Q-value context (trajectory pattern detection) 3. Diffusion denoising → iterative Q-refinement (uncertainty-conditioned) 4. TGNN gating → cross-branch graph message passing (structural coordination) 5. Liquid RK4 ODE → higher-order training dynamics (adaptive Euler/RK4) 6. TFT quantile outputs → uncertainty-driven exploration (replaces epsilon-greedy) 7. TLOB MBP-10 → direct microstructure injection to order/urgency branches Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -0,0 +1,357 @@
|
||||
# Supervised Architecture Transfer — Design Spec
|
||||
|
||||
**Goal:** Transfer 7 architectural concepts from the supervised model suite (KAN, xLSTM, Diffusion, TGNN, Liquid ODE, TFT Quantile, TLOB) into the DQN training pipeline to improve action selection quality, representational power, and training dynamics.
|
||||
|
||||
**Problem:** The DQN has fixed sigmoid activations (no adaptive shapes), no temporal memory of Q-value trajectories, single-shot Q-estimation (no refinement), flat Q-value coordination (no structural knowledge of branch dependencies), first-order training dynamics (Euler), mean-only action selection (ignores distributional uncertainty), and uniform feature routing (all branches see same trunk output).
|
||||
|
||||
**Architecture:** 7 independent components that compose through the existing forward/backward/training pipeline. Each borrows a specific concept from a supervised architecture already implemented in `crates/ml-supervised/`.
|
||||
|
||||
---
|
||||
|
||||
## Component 1: KAN Spline Activations → Adaptive Branch Gates
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/kan/layer.rs` — B-spline basis with learnable coefficients + residual connection.
|
||||
|
||||
### Problem
|
||||
The GLU gate uses `sigmoid(pre)` — a fixed S-curve. Every gate neuron has the same activation shape. Sigmoid can't learn sharp thresholds, asymmetric responses, or multi-modal gating patterns.
|
||||
|
||||
### Solution
|
||||
Replace sigmoid in `glu_combine` with a learned 1D B-spline per gate neuron:
|
||||
|
||||
```
|
||||
Current: gate = sigmoid(pre)
|
||||
KAN: gate = clamp(sum_k coeff[k] * B_k(pre) + residual_w * pre, 0, 1)
|
||||
```
|
||||
|
||||
- 8 cubic B-splines on [-4, 4] per neuron (same grid as KAN default)
|
||||
- `coeff[8]` learnable per neuron, initialized to approximate sigmoid
|
||||
- Residual connection: `residual_w * pre` (linear bypass, from KAN's residual_weight)
|
||||
- Clamped to [0, 1] (gate must be a probability)
|
||||
|
||||
### Parameters
|
||||
- Per branch: AH × 8 (spline coeffs) + AH × 1 (residual weight) = AH × 9
|
||||
- Total: 4 × 256 × 9 = 9,216 params
|
||||
- Stored in params_buf (indices 42-49: spline_coeff_0..3, residual_w_0..3)
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `kan_gate_combine`: evaluates B-splines on pre-sigmoid input, combines with coefficients, applies residual, clamps to [0,1], multiplies with value
|
||||
- Replaces `glu_combine` kernel call (same interface: output, gate_pre, value, n)
|
||||
- New CUDA kernel `kan_gate_backward`: gradient through spline basis + residual + clamp
|
||||
- B-spline basis functions are FIXED (not learned) — only coefficients are trained
|
||||
- Spline evaluation: 8 basis functions × 256 neurons = 2048 evaluations per sample per branch. Each is a cubic polynomial (4 multiplies) = 8192 FLOPs per branch. Negligible vs the GEMM cost.
|
||||
|
||||
---
|
||||
|
||||
## Component 2: xLSTM Matrix Memory → Temporal Q-Value Context
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/xlstm/mlstm.rs` — matrix memory C_t with exponential input gate, covariance-based update.
|
||||
|
||||
### Problem
|
||||
Training dynamics (liquid tau, backtracking) react to CURRENT Q-stats only. No memory of Q-value TRAJECTORY patterns. "Q-gap grew steadily for 5 steps then collapsed" is a pattern that predicts the plateau — but without temporal memory, each step is evaluated independently.
|
||||
|
||||
### Solution
|
||||
Lightweight mLSTM cell on CPU that processes Q-stats every 50 training steps:
|
||||
|
||||
```
|
||||
Input: q_input[11] = [avg_max_q, q_min, q_max, q_mean, q_var, entropy, utilization, q_gap_0..3]
|
||||
|
||||
k_t = W_k @ q_input // [head_dim] key
|
||||
v_t = W_v @ q_input // [head_dim] value
|
||||
q_t = W_q @ q_input // [head_dim] query
|
||||
i_t = exp(W_i @ q_input) // exponential input gate (xLSTM novelty)
|
||||
f_t = sigmoid(W_f @ q_input) // forget gate
|
||||
o_t = sigmoid(W_o @ q_input) // output gate
|
||||
|
||||
C_t = f_t * C_{t-1} + i_t * outer(v_t, k_t) // [8, 8] matrix update
|
||||
n_t = f_t * n_{t-1} + i_t * k_t // normalizer
|
||||
h_t = o_t * (C_t @ q_t) / max(|n_t @ q_t|, 1) // normalized retrieval
|
||||
|
||||
context = h_t // [8] context vector
|
||||
```
|
||||
|
||||
### Context usage
|
||||
- `liquid_tau`: context modulates target (1.0 → context-dependent target)
|
||||
- `selectivity`: context concatenated with h_s2 for richer priority prediction
|
||||
- `backtracking`: context divergence (|h_t - h_{t-1}|) as additional plateau signal
|
||||
- `adaptive RK4`: switch to RK4 when context norm exceeds threshold
|
||||
|
||||
### Parameters
|
||||
- 6 weight matrices: W_k, W_v, W_q, W_i, W_f, W_o each [11, 8] = 528 total
|
||||
- C matrix: [8, 8] = 64 floats (persistent state)
|
||||
- n vector: [8] = 8 floats (persistent state)
|
||||
- All CPU-side. Zero GPU overhead.
|
||||
|
||||
### Implementation
|
||||
- `QLSTMCell` struct in `crates/ml/src/trainers/dqn/trainer/mod.rs`
|
||||
- Called from `reduce_current_q_stats` path, after Q-stats are downloaded
|
||||
- Context vector stored as `[f32; 8]` field on DQNTrainer
|
||||
- Weights trained by simple gradient descent on a prediction loss: "predict next step's q_mean from context" (self-supervised)
|
||||
|
||||
---
|
||||
|
||||
## Component 3: Diffusion Denoising → Iterative Q-Refinement
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/diffusion/denoiser.rs` — SiLU-activated DenoiserBlock with time conditioning and residual.
|
||||
|
||||
### Problem
|
||||
Q-values are computed in one forward pass. Noisy states produce noisy Q-values. The C51 distribution captures this uncertainty (high Var[Q]) but the E[Q] readout doesn't use it for refinement.
|
||||
|
||||
### Solution
|
||||
K=3 denoising steps that iteratively refine Q-values conditioned on distributional uncertainty:
|
||||
|
||||
```
|
||||
Q_0 = Q_raw[B, 12]
|
||||
var = Var[Q][B, 12] // already computed (Layer 4)
|
||||
|
||||
For k = 1..3:
|
||||
noise_level = sqrt(var) * (1 - k/3) // decreasing schedule
|
||||
input = [Q_{k-1}; noise_level] // [B, 24]
|
||||
h = SiLU(W1_k @ input + b1_k) // [B, 24] hidden
|
||||
residual = W2_k @ h + b2_k // [B, 12] project back
|
||||
Q_k = Q_{k-1} + 0.1 * residual // small step
|
||||
|
||||
Q_refined = Q_3
|
||||
```
|
||||
|
||||
### Parameters
|
||||
- Per step: W1[24, 24] + b1[24] + W2[12, 24] + b2[12] = 924
|
||||
- 3 steps: 2,772 total
|
||||
- Separate buffer + Adam (like Q-attention)
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `q_denoise_step`: one thread per sample, evaluates the FC-SiLU-FC block
|
||||
- SiLU: `x * sigmoid(x)` — new activation kernel (reuse from diffusion)
|
||||
- Runs after Q-attention + graph message passing, before action selection
|
||||
- 3 sequential kernel launches (steps are dependent)
|
||||
|
||||
---
|
||||
|
||||
## Component 4: TGNN GatingMechanism → Cross-Branch Graph Message Passing
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/tgnn/gating.rs` — `GatingMechanism` with Q/K/V attention + temperature-scaled softmax on graph edges.
|
||||
|
||||
### Problem
|
||||
Cross-branch Q-attention treats 12 Q-values as a flat vector. No structural knowledge that Q[0:3] are direction and Q[3:6] are magnitude. Branch dependencies are domain-specific (direction→magnitude, not urgency→direction).
|
||||
|
||||
### Solution
|
||||
4-node directed graph with learned gated message passing:
|
||||
|
||||
```
|
||||
Edges (domain knowledge):
|
||||
direction → magnitude (position sizing depends on direction)
|
||||
direction → order (order type depends on market direction)
|
||||
magnitude → urgency (position size affects timing pressure)
|
||||
order → urgency (limit vs market affects urgency)
|
||||
|
||||
Per edge (src → dst):
|
||||
gate = sigmoid(W_gate @ [Q_src; Q_dst]) // [1] relevance
|
||||
message = gate * (W_msg @ Q_src) // [3] gated message
|
||||
|
||||
Per node d:
|
||||
Q_d_new = Q_d + mean(incoming messages) // residual update
|
||||
```
|
||||
|
||||
### Parameters
|
||||
- Per edge: W_gate[1, 6] + W_msg[3, 3] = 15
|
||||
- 4 edges: 60 total. Negligible.
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `branch_graph_message_pass`: one thread per sample, 4 hardcoded edges
|
||||
- Runs after Q-attention, before diffusion refinement
|
||||
- In-place update with residual (Q_raw preserved)
|
||||
|
||||
---
|
||||
|
||||
## Component 5: Liquid RK4 ODE Solver → Higher-Order Training Dynamics
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/liquid/ode_solvers.rs` — `RK4Solver` with 4th-order accuracy, `AdaptiveSolver` for regime-dependent switching.
|
||||
|
||||
### Problem
|
||||
Liquid tau uses Euler integration. Euler is first-order — error accumulates linearly over training steps. Per-branch modulation oscillates instead of smoothly converging.
|
||||
|
||||
### Solution
|
||||
Replace Euler with adaptive RK4/Euler switching:
|
||||
|
||||
```
|
||||
// ODE: dx/dt = (1/tau) * (target - x)
|
||||
// f(x) = (1/tau_d) * (target - x)
|
||||
|
||||
if context_norm > rk4_threshold: // xLSTM context detects regime transition
|
||||
// RK4 step
|
||||
k1 = f(x)
|
||||
k2 = f(x + 0.5*k1)
|
||||
k3 = f(x + 0.5*k2)
|
||||
k4 = f(x + k3)
|
||||
x_new = x + (k1 + 2*k2 + 2*k3 + k4) / 6
|
||||
else:
|
||||
// Euler step (fast, accurate enough in stable regime)
|
||||
x_new = x + f(x)
|
||||
```
|
||||
|
||||
### Parameters
|
||||
Zero — algorithm change only. CPU-side, runs every 50 training steps.
|
||||
|
||||
### Implementation
|
||||
- ~20 lines replacing Euler step in `update_liquid_tau`
|
||||
- `rk4_threshold: f32` field (default 0.5, tunable)
|
||||
- xLSTM context norm drives the switching
|
||||
|
||||
---
|
||||
|
||||
## Component 6: TFT Quantile Outputs → Uncertainty-Driven Exploration
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/tft/quantile_outputs.rs` — separate quantile projections with monotonicity constraints.
|
||||
|
||||
### Problem
|
||||
Action selection uses E[Q] (mean of C51 distribution). Throws away distributional shape. An action with E[Q]=0.1, Var[Q]=0.001 (confident) is treated identically to E[Q]=0.1, Var[Q]=1.0 (uncertain). No principled exploration — relies on epsilon-greedy / Boltzmann temperature.
|
||||
|
||||
### Solution
|
||||
Extract quantile Q-values directly from C51 atom CDF (zero extra parameters):
|
||||
|
||||
```
|
||||
// From C51 atoms: p[j] probabilities, z[j] support values
|
||||
// CDF[j] = sum(p[0..j])
|
||||
|
||||
Q_10th[a] = z[j] where CDF[j] first exceeds 0.10 // pessimistic
|
||||
Q_50th[a] = z[j] where CDF[j] first exceeds 0.50 // median
|
||||
Q_90th[a] = z[j] where CDF[j] first exceeds 0.90 // optimistic
|
||||
|
||||
// Adaptive quantile blend using IQN readiness:
|
||||
alpha = iqn_readiness // 0=exploring, 1=exploiting
|
||||
Q_select[a] = (1-alpha) * Q_90th[a] + alpha * Q_10th[a]
|
||||
```
|
||||
|
||||
### Properties
|
||||
- OFU (Optimism in the Face of Uncertainty): uncertain actions explored naturally
|
||||
- No epsilon-greedy needed — exploration is distributional
|
||||
- As training converges (readiness→1), shifts from optimistic to pessimistic (conservative)
|
||||
- Uses existing C51 atoms — zero extra parameters, zero extra forward pass
|
||||
- Replaces the Boltzmann-on-E[Q] action selection
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `quantile_q_select`: reads atom distributions, computes CDF, extracts 3 quantiles, blends
|
||||
- Replaces `compute_expected_q` output for action selection (E[Q] still computed for Q-stats/loss)
|
||||
- Runs in experience collection (per-timestep) — must be efficient
|
||||
- readiness from existing `iqn_readiness` field
|
||||
|
||||
---
|
||||
|
||||
## Component 7: TLOB MBP-10 Features → Direct Microstructure Injection
|
||||
|
||||
### Source
|
||||
`crates/ml-supervised/src/tlob/mbp10_feature_extractor.rs` — 51 LOB features from MBP-10 snapshots.
|
||||
|
||||
### Problem
|
||||
All branches see the same h_s2 trunk output. Order-type decisions depend on microstructure (spread, depth, queue) but these signals are diluted through the shared trunk. Direction decisions don't need microstructure — they need momentum. The trunk can't specialize for both.
|
||||
|
||||
### Solution
|
||||
Route 3 OFI features directly to order branch (d=2) and 3 to urgency branch (d=3), bypassing the trunk:
|
||||
|
||||
```
|
||||
// Already in feature vector (OFI dims 42-49):
|
||||
// order_extra = [bid_ask_spread, depth_imbalance_L1, queue_pressure]
|
||||
// urgency_extra = [spread_velocity, depth_change_rate, trade_arrival_rate]
|
||||
|
||||
// Order branch: h_input = [vsn_masked; ofi_order[3]] → [B, SH2+3]
|
||||
// Urgency branch: h_input = [vsn_masked; ofi_urgency[3]] → [B, SH2+3]
|
||||
```
|
||||
|
||||
### Parameters
|
||||
- `w_b2fc`: `[AH, SH2]` → `[AH, SH2+3]` (+768)
|
||||
- `w_b3fc`: `[AH, SH2]` → `[AH, SH2+3]` (+768)
|
||||
- `w_gate_2`: `[AH, SH2]` → `[AH, SH2+3]` (+768)
|
||||
- `w_gate_3`: `[AH, SH2]` → `[AH, SH2+3]` (+768)
|
||||
- Total: 3,072 extra params
|
||||
|
||||
### Implementation
|
||||
- `concat_ofi_features` CUDA kernel: copies vsn_masked + appends 3 OFI features per sample
|
||||
- Same pattern as Layer 3 mag_concat_qdir but for branches 2 and 3
|
||||
- OFI feature indices are FIXED (hardcoded from the feature vector layout)
|
||||
- Two new concat buffers: `ord_concat_buf[B, SH2+3]` and `urg_concat_buf[B, SH2+3]`
|
||||
- param_sizes indices [16], [20], [38], [40] widened by 3
|
||||
|
||||
---
|
||||
|
||||
## Full Pipeline (after all 7 components)
|
||||
|
||||
```
|
||||
INPUT: states[B, SD]
|
||||
|
||||
── Shared Trunk ──────────────────────────────────────
|
||||
h_s1 = ReLU(W_s1 @ states + b_s1)
|
||||
h_s2 = ReLU(W_s2 @ h_s1 + b_s2)
|
||||
h_s2 = attention(h_s2) [existing]
|
||||
|
||||
── Value Head ────────────────────────────────────────
|
||||
v_logits = W_v2 @ ReLU(W_v1 @ h_s2 + b_v1) + b_v2
|
||||
|
||||
── Branch Heads ──────────────────────────────────────
|
||||
For each branch d:
|
||||
h_masked = h_s2 * sigmoid(h_s2 @ W_vsn1_d @ W_vsn2_d) [CPBI VSN]
|
||||
input_d = concat(h_masked, extra_d) [Layer 3 + Component 7]
|
||||
gate = KAN_spline(W_gate_d @ input_d + b_gate_d) [Component 1]
|
||||
value = W_bdf_d @ input_d + b_bdf_d
|
||||
h_bd = gate * value [GLU with KAN gate]
|
||||
adv_logits_d = W_bdo_d @ h_bd + b_bdo_d
|
||||
|
||||
── Q-Value Pipeline ──────────────────────────────────
|
||||
Q_raw[B, 12] = compute_expected_q(v_logits, adv_logits)
|
||||
Q_coord = cross_branch_q_attention(Q_raw) [CPBI Q-attn]
|
||||
Q_graph = branch_graph_message_pass(Q_coord) [Component 4]
|
||||
Q_refined = diffusion_denoise(Q_graph, Var[Q], K=3) [Component 3]
|
||||
Q_select = quantile_blend(atoms, iqn_readiness) [Component 6]
|
||||
|
||||
── Action Selection ──────────────────────────────────
|
||||
action = boltzmann(Q_select) [quantile-driven]
|
||||
|
||||
── Training Dynamics ─────────────────────────────────
|
||||
q_stats → xLSTM_context[8] [Component 2]
|
||||
liquid_tau(per_branch, context, adaptive_rk4) [Component 5]
|
||||
selectivity(h_s2, context) → PER priorities [CPBI selective]
|
||||
trajectory_backtracking(context_divergence) [CPBI backtrack]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. Component 6 (Quantile Q-select) — zero params, uses existing atoms, highest OOS impact
|
||||
2. Component 4 (Graph message passing) — 60 params, one kernel, structural coordination
|
||||
3. Component 1 (KAN spline gates) — replaces sigmoid in existing GLU, localized change
|
||||
4. Component 3 (Diffusion refinement) — 2,772 params, separate Adam, post-Q-pipeline
|
||||
5. Component 7 (TLOB injection) — widens 4 weight tensors, same pattern as Layer 3
|
||||
6. Component 2 (xLSTM context) — CPU-only, feeds into Components 5 and existing systems
|
||||
7. Component 5 (RK4 ODE) — depends on Component 2, pure algorithm change
|
||||
|
||||
---
|
||||
|
||||
## Success Criteria
|
||||
|
||||
| Metric | Current (CPBI baseline) | Target | Source |
|
||||
|--------|------------------------|--------|--------|
|
||||
| val_Sharpe plateau | Epoch 22 freeze (train-w6qfd) | No freeze through 200 epochs | All components |
|
||||
| OOS Sharpe | ~0 oscillating | >0 sustained 20+ epochs | Components 1, 3, 6 |
|
||||
| Exploration quality | Epsilon-greedy / Boltzmann | Distributional (no epsilon) | Component 6 |
|
||||
| Branch Q-gap variance | Uniform across branches | Direction > urgency | Components 4, 7 |
|
||||
| Training stability | Euler oscillation | RK4 smooth convergence | Components 2, 5 |
|
||||
| Gate activation diversity | Uniform sigmoid | Per-neuron learned shapes | Component 1 |
|
||||
|
||||
---
|
||||
|
||||
## Risks and Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|-----------|
|
||||
| KAN spline overflow | Clamp gate output to [0, 1] |
|
||||
| xLSTM C matrix divergence | Forget gate f_t ∈ (0,1) prevents unbounded growth |
|
||||
| Diffusion refinement distorts Q ordering | Small step (alpha=0.1) + residual preserves ranking |
|
||||
| Graph message passing creates cycles | Directed graph, no back-edges, residual connection |
|
||||
| Quantile blend too optimistic | IQN readiness gates toward pessimistic as training converges |
|
||||
| TLOB feature injection breaks graph capture | Features from existing buffer, no new GPU allocation needed |
|
||||
| Too many components interact | Each component is independently testable; deploy incrementally |
|
||||
Reference in New Issue
Block a user