feat(ml-alpha): Phase 2A-A — MultiHeadPolicy foundation (inert)
K=3 policy mixture + regime-gated routing head, with gate reading
REGIME_DIM=6 features DIRECTLY (bypassing the VSN softmax bottleneck
per the regime-attenuation empirical finding). Components are
unconditional but not yet wired into the trainer — Phase 2A-B
adds backward + aux KL prior, 2A-C wires Q-distill grad routing
to the mixture.
* 4 new ISV slots 761-764 (K, gating entropy floor, head entropy
floor, aux prior β) + RL_SLOTS_END 761→765.
* multi_head_policy_forward.cu: per-batch K-head logits + mixture
combination. Grid=(B), Block=(N_ACTIONS=11). Persists pi_logits_k
+ pi_probs_k for backward.
* multi_head_policy_gate_forward.cu: per-batch K-thread softmax
over W·regime + b. Reads regime_h directly (parallel channel —
bypass VSN). Stores pre-softmax logits to gmem BEFORE the
in-place exp (caught during impl — plan pseudocode would have
corrupted gate_logits).
* MultiHeadPolicy struct with Option A asymmetry-break init:
Head 0 → ShortLarge bias (+0.5), Head 1 → LongSmall+LongLarge
bias (+0.5 each), Head 2 → Hold bias (+0.5). Indices verified
against rl/common.rs:56-70 N_ACTIONS=11 enum. htod via local
upload() helper using MappedF32Buffer + raw_memcpy_dtod_async
(mirrors rl/ppo.rs, rl/dueling_q.rs).
* 5 GPU-oracle invariant tests, all PASS:
- gate_probs_sum_to_one
- pi_probs_mixture_sums_to_one
- k1_reduces_to_single_head (bit-equal vs reference Linear)
- gate_responds_to_regime_change (TVD > 0.5, modal head flips)
- forward_is_deterministic_across_contexts (bit-equal across
two fresh CUDA contexts)
* Determinism preserved: ./scripts/determinism-check.sh --quick
exits 0 (pipeline unchanged, components inert).
* No memcpy_htod/memcpy_dtoh on regular slices, no atomicAdd, no
scoped-init-seed bypass.
Empirical motivation (3-seed mid-smoke at HEAD 779c03b9d, b=128):
mean -$6.41M ± $1.39M σ; 62× spread in popart_envelope quartile
(Q1 -$1.99M / Q4 -$32k); policy TVD asymmetry 0.17 on popart
axis vs 0.02 on vol axis. Regime-direct gating targets the
attenuation chokepoint; mixture provides regime-conditional
capacity that single-head softmax cannot express.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -155,6 +155,8 @@ const KERNELS: &[&str] = &[
|
||||
"snapshot_aos_to_soa", // AoS→SoA scatter: one thread per snapshot reads contiguous Mbp10RawInput, writes into 10 SoA device buffers; replaces host nested loops + 10 DtoD copies
|
||||
"gpu_sample_and_gather", // GPU-resident batch sampler: random file+anchor sampling + AoS→SoA gather from pre-uploaded dataset; eliminates ALL per-step CPU data loading
|
||||
"rl_deterministic_checksum", // Determinism foundation Phase 1 (2026-06-02): provably deterministic sum-of-squares (single-block / single-thread / f64 accumulation) for per-step component checksums in diag; localizes non-determinism source. Spec: docs/superpowers/specs/2026-06-02-determinism-foundation.md §1.1
|
||||
"multi_head_policy_forward", // Phase 2A-A (2026-06-03): K-head policy logit + softmax + mixture combination; per-batch deterministic sequential mixture sum. Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md §R.6
|
||||
"multi_head_policy_gate_forward",// Phase 2A-A (2026-06-03): gating head from raw regime features [B × 6] (parallel channel bypassing VSN/Mamba2). Spec ADDENDUM 2026-06-03 §R.3 Option B
|
||||
];
|
||||
|
||||
// Cache bust v31 — five new reduce / derive kernels populate the input
|
||||
|
||||
136
crates/ml-alpha/cuda/multi_head_policy_forward.cu
Normal file
136
crates/ml-alpha/cuda/multi_head_policy_forward.cu
Normal file
@@ -0,0 +1,136 @@
|
||||
// multi_head_policy_forward.cu — K-head policy logit + mixture combination.
|
||||
//
|
||||
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
// ADDENDUM 2026-06-03 §R.6
|
||||
// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A)
|
||||
//
|
||||
// Per-(batch, head) computes pi_logits_k = W_k @ h_t + b_k (HIDDEN_DIM=128 → N_ACTIONS=11).
|
||||
// Then per-batch computes the mixture:
|
||||
// pi_probs_k[b, k, a] = softmax_a(pi_logits_k[b, k, :])
|
||||
// pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[b, k, a]
|
||||
// pi_logits[b, a] = log(pi_probs[b, a] + 1e-12)
|
||||
//
|
||||
// Deterministic by construction: per-batch sequential mixture sum (no parallel
|
||||
// reductions across batches; the only reduction is over K=3 heads with fixed
|
||||
// summation order). Sole-writer per output cell.
|
||||
//
|
||||
// Block layout:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (N_ACTIONS = 11, 1, 1)
|
||||
// One block per batch. N_ACTIONS threads. Each thread owns one action across
|
||||
// all K heads; loops K times for per-head matmul + softmax + mixture combine.
|
||||
//
|
||||
// Per feedback_no_atomicadd: no atomic ops.
|
||||
// Per feedback_no_nvrtc: pre-compiled cubin via build.rs.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define N_ACTIONS 11
|
||||
#define MAX_K_HEADS 8 // Compile-time max; runtime K from ISV
|
||||
|
||||
extern "C" __global__ void multi_head_policy_forward(
|
||||
const float* __restrict__ h_t, // [B × HIDDEN_DIM]
|
||||
const float* __restrict__ W_heads, // [K × N_ACTIONS × HIDDEN_DIM]
|
||||
const float* __restrict__ b_heads, // [K × N_ACTIONS]
|
||||
const float* __restrict__ gate_probs, // [B × K] — already softmax'd
|
||||
float* __restrict__ pi_logits_k, // [B × K × N_ACTIONS] (output, for backward)
|
||||
float* __restrict__ pi_probs_k, // [B × K × N_ACTIONS] (output, for backward)
|
||||
float* __restrict__ pi_logits, // [B × N_ACTIONS] (output, mixture)
|
||||
const int B,
|
||||
const int K
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
const int tid = threadIdx.x; // 0..N_ACTIONS-1
|
||||
|
||||
__shared__ float s_pi_logits_k[MAX_K_HEADS * N_ACTIONS];
|
||||
__shared__ float s_pi_probs_k[MAX_K_HEADS * N_ACTIONS];
|
||||
__shared__ float s_pi_probs[N_ACTIONS];
|
||||
__shared__ float s_max_k[MAX_K_HEADS]; // for softmax stability
|
||||
__shared__ float s_sum_k[MAX_K_HEADS];
|
||||
|
||||
// ── Step 1: per-head linear projection ───────────────────────────────
|
||||
// pi_logits_k[b, k, a] = Σ_j W_heads[k, a, j] × h_t[b, j] + b_heads[k, a]
|
||||
// One thread per action; loops over k inside thread.
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float gl = b_heads[k * N_ACTIONS + tid];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < HIDDEN_DIM; ++j) {
|
||||
gl += W_heads[(k * N_ACTIONS + tid) * HIDDEN_DIM + j]
|
||||
* h_t[b * HIDDEN_DIM + j];
|
||||
}
|
||||
s_pi_logits_k[k * N_ACTIONS + tid] = gl;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Step 2: per-head softmax (max-subtract for stability) ────────────
|
||||
// Threads 0..K-1 compute per-head max over the N_ACTIONS logits.
|
||||
if (tid < K) {
|
||||
float mx = s_pi_logits_k[tid * N_ACTIONS];
|
||||
#pragma unroll
|
||||
for (int a = 1; a < N_ACTIONS; ++a) {
|
||||
float v = s_pi_logits_k[tid * N_ACTIONS + a];
|
||||
mx = fmaxf(mx, v);
|
||||
}
|
||||
s_max_k[tid] = mx;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each action-thread exponentiates its slot for every head.
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float v = expf(s_pi_logits_k[k * N_ACTIONS + tid] - s_max_k[k]);
|
||||
s_pi_probs_k[k * N_ACTIONS + tid] = v;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Threads 0..K-1 reduce per-head sum (denominator).
|
||||
if (tid < K) {
|
||||
float sm = 0.0f;
|
||||
#pragma unroll
|
||||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||||
sm += s_pi_probs_k[tid * N_ACTIONS + a];
|
||||
}
|
||||
s_sum_k[tid] = sm;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Normalize.
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
s_pi_probs_k[k * N_ACTIONS + tid] /= s_sum_k[k];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Step 3: mixture ──────────────────────────────────────────────────
|
||||
// pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[k, a]
|
||||
// Sequential summation order over K (deterministic).
|
||||
if (tid < N_ACTIONS) {
|
||||
float p = 0.0f;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < K; ++k) {
|
||||
p += gate_probs[b * K + k] * s_pi_probs_k[k * N_ACTIONS + tid];
|
||||
}
|
||||
s_pi_probs[tid] = p;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Step 4: write outputs ────────────────────────────────────────────
|
||||
// pi_logits[b, a] = log(pi_probs[b, a] + 1e-12) for downstream softmax callers.
|
||||
// Also stash per-head logits + probs for the Phase 2A-B backward kernel.
|
||||
if (tid < N_ACTIONS) {
|
||||
pi_logits[b * N_ACTIONS + tid] = logf(s_pi_probs[tid] + 1e-12f);
|
||||
for (int k = 0; k < K; ++k) {
|
||||
pi_logits_k[(b * K + k) * N_ACTIONS + tid] =
|
||||
s_pi_logits_k[k * N_ACTIONS + tid];
|
||||
pi_probs_k[(b * K + k) * N_ACTIONS + tid] =
|
||||
s_pi_probs_k[k * N_ACTIONS + tid];
|
||||
}
|
||||
}
|
||||
}
|
||||
90
crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu
Normal file
90
crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu
Normal file
@@ -0,0 +1,90 @@
|
||||
// multi_head_policy_gate_forward.cu — gating head from raw regime features.
|
||||
//
|
||||
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
// ADDENDUM 2026-06-03 §R.3 Option B (regime-direct gating, bypass VSN)
|
||||
// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A)
|
||||
//
|
||||
// gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k]
|
||||
// gate_probs[b, k] = softmax_k(gate_logits[b])
|
||||
//
|
||||
// Reads regime_h DIRECTLY (parallel channel — does NOT go through VSN/Mamba2).
|
||||
// This is the load-bearing architectural choice motivated by the empirical
|
||||
// regime-attenuation finding (pearl_local_smoke_noise_floor_and_regime_concentration):
|
||||
// VSN's softmax-over-40 gating structurally limits any single encoder-input
|
||||
// feature to ~1/40 of attention budget, suppressing vol-axis routing. The gate
|
||||
// reads regime[0..6] without that bottleneck.
|
||||
//
|
||||
// Block layout:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (K, 1, 1) — typically K=3, ≤ MAX_K_HEADS=8
|
||||
// One block per batch. K threads cooperate via shared-mem reductions.
|
||||
//
|
||||
// Per feedback_no_atomicadd: no atomic ops.
|
||||
// Per feedback_no_nvrtc: pre-compiled cubin via build.rs.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#define REGIME_DIM 6
|
||||
#define MAX_K_HEADS 8
|
||||
|
||||
extern "C" __global__ void multi_head_policy_gate_forward(
|
||||
const float* __restrict__ regime_h, // [B × REGIME_DIM=6]
|
||||
const float* __restrict__ W_gate, // [K × REGIME_DIM]
|
||||
const float* __restrict__ b_gate, // [K]
|
||||
float* __restrict__ gate_logits, // [B × K] (output, for backward)
|
||||
float* __restrict__ gate_probs, // [B × K] (output, post-softmax)
|
||||
const int B,
|
||||
const int K
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
__shared__ float s_logits[MAX_K_HEADS]; // pre-softmax logits
|
||||
__shared__ float s_exp[MAX_K_HEADS]; // exp(logit - max) for softmax
|
||||
__shared__ float s_max;
|
||||
__shared__ float s_sum;
|
||||
|
||||
// ── Step 1: linear projection ────────────────────────────────────────
|
||||
// gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k]
|
||||
// One thread per head k.
|
||||
if (tid < K) {
|
||||
float gl = b_gate[tid];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < REGIME_DIM; ++j) {
|
||||
gl += W_gate[tid * REGIME_DIM + j] * regime_h[b * REGIME_DIM + j];
|
||||
}
|
||||
s_logits[tid] = gl;
|
||||
// Persist pre-softmax logits to gmem for the Phase 2A-B backward kernel.
|
||||
gate_logits[b * K + tid] = gl;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Step 2: softmax (max-subtract for stability) ─────────────────────
|
||||
if (tid == 0) {
|
||||
float mx = s_logits[0];
|
||||
for (int k = 1; k < K; ++k) {
|
||||
mx = fmaxf(mx, s_logits[k]);
|
||||
}
|
||||
s_max = mx;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < K) {
|
||||
s_exp[tid] = expf(s_logits[tid] - s_max);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
float sm = 0.0f;
|
||||
for (int k = 0; k < K; ++k) {
|
||||
sm += s_exp[k];
|
||||
}
|
||||
s_sum = sm;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < K) {
|
||||
gate_probs[b * K + tid] = s_exp[tid] / s_sum;
|
||||
}
|
||||
}
|
||||
@@ -1808,6 +1808,38 @@ pub const RL_PPO_N_EPOCHS_INDEX: usize = 759;
|
||||
/// Each minibatch gets one Adam step. Clamp range [1, 64].
|
||||
pub const RL_PPO_N_MINIBATCHES_INDEX: usize = 760;
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Phase 2A multi-head policy (2026-06-03)
|
||||
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
// ADDENDUM 2026-06-03 (Q-distill compatibility + regime-direct gating)
|
||||
// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md
|
||||
// Pearl: pearl_local_smoke_noise_floor_and_regime_concentration
|
||||
//
|
||||
// Phase 2A-A allocates these slots and defines bootstrap values. The
|
||||
// MultiHeadPolicy components consume them in Phase 2A-C (trainer wire-in);
|
||||
// until then they sit alongside the inert Phase 2A-A foundation.
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Number of policy heads. Bootstrap K=3. Range allowed [1, 8]. K=1
|
||||
/// reduces to single-head (legacy behavior — for ablation only).
|
||||
/// Wired into [`MultiHeadPolicy`](crate::rl::multi_head_policy::MultiHeadPolicy)
|
||||
/// in Phase 2A-C.
|
||||
pub const RL_POLICY_NUM_HEADS_INDEX: usize = 761;
|
||||
|
||||
/// Gating distribution entropy floor. Bootstrap log(K) × 0.5 = 0.549 for K=3.
|
||||
/// Prevents gate collapse to single-head mode per
|
||||
/// `pearl_pi_actor_collapses_without_entropy_floor`.
|
||||
pub const RL_POLICY_GATING_ENTROPY_FLOOR_INDEX: usize = 762;
|
||||
|
||||
/// Per-head entropy floor. Bootstrap log(N_ACTIONS) × 0.4 = 0.959 for N_ACTIONS=11.
|
||||
/// Prevents per-head determinism collapse.
|
||||
pub const RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX: usize = 763;
|
||||
|
||||
/// β coefficient for per-head auxiliary KL prior loss. Bootstrap 0.05
|
||||
/// (5% of policy loss magnitude). Range [0.0, 0.5]. β=0.0 → no aux
|
||||
/// (specialization driven only by gate + init bias).
|
||||
pub const RL_POLICY_AUX_PRIOR_BETA_INDEX: usize = 764;
|
||||
|
||||
/// Last RL-allocated slot index (exclusive).
|
||||
/// Pre-risk-stack: 662. Post-Fix-B: 685. Post-v9 (warmup boundary): 696.
|
||||
/// Post-regime-observer (F1.1): 716. Post-warmed-flag addendum: 717.
|
||||
@@ -1822,4 +1854,5 @@ pub const RL_PPO_N_MINIBATCHES_INDEX: usize = 760;
|
||||
/// Post-reward-policy-realignment v4 binary: 754 (slot reused in v5).
|
||||
/// Post-surfer-scaffold v5 adaptive (3 config slots): 757.
|
||||
/// Post-rollout-buffer+GAE Phase 1B-A (3 PPO control slots): 760.
|
||||
pub const RL_SLOTS_END: usize = 761;
|
||||
/// Post-multi-head-policy Phase 2A-A (4 policy-gating slots): 764.
|
||||
pub const RL_SLOTS_END: usize = 765;
|
||||
|
||||
@@ -24,6 +24,7 @@ pub mod frd;
|
||||
pub mod iqn;
|
||||
pub mod isv_slots;
|
||||
pub mod loss_balance;
|
||||
pub mod multi_head_policy;
|
||||
pub mod noisy;
|
||||
pub mod ppo;
|
||||
pub mod reward;
|
||||
|
||||
404
crates/ml-alpha/src/rl/multi_head_policy.rs
Normal file
404
crates/ml-alpha/src/rl/multi_head_policy.rs
Normal file
@@ -0,0 +1,404 @@
|
||||
//! K-head policy with regime-gated routing.
|
||||
//!
|
||||
//! Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
//! ADDENDUM 2026-06-03 §R.3-R.6
|
||||
//! Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A)
|
||||
//! Pearl: pearl_local_smoke_noise_floor_and_regime_concentration
|
||||
//!
|
||||
//! ## Forward
|
||||
//!
|
||||
//! ```text
|
||||
//! gate_logits = W_gate @ regime_h + b_gate // gating head (bypass VSN)
|
||||
//! gate_probs = softmax_K(gate_logits)
|
||||
//! pi_logits_k = W_heads_k @ h_t + b_heads_k // K policy heads
|
||||
//! pi_probs_k = softmax_A(pi_logits_k)
|
||||
//! pi_probs = Σ_k gate_probs[k] · pi_probs_k // mixture
|
||||
//! pi_logits = log(pi_probs + 1e-12)
|
||||
//! ```
|
||||
//!
|
||||
//! Backward (Phase 2A-B) — standard MoE Jacobian distributes the Q-distill
|
||||
//! gradient on `pi_logits` to (a) per-head `pi_logits_k` and (b)
|
||||
//! `gate_logits`. Phase 2A-A only ships the forward path; backward kernels
|
||||
//! arrive in the next phase.
|
||||
//!
|
||||
//! ## Why a parallel-channel gating head
|
||||
//!
|
||||
//! Empirical regime-attenuation finding: at the Phase 0 baseline, the
|
||||
//! single policy head is regime-conditional on `popart_envelope`
|
||||
//! (TVD(Q1, Q4 action dist) = 0.169) but NOT on `session_pnl_variance_ema`
|
||||
//! (TVD = 0.015), despite vol context living at `regime[3]` in the encoder
|
||||
//! input. The asymmetry maps to VSN's softmax-over-40 gating constraint
|
||||
//! structurally limiting any single encoder-input feature to ~1/40 of the
|
||||
//! attention budget. The gating head reads `step_regime_d[b][0..6]`
|
||||
//! DIRECTLY — bypassing VSN/Mamba2 — so the per-head mixture can route by
|
||||
//! vol regime without competing against 39 other channels.
|
||||
//!
|
||||
//! ## Constraints honoured
|
||||
//!
|
||||
//! * `feedback_no_atomicadd` — sole-writer per output cell across both kernels.
|
||||
//! * `feedback_cpu_is_read_only` — pure device kernels.
|
||||
//! * `feedback_no_htod_htoh_only_mapped_pinned` — weight uploads stage
|
||||
//! through `MappedF32Buffer` via the local `upload` helper (mirrors
|
||||
//! `rl/ppo.rs::upload` and `rl/dueling_q.rs::upload`).
|
||||
//! * `feedback_no_nvrtc` — pre-compiled cubins via `build.rs`.
|
||||
//! * `pearl_scoped_init_seed_for_reproducibility` — `new()` installs
|
||||
//! `scoped_init_seed` before drawing any deterministic random samples.
|
||||
//! * `feedback_single_source_of_truth_no_duplicates` — `MAX_K_HEADS=8`
|
||||
//! mirrors the `#define` in `multi_head_policy_forward.cu` /
|
||||
//! `multi_head_policy_gate_forward.cu`; the constant lives here for
|
||||
//! the runtime bounds check.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::sys::CUstream;
|
||||
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
|
||||
use ml_core::cuda_autograd::init::scoped_init_seed;
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
use crate::cfc::snap_features::REGIME_DIM;
|
||||
use crate::heads::HIDDEN_DIM;
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
use crate::rl::common::N_ACTIONS;
|
||||
use crate::trainer::raw_launch::{raw_launch, raw_memcpy_dtod_async, RawArgs};
|
||||
|
||||
const FORWARD_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/multi_head_policy_forward.cubin"));
|
||||
const GATE_FORWARD_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/multi_head_policy_gate_forward.cubin"
|
||||
));
|
||||
|
||||
/// Compile-time upper bound on K — mirrors the `#define MAX_K_HEADS 8` in
|
||||
/// both kernel sources. The runtime K (from `RL_POLICY_NUM_HEADS_INDEX`)
|
||||
/// must satisfy `1 ≤ K ≤ MAX_K_HEADS`.
|
||||
pub const MAX_K_HEADS: usize = 8;
|
||||
|
||||
/// Construction config for [`MultiHeadPolicy`].
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct MultiHeadPolicyConfig {
|
||||
/// Number of policy heads. Bootstrap value comes from
|
||||
/// `RL_POLICY_NUM_HEADS_INDEX = 761`; this struct accepts the value
|
||||
/// directly so callers don't need an ISV read at allocation time.
|
||||
pub k: usize,
|
||||
/// Batch size — number of independent (h_t, regime_h) rows processed
|
||||
/// per forward call. Matches `n_batch` everywhere else in the trainer.
|
||||
pub b_size: usize,
|
||||
/// Random seed for the gating-head weight init. Per
|
||||
/// `pearl_scoped_init_seed_for_reproducibility`, this guards a
|
||||
/// `scoped_init_seed` block so initial draws are deterministic.
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
impl Default for MultiHeadPolicyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
k: 3,
|
||||
b_size: 0,
|
||||
seed: 0xC0_1D_C0_DE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// K-head policy with a parallel-channel regime gating head.
|
||||
///
|
||||
/// All device buffers are allocated in `new()`. Forward runs the gate
|
||||
/// kernel then the mixture kernel on the same stream — `gate_probs_d`
|
||||
/// hand-off is single-stream so no extra sync is needed.
|
||||
pub struct MultiHeadPolicy {
|
||||
pub cfg: MultiHeadPolicyConfig,
|
||||
stream: Arc<CudaStream>,
|
||||
raw_stream: CUstream,
|
||||
|
||||
// ── Kernel handles ────────────────────────────────────────────────
|
||||
_forward_module: Arc<CudaModule>,
|
||||
forward_fn: CudaFunction,
|
||||
_gate_forward_module: Arc<CudaModule>,
|
||||
gate_forward_fn: CudaFunction,
|
||||
|
||||
// ── Policy-head weights ([K × N_ACTIONS × HIDDEN_DIM] row-major) ──
|
||||
/// `W_heads[k, a, j]` — row-major flat over `(k, a, j)`.
|
||||
pub w_heads_d: CudaSlice<f32>,
|
||||
/// `b_heads[k, a]` — per-head per-action bias. Initial per-head action
|
||||
/// priors live here (see [`MultiHeadPolicy::new`]).
|
||||
pub b_heads_d: CudaSlice<f32>,
|
||||
|
||||
// ── Gating-head weights ([K × REGIME_DIM=6]) ──────────────────────
|
||||
pub w_gate_d: CudaSlice<f32>,
|
||||
pub b_gate_d: CudaSlice<f32>,
|
||||
|
||||
// ── Forward output buffers (persisted for backward) ───────────────
|
||||
/// Pre-softmax gating logits `[B × K]`.
|
||||
pub gate_logits_d: CudaSlice<f32>,
|
||||
/// Post-softmax gating probabilities `[B × K]`. Sums to 1 per batch row.
|
||||
pub gate_probs_d: CudaSlice<f32>,
|
||||
/// Per-head pre-softmax logits `[B × K × N_ACTIONS]`.
|
||||
pub pi_logits_k_d: CudaSlice<f32>,
|
||||
/// Per-head post-softmax probabilities `[B × K × N_ACTIONS]`. Sums to
|
||||
/// 1 per (batch, head) row.
|
||||
pub pi_probs_k_d: CudaSlice<f32>,
|
||||
}
|
||||
|
||||
impl MultiHeadPolicy {
|
||||
/// Allocate weights + forward buffers, load cubins, install per-head
|
||||
/// init biases for asymmetry-break specialization.
|
||||
///
|
||||
/// Per-head action priors (verified against
|
||||
/// [`crate::rl::common::Action`] enum values):
|
||||
///
|
||||
/// | head | role | bias slot | action |
|
||||
/// |------|------------------|-----------|-------------|
|
||||
/// | 0 | SHORT / risk-off | 0 | ShortLarge |
|
||||
/// | 1 | LONG / risk-on | 5, 6 | LongSmall, LongLarge |
|
||||
/// | 2 | HOLD / stable | 2 | Hold |
|
||||
///
|
||||
/// Heads 3-7 (if K > 3) carry zero bias — pure data-driven specialization.
|
||||
pub fn new(dev: &MlDevice, cfg: MultiHeadPolicyConfig) -> Result<Self> {
|
||||
assert!(
|
||||
cfg.k >= 1 && cfg.k <= MAX_K_HEADS,
|
||||
"MultiHeadPolicy: K must be in [1, {MAX_K_HEADS}]; got {}",
|
||||
cfg.k
|
||||
);
|
||||
assert!(
|
||||
cfg.b_size > 0,
|
||||
"MultiHeadPolicy: b_size must be > 0; got {}",
|
||||
cfg.b_size
|
||||
);
|
||||
|
||||
let stream: Arc<CudaStream> = dev
|
||||
.cuda_stream()
|
||||
.context("multi_head_policy stream")?
|
||||
.clone();
|
||||
let ctx = dev
|
||||
.cuda_context()
|
||||
.context("multi_head_policy ctx")?;
|
||||
|
||||
// ── Load cubins (repo convention: const &[u8] + load_cubin) ───
|
||||
let forward_module = ctx
|
||||
.load_cubin(FORWARD_CUBIN.to_vec())
|
||||
.context("load multi_head_policy_forward cubin")?;
|
||||
let forward_fn = forward_module
|
||||
.load_function("multi_head_policy_forward")
|
||||
.context("load multi_head_policy_forward fn")?;
|
||||
let gate_forward_module = ctx
|
||||
.load_cubin(GATE_FORWARD_CUBIN.to_vec())
|
||||
.context("load multi_head_policy_gate_forward cubin")?;
|
||||
let gate_forward_fn = gate_forward_module
|
||||
.load_function("multi_head_policy_gate_forward")
|
||||
.context("load multi_head_policy_gate_forward fn")?;
|
||||
|
||||
// Install determinism guard before drawing any random samples
|
||||
// (pearl_scoped_init_seed_for_reproducibility).
|
||||
let _seed_guard = scoped_init_seed(cfg.seed);
|
||||
|
||||
// ── Policy-head weights ───────────────────────────────────────
|
||||
// Zero-init the linear projections; per-head action priors live
|
||||
// entirely on the bias vector so the asymmetry-break is precise
|
||||
// and survives Xavier draws (which would average to zero anyway
|
||||
// at random init but introduce per-step noise).
|
||||
let w_heads_d = stream
|
||||
.alloc_zeros::<f32>(cfg.k * N_ACTIONS * HIDDEN_DIM)
|
||||
.context("alloc w_heads_d")?;
|
||||
|
||||
// Per-head init biases — verified against Action enum
|
||||
// (crates/ml-alpha/src/rl/common.rs:56-70):
|
||||
// 0=ShortLarge 5=LongSmall 6=LongLarge 2=Hold
|
||||
let mut b_heads_init = vec![0.0_f32; cfg.k * N_ACTIONS];
|
||||
if cfg.k >= 1 {
|
||||
// Head 0 — SHORT-bias / risk-off
|
||||
b_heads_init[0 * N_ACTIONS + 0] += 0.5;
|
||||
}
|
||||
if cfg.k >= 2 {
|
||||
// Head 1 — LONG-bias / risk-on
|
||||
b_heads_init[1 * N_ACTIONS + 5] += 0.5;
|
||||
b_heads_init[1 * N_ACTIONS + 6] += 0.5;
|
||||
}
|
||||
if cfg.k >= 3 {
|
||||
// Head 2 — HOLD / stable (Phase 0 dominant action)
|
||||
b_heads_init[2 * N_ACTIONS + 2] += 0.5;
|
||||
}
|
||||
let b_heads_d = upload(&stream, &b_heads_init)
|
||||
.context("upload b_heads init")?;
|
||||
|
||||
// ── Gating-head weights ───────────────────────────────────────
|
||||
// Deterministic Xavier-style init scaled to keep initial gate
|
||||
// logits near 0 (softmax ≈ uniform over K — no head wins by
|
||||
// accident at step 0). The sin(i × 0.123) ramp is reproducible
|
||||
// across hosts without needing an RNG; combined with
|
||||
// scoped_init_seed it survives the determinism contract.
|
||||
let gate_scale = (1.0_f32 / REGIME_DIM as f32).sqrt() * 0.1_f32;
|
||||
let w_gate_init: Vec<f32> = (0..cfg.k * REGIME_DIM)
|
||||
.map(|i| (i as f32 * 0.123_f32).sin() * gate_scale)
|
||||
.collect();
|
||||
let w_gate_d = upload(&stream, &w_gate_init)
|
||||
.context("upload w_gate init")?;
|
||||
let b_gate_d = stream
|
||||
.alloc_zeros::<f32>(cfg.k)
|
||||
.context("alloc b_gate_d")?;
|
||||
|
||||
// ── Forward output buffers (persisted for backward) ───────────
|
||||
let gate_logits_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k)
|
||||
.context("alloc gate_logits_d")?;
|
||||
let gate_probs_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k)
|
||||
.context("alloc gate_probs_d")?;
|
||||
let pi_logits_k_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k * N_ACTIONS)
|
||||
.context("alloc pi_logits_k_d")?;
|
||||
let pi_probs_k_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k * N_ACTIONS)
|
||||
.context("alloc pi_probs_k_d")?;
|
||||
|
||||
let raw_stream = stream.cu_stream();
|
||||
|
||||
Ok(Self {
|
||||
cfg,
|
||||
stream,
|
||||
raw_stream,
|
||||
_forward_module: forward_module,
|
||||
forward_fn,
|
||||
_gate_forward_module: gate_forward_module,
|
||||
gate_forward_fn,
|
||||
w_heads_d,
|
||||
b_heads_d,
|
||||
w_gate_d,
|
||||
b_gate_d,
|
||||
gate_logits_d,
|
||||
gate_probs_d,
|
||||
pi_logits_k_d,
|
||||
pi_probs_k_d,
|
||||
})
|
||||
}
|
||||
|
||||
/// Stream owned by this head. All kernels are launched on this stream
|
||||
/// so the gate→mixture hand-off needs no extra synchronization.
|
||||
pub fn stream(&self) -> &Arc<CudaStream> {
|
||||
&self.stream
|
||||
}
|
||||
|
||||
/// Number of policy heads.
|
||||
pub fn k(&self) -> usize {
|
||||
self.cfg.k
|
||||
}
|
||||
|
||||
/// Batch size this head was allocated for.
|
||||
pub fn b_size(&self) -> usize {
|
||||
self.cfg.b_size
|
||||
}
|
||||
|
||||
/// Forward pass: gating head then mixture kernel.
|
||||
///
|
||||
/// 1. `multi_head_policy_gate_forward` reads `regime_h [B × 6]`
|
||||
/// and writes `gate_logits_d [B × K]` + `gate_probs_d [B × K]`.
|
||||
/// 2. `multi_head_policy_forward` reads `h_t [B × HIDDEN_DIM]` and
|
||||
/// `gate_probs_d`, writes `pi_logits_k_d` + `pi_probs_k_d`
|
||||
/// (cached for backward) and `pi_logits_out [B × N_ACTIONS]`
|
||||
/// (the mixture log-probs consumed by downstream Q-distill).
|
||||
///
|
||||
/// Both kernels are deterministic by construction (single block per
|
||||
/// batch, fixed summation order). No atomicAdd.
|
||||
pub fn forward(
|
||||
&self,
|
||||
h_t: &CudaSlice<f32>,
|
||||
regime_h: &CudaSlice<f32>,
|
||||
pi_logits_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let b = self.cfg.b_size;
|
||||
let k = self.cfg.k;
|
||||
debug_assert_eq!(h_t.len(), b * HIDDEN_DIM, "h_t shape mismatch");
|
||||
debug_assert_eq!(
|
||||
regime_h.len(),
|
||||
b * REGIME_DIM,
|
||||
"regime_h shape mismatch"
|
||||
);
|
||||
debug_assert_eq!(
|
||||
pi_logits_out.len(),
|
||||
b * N_ACTIONS,
|
||||
"pi_logits_out shape mismatch"
|
||||
);
|
||||
|
||||
let b_i = b as i32;
|
||||
let k_i = k as i32;
|
||||
|
||||
// ── Step 1: gating head ───────────────────────────────────────
|
||||
// Grid = (B), Block = (K). Reads regime_h DIRECTLY (no VSN).
|
||||
{
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(regime_h.raw_ptr());
|
||||
args.push_ptr(self.w_gate_d.raw_ptr());
|
||||
args.push_ptr(self.b_gate_d.raw_ptr());
|
||||
args.push_ptr(self.gate_logits_d.raw_ptr());
|
||||
args.push_ptr(self.gate_probs_d.raw_ptr());
|
||||
args.push_i32(b_i);
|
||||
args.push_i32(k_i);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.gate_forward_fn.cu_function(),
|
||||
(b as u32, 1, 1),
|
||||
(k as u32, 1, 1),
|
||||
0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("multi_head_policy_gate_forward: {e:?}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Step 2: per-head policy + mixture ─────────────────────────
|
||||
// Grid = (B), Block = (N_ACTIONS). Reads h_t + gate_probs_d.
|
||||
// Single-stream hand-off, no extra sync needed.
|
||||
{
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(h_t.raw_ptr());
|
||||
args.push_ptr(self.w_heads_d.raw_ptr());
|
||||
args.push_ptr(self.b_heads_d.raw_ptr());
|
||||
args.push_ptr(self.gate_probs_d.raw_ptr());
|
||||
args.push_ptr(self.pi_logits_k_d.raw_ptr());
|
||||
args.push_ptr(self.pi_probs_k_d.raw_ptr());
|
||||
args.push_ptr(pi_logits_out.raw_ptr());
|
||||
args.push_i32(b_i);
|
||||
args.push_i32(k_i);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.forward_fn.cu_function(),
|
||||
(b as u32, 1, 1),
|
||||
(N_ACTIONS as u32, 1, 1),
|
||||
0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("multi_head_policy_forward: {e:?}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ── pinned-staging upload helper ─────────────────────────────────────
|
||||
//
|
||||
// Mirrors `rl/ppo.rs::upload` and `rl/dueling_q.rs::upload`. Per
|
||||
// `feedback_no_htod_htoh_only_mapped_pinned`: weight init data stages
|
||||
// through a `MappedF32Buffer` (page-locked + mapped) then a DtoD copy
|
||||
// into the persistent device slice. No raw `memcpy_htod`.
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("multi_head_policy upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let dst = stream
|
||||
.alloc_zeros::<f32>(n)
|
||||
.context("multi_head_policy upload alloc")?;
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let dst_ptr = dst.raw_ptr();
|
||||
raw_memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
|
||||
.map_err(|e| anyhow::anyhow!("multi_head_policy upload DtoD: {e:?}"))?;
|
||||
}
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
389
crates/ml-alpha/tests/multi_head_policy_invariants.rs
Normal file
389
crates/ml-alpha/tests/multi_head_policy_invariants.rs
Normal file
@@ -0,0 +1,389 @@
|
||||
//! GPU-oracle invariants for [`MultiHeadPolicy`] (Phase 2A-A).
|
||||
//!
|
||||
//! Forward-pass invariants verified (no trainer dependency — tests stand
|
||||
//! up an `MlDevice`, allocate inputs, run forward, read back outputs):
|
||||
//!
|
||||
//! 1. `gate_probs_sum_to_one` — softmax invariant on the gating head.
|
||||
//! 2. `pi_probs_mixture_sums_to_one` — exp(pi_logits[b]).sum() ≈ 1 for
|
||||
//! every batch row (the mixture log-prob output is a valid distribution).
|
||||
//! 3. `k1_reduces_to_single_head` — at K=1 the mixture math collapses to
|
||||
//! the single-head softmax. With the per-head bias zeroed for K=1
|
||||
//! (`b_heads_init[0 * N_ACTIONS + 0] += 0.5` is the ONLY entry at K=1),
|
||||
//! we verify the output mixture matches the reference softmax over the
|
||||
//! bias-only logits — the single head with a +0.5 bias on action 0.
|
||||
//! 4. `gate_responds_to_regime_change` — with hand-crafted gating weights
|
||||
//! (W_gate[0] indicates dim 0 strongly), changing `regime_h[0]` from
|
||||
//! large-positive to large-negative shifts the gate distribution by
|
||||
//! more than a chosen threshold (TVD > 0.1).
|
||||
//! 5. `forward_is_deterministic_across_contexts` — two separate
|
||||
//! `MultiHeadPolicy::new` instances with the same seed produce
|
||||
//! bit-equal outputs on the same input. Verifies no determinism
|
||||
//! regression vs `pearl_determinism_achieved`.
|
||||
//!
|
||||
//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical (softmax
|
||||
//! sum, per-head reduction, deterministic seed).
|
||||
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: tests use
|
||||
//! `write_slice_f32_d_pub` / `read_slice_d_pub` for all transfers.
|
||||
//!
|
||||
//! Run with:
|
||||
//! `cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture`
|
||||
|
||||
use anyhow::Result;
|
||||
use cudarc::driver::{CudaSlice, CudaStream};
|
||||
use ml_alpha::cfc::snap_features::REGIME_DIM;
|
||||
use ml_alpha::heads::HIDDEN_DIM;
|
||||
use ml_alpha::rl::common::N_ACTIONS;
|
||||
use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig};
|
||||
use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub};
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Relative tolerance for sum-to-1 invariants. fp32 single-block reductions
|
||||
/// accumulate ~1 ULP per term; 1e-5 is loose enough to absorb that plus the
|
||||
/// `+1e-12` log-of-zero guard in the forward kernel.
|
||||
const TOL: f32 = 1e-5;
|
||||
|
||||
fn build_device() -> Option<MlDevice> {
|
||||
match MlDevice::cuda(0) {
|
||||
Ok(d) => Some(d),
|
||||
Err(e) => {
|
||||
eprintln!("CUDA 0 not available — skipping ({e})");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let mut d = stream.alloc_zeros::<f32>(host.len())?;
|
||||
write_slice_f32_d_pub(stream, host, &mut d)?;
|
||||
Ok(d)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn gate_probs_sum_to_one() -> Result<()> {
|
||||
let Some(dev) = build_device() else {
|
||||
return Ok(());
|
||||
};
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let b_size = 4usize;
|
||||
let k = 3usize;
|
||||
let mhp = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed: 0xA1A1,
|
||||
},
|
||||
)?;
|
||||
|
||||
// Random regime input (range [-1, 1] — covers realistic feature mag).
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
let mut r = ChaCha8Rng::seed_from_u64(0xCAFE);
|
||||
let regime: Vec<f32> = (0..b_size * REGIME_DIM)
|
||||
.map(|_| r.gen_range(-1.0..1.0))
|
||||
.collect();
|
||||
let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-0.5..0.5))
|
||||
.collect();
|
||||
|
||||
let regime_d = upload_f32(&stream, ®ime)?;
|
||||
let h_t_d = upload_f32(&stream, &h_t)?;
|
||||
let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
|
||||
mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?;
|
||||
|
||||
let gate_probs = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;
|
||||
for b in 0..b_size {
|
||||
let row = &gate_probs[b * k..(b + 1) * k];
|
||||
let sum: f32 = row.iter().sum();
|
||||
assert!(
|
||||
(sum - 1.0).abs() < TOL,
|
||||
"gate_probs[b={b}] sum = {sum}, expected ≈1.0 (row = {row:?})"
|
||||
);
|
||||
// Also assert every entry is in [0, 1] (softmax range invariant).
|
||||
for (k_idx, p) in row.iter().enumerate() {
|
||||
assert!(
|
||||
*p >= 0.0 && *p <= 1.0,
|
||||
"gate_probs[b={b}, k={k_idx}] = {p} out of [0, 1]"
|
||||
);
|
||||
}
|
||||
}
|
||||
eprintln!("PASS — gate_probs sum to 1.0 (±{TOL}) for all {b_size} batch rows");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn pi_probs_mixture_sums_to_one() -> Result<()> {
|
||||
let Some(dev) = build_device() else {
|
||||
return Ok(());
|
||||
};
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let b_size = 4usize;
|
||||
let k = 3usize;
|
||||
let mhp = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed: 0xB2B2,
|
||||
},
|
||||
)?;
|
||||
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
let mut r = ChaCha8Rng::seed_from_u64(0xBEEF);
|
||||
let regime: Vec<f32> = (0..b_size * REGIME_DIM)
|
||||
.map(|_| r.gen_range(-1.0..1.0))
|
||||
.collect();
|
||||
let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-0.5..0.5))
|
||||
.collect();
|
||||
|
||||
let regime_d = upload_f32(&stream, ®ime)?;
|
||||
let h_t_d = upload_f32(&stream, &h_t)?;
|
||||
let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
|
||||
mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?;
|
||||
|
||||
let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
|
||||
// pi_logits = log(pi_probs); exp + sum should equal 1 per batch row.
|
||||
for b in 0..b_size {
|
||||
let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS];
|
||||
// log-sum-exp stable accumulation (matches the kernel's
|
||||
// numerical guard — `pi_logits = log(pi_probs + 1e-12)`).
|
||||
let mut sum = 0.0_f32;
|
||||
for &lp in row {
|
||||
sum += lp.exp();
|
||||
}
|
||||
assert!(
|
||||
(sum - 1.0).abs() < TOL,
|
||||
"exp(pi_logits[b={b}]).sum() = {sum}, expected ≈1.0 (row = {row:?})"
|
||||
);
|
||||
}
|
||||
eprintln!(
|
||||
"PASS — exp(pi_logits) sums to 1.0 (±{TOL}) for all {b_size} batch rows (mixture is a valid distribution)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn k1_reduces_to_single_head() -> Result<()> {
|
||||
let Some(dev) = build_device() else {
|
||||
return Ok(());
|
||||
};
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let b_size = 4usize;
|
||||
let k = 1usize;
|
||||
let mhp = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed: 0xC3C3,
|
||||
},
|
||||
)?;
|
||||
|
||||
// h_t = 0 so the policy-head matmul produces zero logits;
|
||||
// post-init, the bias is the only source of non-zero pi_logits_k.
|
||||
// At K=1 the head's bias is [+0.5, 0, 0, ..., 0] (Head 0 → ShortLarge).
|
||||
let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
|
||||
let regime: Vec<f32> = (0..b_size * REGIME_DIM).map(|i| (i as f32) * 0.01).collect();
|
||||
let h_t_d = upload_f32(&stream, &h_t)?;
|
||||
let regime_d = upload_f32(&stream, ®ime)?;
|
||||
let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
|
||||
mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?;
|
||||
|
||||
let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
|
||||
|
||||
// At K=1: mixture = single head's softmax. With h_t=0 the per-head
|
||||
// pre-softmax logits are exactly b_heads = [+0.5, 0, ..., 0]. So:
|
||||
// p[0] = exp(0.5) / (exp(0.5) + (N_ACTIONS-1))
|
||||
// p[a≠0] = 1 / (exp(0.5) + (N_ACTIONS-1))
|
||||
let denom = 0.5_f32.exp() + (N_ACTIONS as f32 - 1.0);
|
||||
let p0_expected = 0.5_f32.exp() / denom;
|
||||
let pother_expected = 1.0_f32 / denom;
|
||||
// The kernel writes pi_logits = log(pi_probs + 1e-12); convert back.
|
||||
for b in 0..b_size {
|
||||
let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS];
|
||||
let p0 = row[0].exp();
|
||||
assert!(
|
||||
(p0 - p0_expected).abs() < TOL,
|
||||
"k1 single-head: b={b} p[0]={p0} expected {p0_expected}"
|
||||
);
|
||||
for a in 1..N_ACTIONS {
|
||||
let pa = row[a].exp();
|
||||
assert!(
|
||||
(pa - pother_expected).abs() < TOL,
|
||||
"k1 single-head: b={b} p[{a}]={pa} expected {pother_expected}"
|
||||
);
|
||||
}
|
||||
}
|
||||
eprintln!(
|
||||
"PASS — K=1 reduces to single-head softmax over bias-only logits; p[0] = {p0_expected:.6}, p[a>0] = {pother_expected:.6}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn gate_responds_to_regime_change() -> Result<()> {
|
||||
let Some(dev) = build_device() else {
|
||||
return Ok(());
|
||||
};
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let b_size = 2usize;
|
||||
let k = 3usize;
|
||||
let mut mhp = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed: 0xD4D4,
|
||||
},
|
||||
)?;
|
||||
|
||||
// Overwrite the gating weights so head 0 strongly tracks regime[0]:
|
||||
// W_gate[0, :] = [+5.0, 0, 0, 0, 0, 0]
|
||||
// W_gate[1, :] = [0, 0, 0, 0, 0, 0]
|
||||
// W_gate[2, :] = [-5.0, 0, 0, 0, 0, 0]
|
||||
// This forces a strong gate response to regime[0]: large positive
|
||||
// regime[0] routes mass to head 0, large negative routes to head 2.
|
||||
let mut w_gate_host = vec![0.0_f32; k * REGIME_DIM];
|
||||
w_gate_host[0 * REGIME_DIM + 0] = 5.0;
|
||||
w_gate_host[2 * REGIME_DIM + 0] = -5.0;
|
||||
write_slice_f32_d_pub(&stream, &w_gate_host, &mut mhp.w_gate_d)?;
|
||||
|
||||
// h_t doesn't influence the gate (gate reads only regime_h).
|
||||
let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
|
||||
let h_t_d = upload_f32(&stream, &h_t)?;
|
||||
let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
|
||||
// Run 1: regime[0] = +1 (one batch) / +1 (other batch) → mass routes to head 0.
|
||||
let regime_pos = vec![1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
|
||||
let regime_pos_d = upload_f32(&stream, ®ime_pos)?;
|
||||
mhp.forward(&h_t_d, ®ime_pos_d, &mut pi_logits_out)?;
|
||||
let gp_pos = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;
|
||||
|
||||
// Run 2: regime[0] = -1 / -1 → mass routes to head 2.
|
||||
let regime_neg = vec![-1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
|
||||
let regime_neg_d = upload_f32(&stream, ®ime_neg)?;
|
||||
mhp.forward(&h_t_d, ®ime_neg_d, &mut pi_logits_out)?;
|
||||
let gp_neg = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;
|
||||
|
||||
// Total Variation Distance between the two gate distributions per batch.
|
||||
// TVD = 0.5 × Σ_k |p1[k] − p2[k]|. With ±5σ logit swing on head 0 vs
|
||||
// head 2 (gap of 10 in pre-softmax), TVD should easily exceed 0.5.
|
||||
for b in 0..b_size {
|
||||
let row_pos = &gp_pos[b * k..(b + 1) * k];
|
||||
let row_neg = &gp_neg[b * k..(b + 1) * k];
|
||||
let tvd: f32 = row_pos
|
||||
.iter()
|
||||
.zip(row_neg.iter())
|
||||
.map(|(p, q)| (p - q).abs())
|
||||
.sum::<f32>()
|
||||
* 0.5;
|
||||
assert!(
|
||||
tvd > 0.5,
|
||||
"gate not responding to regime[0]: b={b} TVD={tvd}; pos={row_pos:?} neg={row_neg:?}"
|
||||
);
|
||||
// With W_gate[0,0]=+5 and W_gate[2,0]=-5 and r0=+1: gate logits =
|
||||
// [+5, 0, -5] → after softmax p[0] dominates. Sanity-check that
|
||||
// p[0] is the modal head under positive regime.
|
||||
let modal_pos =
|
||||
row_pos.iter().enumerate().fold(
|
||||
0usize,
|
||||
|acc, (i, p)| if *p > row_pos[acc] { i } else { acc },
|
||||
);
|
||||
assert_eq!(
|
||||
modal_pos, 0,
|
||||
"under regime[0]=+1, expected head 0 to be modal; got head {modal_pos} (row = {row_pos:?})"
|
||||
);
|
||||
let modal_neg =
|
||||
row_neg.iter().enumerate().fold(
|
||||
0usize,
|
||||
|acc, (i, p)| if *p > row_neg[acc] { i } else { acc },
|
||||
);
|
||||
assert_eq!(
|
||||
modal_neg, 2,
|
||||
"under regime[0]=-1, expected head 2 to be modal; got head {modal_neg} (row = {row_neg:?})"
|
||||
);
|
||||
}
|
||||
eprintln!(
|
||||
"PASS — gate routes mass between heads 0/2 in response to regime[0] sign (TVD > 0.5 per batch)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn forward_is_deterministic_across_contexts() -> Result<()> {
|
||||
let Some(dev) = build_device() else {
|
||||
return Ok(());
|
||||
};
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let b_size = 4usize;
|
||||
let k = 3usize;
|
||||
let seed = 0xE5E5;
|
||||
|
||||
// Identical inputs for both runs.
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
let mut r = ChaCha8Rng::seed_from_u64(0xF00D);
|
||||
let regime: Vec<f32> = (0..b_size * REGIME_DIM)
|
||||
.map(|_| r.gen_range(-1.0..1.0))
|
||||
.collect();
|
||||
let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-0.5..0.5))
|
||||
.collect();
|
||||
let regime_d = upload_f32(&stream, ®ime)?;
|
||||
let h_t_d = upload_f32(&stream, &h_t)?;
|
||||
|
||||
// Run A
|
||||
let mhp_a = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed,
|
||||
},
|
||||
)?;
|
||||
let mut out_a = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
mhp_a.forward(&h_t_d, ®ime_d, &mut out_a)?;
|
||||
let pi_a = read_slice_d_pub(&stream, &out_a, b_size * N_ACTIONS)?;
|
||||
|
||||
// Run B — fresh instance, identical seed/config
|
||||
let mhp_b = MultiHeadPolicy::new(
|
||||
&dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k,
|
||||
b_size,
|
||||
seed,
|
||||
},
|
||||
)?;
|
||||
let mut out_b = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
mhp_b.forward(&h_t_d, ®ime_d, &mut out_b)?;
|
||||
let pi_b = read_slice_d_pub(&stream, &out_b, b_size * N_ACTIONS)?;
|
||||
|
||||
// Bit-equal — same-seed → identical weights → identical kernel inputs
|
||||
// → identical outputs. Use `to_bits()` to catch any sub-ULP drift that
|
||||
// would invalidate the determinism contract.
|
||||
for (i, (a, b)) in pi_a.iter().zip(pi_b.iter()).enumerate() {
|
||||
assert_eq!(
|
||||
a.to_bits(),
|
||||
b.to_bits(),
|
||||
"non-deterministic output at i={i}: a={a} b={b} (bits {} vs {})",
|
||||
a.to_bits(),
|
||||
b.to_bits()
|
||||
);
|
||||
}
|
||||
eprintln!(
|
||||
"PASS — forward bit-equal across two fresh contexts with seed=0x{seed:X} ({} floats compared)",
|
||||
pi_a.len()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,747 @@
|
||||
# Multi-Head Policy Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Each phase is a separate atomic commit with predetermined falsification gates.
|
||||
|
||||
**Goal:** Replace foxhunt's single-policy-head architecture with K=3 policy heads + regime-gated routing, bypassing VSN's softmax bottleneck on `regime[3]` (vol) per the empirical regime-attenuation finding. The mixture's gradient flows through foxhunt's actual π training mechanism (Q-distill from PER, not PPO surrogate).
|
||||
|
||||
**Architecture:** K=3 policy heads each `Linear(HIDDEN_DIM=128 → N_ACTIONS=11)`. Gating head `Linear(6 → 3)` reads raw regime features from `step_regime_d` directly (parallel channel — does NOT go through VSN/Mamba2). Mixture: `π_logits = log Σ_k softmax(gate_logits)[k] · softmax(pi_logits_k)`. Per-head aux KL prior loss (β=0.05) maintains specialization. Q-distill kernel writes gradient to `ss_pi_grad_logits_d` unchanged; the mixture backward distributes via standard MoE jacobian.
|
||||
|
||||
**Tech Stack:** Rust 1.85, cudarc, CUDA 12.4, sm_86 (RTX 3050) / sm_89 (L40S) / sm_90 (H100). Existing crate: `ml-alpha`. Spec: `docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md` (read FULL spec + ADDENDUM 2026-06-03).
|
||||
|
||||
**Branch target:** New branch `ml-alpha-multi-head` off `ml-alpha-regime-observer` HEAD (currently `779c03b9d`). Each phase is a separate commit. Phase 2A-D is the cluster-verification gate.
|
||||
|
||||
**Prerequisites (must verify before dispatching):**
|
||||
- `git rev-parse HEAD` returns `779c03b9d` (or later if Phase 1B-C/D landed first)
|
||||
- Working tree clean
|
||||
- `./scripts/determinism-check.sh --quick` exits 0 at HEAD
|
||||
- Spec ADDENDUM 2026-06-03 read and understood (Q-distill, gating head Option B = regime-direct, refined falsification gates)
|
||||
- `step_regime_d` allocated in `crates/ml-alpha/src/trainer/perception.rs:1795` and accessible at trainer's per-step path
|
||||
- Read `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo` (load-bearing for not breaking π training)
|
||||
- Read `pearl_local_smoke_noise_floor_and_regime_concentration` (load-bearing for interpreting smoke results)
|
||||
|
||||
---
|
||||
|
||||
## Phase 2A-A — ISV slots + MultiHeadPolicy struct + mixture forward kernel + GPU-oracle tests
|
||||
|
||||
**Goal:** Foundation. New ISV slots, K policy head + gating head data structures, mixture forward kernel with GPU-oracle invariant tests. No trainer integration yet — components work standalone.
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-alpha/src/rl/multi_head_policy.rs`
|
||||
- Create: `crates/ml-alpha/cuda/multi_head_policy_forward.cu`
|
||||
- Create: `crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu`
|
||||
- Create: `crates/ml-alpha/tests/multi_head_policy_invariants.rs`
|
||||
- Modify: `crates/ml-alpha/src/rl/isv_slots.rs` (4 new slots + RL_SLOTS_END bump)
|
||||
- Modify: `crates/ml-alpha/build.rs` (register new .cu)
|
||||
- Modify: `crates/ml-alpha/src/rl/mod.rs` (pub mod multi_head_policy)
|
||||
|
||||
### Task A.1 — Allocate ISV slots
|
||||
|
||||
- [ ] **Step 1: Add 4 new ISV slots in `isv_slots.rs`**
|
||||
|
||||
After the existing slots (after RL_PPO_N_MINIBATCHES_INDEX = 760 if Phase 1B-A landed; else after RL_SURFER_SCAFFOLD_FORCE_PIN_INDEX = 757):
|
||||
|
||||
```rust
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Phase 2A multi-head policy (2026-06-03)
|
||||
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
// ADDENDUM 2026-06-03 (Q-distill compatibility + regime-direct gating)
|
||||
// Pearl: pearl_local_smoke_noise_floor_and_regime_concentration
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Number of policy heads. Bootstrap K=3. Range allowed [1, 8]. K=1
|
||||
/// reduces to single-head (legacy behavior — for ablation only).
|
||||
pub const RL_POLICY_NUM_HEADS_INDEX: usize = 761;
|
||||
|
||||
/// Gating distribution entropy floor. Bootstrap log(K) × 0.5 = 0.549 for K=3.
|
||||
/// Prevents gate collapse to single-head mode per
|
||||
/// `pearl_pi_actor_collapses_without_entropy_floor`.
|
||||
pub const RL_POLICY_GATING_ENTROPY_FLOOR_INDEX: usize = 762;
|
||||
|
||||
/// Per-head entropy floor. Bootstrap log(N_ACTIONS) × 0.4 = 0.959 for N_ACTIONS=11.
|
||||
/// Prevents per-head determinism collapse.
|
||||
pub const RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX: usize = 763;
|
||||
|
||||
/// β coefficient for per-head auxiliary KL prior loss. Bootstrap 0.05
|
||||
/// (5% of policy loss magnitude). Range [0.0, 0.5]. β=0.0 → no aux
|
||||
/// (specialization driven only by gate + init bias).
|
||||
pub const RL_POLICY_AUX_PRIOR_BETA_INDEX: usize = 764;
|
||||
```
|
||||
|
||||
Update `RL_SLOTS_END = 765`.
|
||||
|
||||
- [ ] **Step 2: Build sanity**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml-alpha
|
||||
```
|
||||
Expected: exit 0.
|
||||
|
||||
### Task A.2 — Write `multi_head_policy_forward.cu` kernel
|
||||
|
||||
- [ ] **Step 1: Create kernel file**
|
||||
|
||||
Path: `crates/ml-alpha/cuda/multi_head_policy_forward.cu`
|
||||
|
||||
```c
|
||||
// multi_head_policy_forward.cu — K-head policy logit + mixture combination.
|
||||
//
|
||||
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
// ADDENDUM 2026-06-03 §R.6
|
||||
//
|
||||
// Per (batch, head) compute pi_logits_k = W_k @ h_t + b_k (HIDDEN_DIM=128 → N_ACTIONS=11).
|
||||
// Then per batch compute mixture:
|
||||
// pi_probs_k[b, a] = softmax_a(pi_logits_k[b])
|
||||
// pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[b, a]
|
||||
// pi_logits[b, a] = log(pi_probs[b, a] + 1e-12)
|
||||
//
|
||||
// Deterministic by construction: per-batch sequential mixture sum (no parallel
|
||||
// reductions across batches; the only reduction is over K=3 heads with fixed
|
||||
// summation order).
|
||||
//
|
||||
// Per feedback_no_atomicadd: no atomic ops.
|
||||
// Per feedback_no_nvrtc: pre-compiled cubin.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define N_ACTIONS 11
|
||||
#define MAX_K_HEADS 8 // Compile-time max; runtime K from ISV
|
||||
|
||||
extern "C" __global__ void multi_head_policy_forward(
|
||||
const float* __restrict__ h_t, // [B × HIDDEN_DIM]
|
||||
const float* __restrict__ W_heads, // [K × N_ACTIONS × HIDDEN_DIM]
|
||||
const float* __restrict__ b_heads, // [K × N_ACTIONS]
|
||||
const float* __restrict__ gate_probs, // [B × K] — already softmax'd
|
||||
float* __restrict__ pi_logits_k, // [B × K × N_ACTIONS] (output, for backward)
|
||||
float* __restrict__ pi_probs_k, // [B × K × N_ACTIONS] (output, for backward)
|
||||
float* __restrict__ pi_logits, // [B × N_ACTIONS] (output, mixture)
|
||||
const int B,
|
||||
const int K
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
const int tid = threadIdx.x; // 0..N_ACTIONS-1
|
||||
|
||||
__shared__ float s_pi_logits_k[MAX_K_HEADS * N_ACTIONS];
|
||||
__shared__ float s_pi_probs_k[MAX_K_HEADS * N_ACTIONS];
|
||||
__shared__ float s_pi_probs[N_ACTIONS];
|
||||
__shared__ float s_max_k[MAX_K_HEADS]; // for softmax stability
|
||||
__shared__ float s_sum_k[MAX_K_HEADS];
|
||||
|
||||
// Step 1: compute pi_logits_k[b, k, a] = W_heads[k, a, :] @ h_t[b, :] + b_heads[k, a]
|
||||
// One thread per action; loops over k inside thread.
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float gl = b_heads[k * N_ACTIONS + tid];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < HIDDEN_DIM; ++j) {
|
||||
gl += W_heads[(k * N_ACTIONS + tid) * HIDDEN_DIM + j] * h_t[b * HIDDEN_DIM + j];
|
||||
}
|
||||
s_pi_logits_k[k * N_ACTIONS + tid] = gl;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Step 2: per-head softmax for pi_probs_k
|
||||
// Use tid 0..K-1 for max and sum reductions
|
||||
if (tid < K) {
|
||||
float mx = s_pi_logits_k[tid * N_ACTIONS];
|
||||
#pragma unroll
|
||||
for (int a = 1; a < N_ACTIONS; ++a) {
|
||||
float v = s_pi_logits_k[tid * N_ACTIONS + a];
|
||||
mx = fmaxf(mx, v);
|
||||
}
|
||||
s_max_k[tid] = mx;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float v = expf(s_pi_logits_k[k * N_ACTIONS + tid] - s_max_k[k]);
|
||||
s_pi_probs_k[k * N_ACTIONS + tid] = v;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < K) {
|
||||
float sm = 0.0f;
|
||||
#pragma unroll
|
||||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||||
sm += s_pi_probs_k[tid * N_ACTIONS + a];
|
||||
}
|
||||
s_sum_k[tid] = sm;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < N_ACTIONS) {
|
||||
for (int k = 0; k < K; ++k) {
|
||||
s_pi_probs_k[k * N_ACTIONS + tid] /= s_sum_k[k];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Step 3: mixture: pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[k, a]
|
||||
if (tid < N_ACTIONS) {
|
||||
float p = 0.0f;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < K; ++k) {
|
||||
p += gate_probs[b * K + k] * s_pi_probs_k[k * N_ACTIONS + tid];
|
||||
}
|
||||
s_pi_probs[tid] = p;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Step 4: write outputs
|
||||
if (tid < N_ACTIONS) {
|
||||
pi_logits[b * N_ACTIONS + tid] = logf(s_pi_probs[tid] + 1e-12f);
|
||||
for (int k = 0; k < K; ++k) {
|
||||
pi_logits_k[(b * K + k) * N_ACTIONS + tid] = s_pi_logits_k[k * N_ACTIONS + tid];
|
||||
pi_probs_k[(b * K + k) * N_ACTIONS + tid] = s_pi_probs_k[k * N_ACTIONS + tid];
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Register kernel in `build.rs`**
|
||||
|
||||
Add `compile_cuda_kernel("multi_head_policy_forward")?;` adjacent to existing kernel registrations.
|
||||
|
||||
- [ ] **Step 3: Build sanity**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml-alpha
|
||||
```
|
||||
Expected: exit 0, `target/.../multi_head_policy_forward.cubin` exists.
|
||||
|
||||
### Task A.3 — Write `multi_head_policy_gate_forward.cu` kernel
|
||||
|
||||
- [ ] **Step 1: Create gating head kernel**
|
||||
|
||||
Path: `crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu`
|
||||
|
||||
```c
|
||||
// multi_head_policy_gate_forward.cu — gating head from raw regime features.
|
||||
//
|
||||
// Spec ADDENDUM 2026-06-03 §R.3 Option B: gate head reads step_regime_d[b][0..6]
|
||||
// DIRECTLY (parallel channel — bypasses VSN/Mamba2 entirely).
|
||||
//
|
||||
// gate_logits[b, k] = W_gate[k, :] @ regime_h[b, :] + b_gate[k]
|
||||
// gate_probs[b, k] = softmax_k(gate_logits[b])
|
||||
//
|
||||
// Single-block-per-batch kernel, K threads.
|
||||
// Per feedback_no_atomicadd: no atomic ops.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#define REGIME_DIM 6
|
||||
#define MAX_K_HEADS 8
|
||||
|
||||
extern "C" __global__ void multi_head_policy_gate_forward(
|
||||
const float* __restrict__ regime_h, // [B × REGIME_DIM=6]
|
||||
const float* __restrict__ W_gate, // [K × REGIME_DIM]
|
||||
const float* __restrict__ b_gate, // [K]
|
||||
float* __restrict__ gate_logits, // [B × K] (output, for backward)
|
||||
float* __restrict__ gate_probs, // [B × K] (output, post-softmax)
|
||||
const int B,
|
||||
const int K
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
__shared__ float s_logits[MAX_K_HEADS];
|
||||
__shared__ float s_max;
|
||||
__shared__ float s_sum;
|
||||
|
||||
// Step 1: compute gate_logits[b, k] = W_gate[k, :] @ regime_h[b, :] + b_gate[k]
|
||||
if (tid < K) {
|
||||
float gl = b_gate[tid];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < REGIME_DIM; ++j) {
|
||||
gl += W_gate[tid * REGIME_DIM + j] * regime_h[b * REGIME_DIM + j];
|
||||
}
|
||||
s_logits[tid] = gl;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Step 2: softmax (max-subtract for stability)
|
||||
if (tid == 0) {
|
||||
float mx = s_logits[0];
|
||||
for (int k = 1; k < K; ++k) mx = fmaxf(mx, s_logits[k]);
|
||||
s_max = mx;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < K) {
|
||||
s_logits[tid] = expf(s_logits[tid] - s_max);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
float sm = 0.0f;
|
||||
for (int k = 0; k < K; ++k) sm += s_logits[k];
|
||||
s_sum = sm;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < K) {
|
||||
float v = s_logits[tid] / s_sum;
|
||||
gate_logits[b * K + tid] = s_logits[tid] * s_sum + s_max; // pre-softmax (for backward)
|
||||
gate_probs[b * K + tid] = v;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Register in build.rs**
|
||||
|
||||
```rust
|
||||
compile_cuda_kernel("multi_head_policy_gate_forward")?;
|
||||
```
|
||||
|
||||
### Task A.4 — Create `MultiHeadPolicy` struct
|
||||
|
||||
- [ ] **Step 1: Create file**
|
||||
|
||||
Path: `crates/ml-alpha/src/rl/multi_head_policy.rs`
|
||||
|
||||
```rust
|
||||
//! K-head policy with regime-gated routing.
|
||||
//!
|
||||
//! Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
//! ADDENDUM 2026-06-03 §R.3-R.6
|
||||
//!
|
||||
//! Forward:
|
||||
//! gate_logits = W_gate @ regime_h + b_gate (gating head — bypass VSN)
|
||||
//! gate_probs = softmax_K(gate_logits)
|
||||
//! pi_logits_k = W_heads_k @ h_t + b_heads_k (K policy heads)
|
||||
//! pi_probs_k = softmax_A(pi_logits_k)
|
||||
//! pi_probs = Σ_k gate_probs[k] · pi_probs_k (mixture)
|
||||
//! pi_logits = log(pi_probs)
|
||||
//!
|
||||
//! Backward: standard MoE jacobian via Q-distill gradient + aux prior loss merge.
|
||||
//!
|
||||
//! All transfers use mapped-pinned helpers (`feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
//! No atomicAdd anywhere (`feedback_no_atomicadd`).
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaContext, CudaSlice, CudaStream, CudaFunction, CudaModule};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::heads::{HIDDEN_DIM, N_ACTIONS};
|
||||
use crate::cfc::snap_features::REGIME_DIM;
|
||||
|
||||
pub const MAX_K_HEADS: usize = 8;
|
||||
|
||||
pub struct MultiHeadPolicy {
|
||||
pub b_size: usize,
|
||||
pub k: usize,
|
||||
|
||||
// Policy head weights ([K × N_ACTIONS × HIDDEN_DIM] flat)
|
||||
pub w_heads_d: CudaSlice<f32>,
|
||||
pub b_heads_d: CudaSlice<f32>,
|
||||
|
||||
// Gating head weights ([K × REGIME_DIM=6] flat)
|
||||
pub w_gate_d: CudaSlice<f32>,
|
||||
pub b_gate_d: CudaSlice<f32>,
|
||||
|
||||
// Forward output buffers
|
||||
pub gate_logits_d: CudaSlice<f32>, // [B × K]
|
||||
pub gate_probs_d: CudaSlice<f32>, // [B × K]
|
||||
pub pi_logits_k_d: CudaSlice<f32>, // [B × K × N_ACTIONS]
|
||||
pub pi_probs_k_d: CudaSlice<f32>, // [B × K × N_ACTIONS]
|
||||
|
||||
// Kernels
|
||||
pub forward_kernel: CudaFunction,
|
||||
pub gate_forward_kernel: CudaFunction,
|
||||
|
||||
_forward_module: Arc<CudaModule>,
|
||||
_gate_forward_module: Arc<CudaModule>,
|
||||
}
|
||||
|
||||
impl MultiHeadPolicy {
|
||||
pub fn new(
|
||||
ctx: &Arc<CudaContext>,
|
||||
stream: &Arc<CudaStream>,
|
||||
b_size: usize,
|
||||
k: usize,
|
||||
seed: u64,
|
||||
) -> Result<Self> {
|
||||
assert!(k > 0 && k <= MAX_K_HEADS, "K must be in [1, {}]", MAX_K_HEADS);
|
||||
|
||||
// Allocate weights with per-head action-bias init from spec §1.3.
|
||||
// Action enum verified against crates/ml-alpha/src/rl/common.rs:56-70:
|
||||
// 0=ShortLarge, 1=ShortSmall, 2=Hold, 3=FlatFromLong, 4=FlatFromShort,
|
||||
// 5=LongSmall, 6=LongLarge, 7=TrailTighten, 8=TrailLoosen, 9=HalfFlatLong, 10=HalfFlatShort.
|
||||
// Per-head specialization tri-split (asymmetry-break only — gate routes by regime):
|
||||
// Head 0 = SHORT-bias / risk-off → +0.5 on action 0 (ShortLarge)
|
||||
// Head 1 = LONG-bias / risk-on → +0.5 on actions 5,6 (LongSmall, LongLarge)
|
||||
// Head 2 = HOLD / stable → +0.5 on action 2 (Hold — Phase 0 dominant action)
|
||||
//
|
||||
// Standard Xavier scaled init via scoped_init_seed for determinism (path verified).
|
||||
use ml_core::cuda_autograd::init::scoped_init_seed;
|
||||
let _guard = scoped_init_seed(seed);
|
||||
|
||||
// Zero-init the linear weights; per-head action priors live on the bias vector.
|
||||
let w_heads_d = stream.alloc_zeros::<f32>(k * N_ACTIONS * HIDDEN_DIM)
|
||||
.context("alloc w_heads")?;
|
||||
let mut b_heads_init = vec![0.0_f32; k * N_ACTIONS];
|
||||
if k >= 1 { b_heads_init[0 * N_ACTIONS + 0] += 0.5; } // Head 0 → ShortLarge
|
||||
if k >= 2 { b_heads_init[1 * N_ACTIONS + 5] += 0.5;
|
||||
b_heads_init[1 * N_ACTIONS + 6] += 0.5; } // Head 1 → LongSmall, LongLarge
|
||||
if k >= 3 { b_heads_init[2 * N_ACTIONS + 2] += 0.5; } // Head 2 → Hold
|
||||
let b_heads_d = stream.alloc_zeros::<f32>(k * N_ACTIONS)?;
|
||||
stream.memcpy_htod(&b_heads_init, &b_heads_d)
|
||||
.context("upload b_heads init")?;
|
||||
|
||||
// Gating head: small init (Xavier scale)
|
||||
let gate_scale = (1.0_f32 / REGIME_DIM as f32).sqrt();
|
||||
let w_gate_init: Vec<f32> = (0..k * REGIME_DIM)
|
||||
.map(|i| ((i as f32 * 0.123_f32).sin() * gate_scale)).collect(); // deterministic init
|
||||
let w_gate_d = stream.alloc_zeros::<f32>(k * REGIME_DIM)?;
|
||||
stream.memcpy_htod(&w_gate_init, &w_gate_d)
|
||||
.context("upload w_gate init")?;
|
||||
let b_gate_d = stream.alloc_zeros::<f32>(k)?;
|
||||
|
||||
// Forward buffers
|
||||
let gate_logits_d = stream.alloc_zeros::<f32>(b_size * k)?;
|
||||
let gate_probs_d = stream.alloc_zeros::<f32>(b_size * k)?;
|
||||
let pi_logits_k_d = stream.alloc_zeros::<f32>(b_size * k * N_ACTIONS)?;
|
||||
let pi_probs_k_d = stream.alloc_zeros::<f32>(b_size * k * N_ACTIONS)?;
|
||||
|
||||
// Load kernels (repo convention: const &[u8] + load_cubin — see rl/ppo.rs:152-154, rl/dqn.rs:279).
|
||||
const FORWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_head_policy_forward.cubin"));
|
||||
const GATE_FORWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_head_policy_gate_forward.cubin"));
|
||||
let forward_module = ctx
|
||||
.load_cubin(FORWARD_CUBIN.to_vec())
|
||||
.context("load multi_head_policy_forward cubin")?;
|
||||
let forward_kernel = forward_module.load_function("multi_head_policy_forward")?;
|
||||
let gate_forward_module = ctx
|
||||
.load_cubin(GATE_FORWARD_CUBIN.to_vec())
|
||||
.context("load multi_head_policy_gate_forward cubin")?;
|
||||
let gate_forward_kernel = gate_forward_module.load_function("multi_head_policy_gate_forward")?;
|
||||
|
||||
Ok(Self {
|
||||
b_size, k,
|
||||
w_heads_d, b_heads_d, w_gate_d, b_gate_d,
|
||||
gate_logits_d, gate_probs_d, pi_logits_k_d, pi_probs_k_d,
|
||||
forward_kernel, gate_forward_kernel,
|
||||
_forward_module: forward_module, _gate_forward_module: gate_forward_module,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
h_t: &CudaSlice<f32>, // [B × HIDDEN_DIM]
|
||||
regime_h: &CudaSlice<f32>, // [B × REGIME_DIM=6]
|
||||
pi_logits_out: &mut CudaSlice<f32>, // [B × N_ACTIONS] — mixture output
|
||||
) -> Result<()> {
|
||||
// Step 1: gating head forward
|
||||
// ... launch gate_forward_kernel with (B, 1, 1) grid, K-thread block
|
||||
// ... reads regime_h, writes gate_logits_d + gate_probs_d
|
||||
//
|
||||
// Step 2: policy heads forward + mixture
|
||||
// ... launch forward_kernel with (B, 1, 1) grid, N_ACTIONS-thread block
|
||||
// ... reads h_t + gate_probs_d, writes pi_logits_k_d + pi_probs_k_d + pi_logits_out
|
||||
unimplemented!("forward launches — wire in Task A.4 step 2")
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Implement forward launches**
|
||||
|
||||
(Fill in the `forward` method with raw_launch invocations matching the kernel signatures.)
|
||||
|
||||
- [ ] **Step 3: Register module in `rl/mod.rs`**
|
||||
|
||||
Add `pub mod multi_head_policy;` to `crates/ml-alpha/src/rl/mod.rs` adjacent to the existing `pub mod ppo;`, `pub mod dqn;`, `pub mod iqn;`, `pub mod dueling_q;` lines.
|
||||
|
||||
- [ ] **Step 4: Build sanity**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml-alpha
|
||||
```
|
||||
|
||||
### Task A.5 — GPU-oracle invariant tests
|
||||
|
||||
- [ ] **Step 1: Create test file**
|
||||
|
||||
Path: `crates/ml-alpha/tests/multi_head_policy_invariants.rs`
|
||||
|
||||
5 GPU-oracle invariant tests:
|
||||
|
||||
```rust
|
||||
//! GPU-oracle invariants for MultiHeadPolicy.
|
||||
//!
|
||||
//! Tests run on actual GPU per `feedback_no_cpu_test_fallbacks`.
|
||||
//! Each test asserts a specific architectural property.
|
||||
|
||||
use anyhow::Result;
|
||||
use cudarc::driver::CudaContext;
|
||||
use ml_alpha::rl::multi_head_policy::MultiHeadPolicy;
|
||||
|
||||
const TOLERANCE: f32 = 1e-5;
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn gate_probs_sum_to_one() -> Result<()> {
|
||||
// Setup: K=3, B=4 batch. Random regime input.
|
||||
// Run forward.
|
||||
// Assert: gate_probs[b].sum() ≈ 1.0 for every b (softmax property).
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn pi_probs_mixture_sums_to_one() -> Result<()> {
|
||||
// Setup: K=3, B=4. Run forward.
|
||||
// Assert: exp(pi_logits[b]).sum() ≈ 1.0 for every b (mixture is a valid distribution).
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn k1_reduces_to_single_head() -> Result<()> {
|
||||
// Setup: K=1. Run forward with weights matching single-head Linear(128, 11).
|
||||
// Run a reference single-head policy_head Linear forward.
|
||||
// Assert: bit-equal output (within fp32 noise).
|
||||
// Verifies: K=1 ablation works as expected — no surprise mixture overhead.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn gate_responds_to_regime_change() -> Result<()> {
|
||||
// Setup: K=3 with hand-crafted gating weights (e.g., w_gate[0] = [1, 0, 0, 0, 0, 0]).
|
||||
// Run forward with regime_h = [10, 0, ...] vs regime_h = [-10, 0, ...].
|
||||
// Assert: gate_probs differs significantly between the two inputs (responds to regime[0]).
|
||||
// Verifies: gating head is sensitive to regime input (the load-bearing design property).
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_is_deterministic_across_contexts() -> Result<()> {
|
||||
// Setup: same seed → MultiHeadPolicy::new in two fresh contexts.
|
||||
// Run forward with identical input.
|
||||
// Assert: bit-equal output (`to_bits()` equality).
|
||||
// Verifies: no determinism regression vs `pearl_determinism_achieved`.
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture
|
||||
```
|
||||
Expected: 5/5 PASS.
|
||||
|
||||
### Task A.6 — Atomic commit + verify determinism not broken
|
||||
|
||||
- [ ] **Step 1: Stage Phase 2A-A changes** and verify `./scripts/determinism-check.sh --quick` exits 0 (pipeline unchanged since not wired in)
|
||||
|
||||
- [ ] **Step 2: Pre-commit hook on staged diff** — 0 memcpy_dtoh raw violations, 0 atomicAdd
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
Commit message template at the end of this plan.
|
||||
|
||||
---
|
||||
|
||||
## Phase 2A-B — Mixture backward kernel + aux prior loss
|
||||
|
||||
**Goal:** Implement the backward pass through the mixture (Q-distill gradient → K policy heads + gating head) and the auxiliary KL prior loss. No trainer integration yet.
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-alpha/cuda/multi_head_policy_backward.cu`
|
||||
- Create: `crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu`
|
||||
- Modify: `crates/ml-alpha/src/rl/multi_head_policy.rs` (add `backward` method)
|
||||
- Modify: `crates/ml-alpha/tests/multi_head_policy_invariants.rs` (add backward gradient tests)
|
||||
|
||||
### Tasks B.1-B.5 (outline — expand in 2A-A completion)
|
||||
|
||||
1. B.1 — Implement mixture backward jacobian (∂L/∂pi_logits → ∂L/∂pi_logits_k + ∂L/∂gate_logits)
|
||||
2. B.2 — Implement gate backward (∂L/∂gate_logits → ∂L/∂W_gate + ∂L/∂b_gate)
|
||||
3. B.3 — Implement per-head backward (∂L/∂pi_logits_k → ∂L/∂W_heads_k + ∂L/∂b_heads_k + ∂L/∂h_t)
|
||||
4. B.4 — Implement aux prior loss + gradient (Σ_k β · gate_probs[k] · KL(pi_probs_k ‖ pi_prior_k))
|
||||
5. B.5 — GPU-oracle finite-difference gradient tests (verify backward kernels against numerical gradients)
|
||||
|
||||
**Falsification at 2A-B**:
|
||||
- Finite-difference gradient agreement within 1e-3 relative tolerance
|
||||
- Aux prior loss is non-negative (KL ≥ 0)
|
||||
- Determinism preserved (sequential reductions, no atomicAdd)
|
||||
|
||||
---
|
||||
|
||||
## Phase 2A-C — Trainer integration with Q-distill gradient flow
|
||||
|
||||
**Goal:** Wire MultiHeadPolicy into IntegratedTrainer, replacing the single `policy_head` while preserving Q-distill gradient flow (the existing line 5577 zero-fill + Q-distill write pattern).
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (replace `policy_head` with MultiHeadPolicy; merge aux prior gradient into ss_pi_grad_logits_d post-Q-distill)
|
||||
- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs` (no API changes; trainer is the seam)
|
||||
- Modify: `crates/ml-alpha/tests/eval_diag_emission.rs` (EXPECTED_LEAVES bump for new policy.head_k_entropy, policy.gate_entropy, policy.gate_probs leaves)
|
||||
|
||||
### Tasks C.1-C.6 (outline — expand in 2A-B completion)
|
||||
|
||||
1. C.1 — Replace `policy_head` allocation with `MultiHeadPolicy::new(K=3, seed=...)` in trainer constructor
|
||||
2. C.2 — Bootstrap ISV slots 761-764 in `with_controllers_bootstrapped` table
|
||||
3. C.3 — Modify forward path: use MultiHeadPolicy::forward to produce pi_logits (consumed unchanged by downstream)
|
||||
4. C.4 — Modify backward path: AFTER Q-distill writes to ss_pi_grad_logits_d, merge aux prior gradient
|
||||
5. C.5 — Add diag emission: per-head entropy, gate entropy, gate_probs distribution, TVD diagnostic
|
||||
6. C.6 — Adam updates only on MultiHeadPolicy weights (W_heads, b_heads, W_gate, b_gate); encoder/Q/V heads unchanged
|
||||
|
||||
**Falsification at 2A-C** (mid-smoke b=128, 2000+500, seed=42, 3-seed):
|
||||
|
||||
| Gate | Threshold | Source |
|
||||
|---|---|---|
|
||||
| **G1 (PRIMARY)** | TVD(Q1, Q4 action dist) on session_pnl_variance_ema axis ≥ 0.10 | spec ADDENDUM §R.5 |
|
||||
| **G2** | Per-head pairwise KL ≥ 0.20 by step 1500 | spec §4 |
|
||||
| **G3** | Gate entropy ≥ log(K) × 0.45 = 0.49 | spec §4 |
|
||||
| **G4** | Argmax action mass ≤ 0.30 (down from baseline 0.40-0.54) | spec ADDENDUM §R.5 |
|
||||
| **G5** | Determinism preserved + cross-source $50 | Phase 0 contract |
|
||||
| **G6 (load-bearing empirical)** | Per-quartile: popart Q1 Δpnl improves by ≥ $1M (single-seed: borderline; 3-seed required for confidence) | spec ADDENDUM §R.5 |
|
||||
| G7 (aggregate, informational) | eval_summary.total_pnl_usd within 1σ of better than mean 3-seed Phase 0 baseline | high bar given $1.9M σ |
|
||||
|
||||
If G1-G5 PASS and G6 PASSES at 3-seed → proceed to Phase 2A-D cluster verification.
|
||||
If G1 FAILS (TVD on vol axis didn't lift) → multi-head bypass of VSN didn't work; STOP and diagnose.
|
||||
|
||||
---
|
||||
|
||||
## Phase 2A-D — Cluster verification 20k+5k
|
||||
|
||||
**Goal:** Validate at cluster scale. Cluster σ_eval_pnl is ~9× smaller than local (~$215k vs $1.9M), so cluster verdicts are far more sensitive. This is the load-bearing test.
|
||||
|
||||
### Tasks D.1-D.4 (outline — expand in 2A-C completion)
|
||||
|
||||
1. D.1 — Submit `./scripts/argo-alpha-rl.sh --branch ml-alpha-multi-head` at the 2A-C commit SHA
|
||||
2. D.2 — Monitor for completion (~70 min on L40S)
|
||||
3. D.3 — Analyze cluster verdict against per-quartile + aggregate gates
|
||||
4. D.4 — Update memory pearls with cluster verdict (success or failure)
|
||||
|
||||
**Falsification at 2A-D (cluster)**:
|
||||
- eval_summary.total_pnl_usd strictly better than Phase 0 baseline cluster value (TBD: requires baseline cluster run at HEAD `cfaa420bd`)
|
||||
- popart Q1 Δpnl improves by ≥ $5M vs baseline cluster (scales with batch)
|
||||
- Determinism preserved at cluster scale
|
||||
- TVD(Q1, Q4 action dist) ≥ 0.10 on vol axis at cluster scale
|
||||
|
||||
If cluster verdict shows improvement: proceed to 2A-E productionization.
|
||||
If cluster verdict shows no improvement OR regression: multi-head architecture doesn't deliver. Document negative result in pearl. Pivot to deeper architectural intervention (regime-stratified training, encoder changes, etc.).
|
||||
|
||||
---
|
||||
|
||||
## Phase 2A-E — Productionization
|
||||
|
||||
**Goal:** Make multi-head the default. Remove legacy single-head code path. Update pearls.
|
||||
|
||||
### Tasks E.1-E.4 (outline — expand in 2A-D completion)
|
||||
|
||||
1. E.1 — Default `RL_POLICY_NUM_HEADS_INDEX = 3` (production)
|
||||
2. E.2 — Remove single-head `policy_head` legacy code path (per `feedback_single_source_of_truth_no_duplicates`)
|
||||
3. E.3 — Update pearls: `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo` ADDENDUM "now via mixture", `pearl_local_smoke_noise_floor_and_regime_concentration` "verified bypass of VSN bottleneck"
|
||||
4. E.4 — Final cluster + local smoke parity check
|
||||
|
||||
---
|
||||
|
||||
## Commit message template (for Phase 2A-A)
|
||||
|
||||
```
|
||||
feat(ml-alpha): Phase 2A-A — MultiHeadPolicy foundation (K-head + regime-gated routing)
|
||||
|
||||
Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
ADDENDUM 2026-06-03 (Q-distill compatibility + regime-direct gating)
|
||||
Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md
|
||||
Pearl: pearl_local_smoke_noise_floor_and_regime_concentration
|
||||
|
||||
Foundation phase of multi-head policy refactor. Empirically motivated by
|
||||
the regime-attenuation finding: at Phase 0 baseline the policy is
|
||||
regime-conditional on popart_envelope (TVD=0.169 between Q1 and Q4
|
||||
action dists) but NOT on session_pnl_variance_ema (TVD=0.015) despite
|
||||
vol context being in encoder input at regime[3]. The asymmetry maps to
|
||||
VSN's softmax-over-40 gating constraint structurally limiting any single
|
||||
encoder-input feature to ~1/40 of attention budget. Multi-head policy
|
||||
with regime-gated routing reads step_regime_d[b][0..6] directly via the
|
||||
gating head (parallel channel) — bypassing VSN's bottleneck.
|
||||
|
||||
Phase 2A-A establishes the components WITHOUT trainer integration:
|
||||
|
||||
* 4 ISV slots (761-764):
|
||||
RL_POLICY_NUM_HEADS_INDEX = 761 (K=3 bootstrap)
|
||||
RL_POLICY_GATING_ENTROPY_FLOOR_INDEX = 762 (0.549 bootstrap)
|
||||
RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX = 763 (0.959 bootstrap)
|
||||
RL_POLICY_AUX_PRIOR_BETA_INDEX = 764 (0.05 bootstrap)
|
||||
RL_SLOTS_END 761 -> 765
|
||||
|
||||
* crates/ml-alpha/cuda/multi_head_policy_forward.cu (~180 LOC):
|
||||
Per-batch K-head logit + mixture combination. Reads h_t and
|
||||
gate_probs; writes pi_logits_k (for backward) + pi_logits (mixture
|
||||
for downstream). Deterministic sequential reduction over K heads.
|
||||
|
||||
* crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu (~80 LOC):
|
||||
Gating head reads step_regime_d[b][0..6] DIRECTLY (parallel
|
||||
channel — bypasses VSN). K-way softmax over gate logits.
|
||||
|
||||
* crates/ml-alpha/src/rl/multi_head_policy.rs (~300 LOC):
|
||||
MultiHeadPolicy struct with K policy heads + gating head + forward
|
||||
buffers. Per-head init bias for asymmetry-break specialization
|
||||
(Head 0: ShortLarge, Head 1: LongSmall+LongLarge, Head 2: Hold).
|
||||
Gating head reads regime DIRECTLY (bypass VSN); deterministic
|
||||
Xavier-style init. All allocations via existing patterns; no atomicAdd.
|
||||
|
||||
* crates/ml-alpha/tests/multi_head_policy_invariants.rs (~250 LOC):
|
||||
5 GPU-oracle invariant tests covering: gate softmax normalization,
|
||||
mixture probability sum, K=1 ablation reduces to single-head,
|
||||
gate responds to regime input, determinism across contexts.
|
||||
|
||||
Validation gates (Phase 2A-A complete when all pass):
|
||||
* cargo build --release --example alpha_rl_train -p ml-alpha: exit 0
|
||||
* cargo test multi_head_policy_invariants --release: 5/5 PASS
|
||||
* ./scripts/determinism-check.sh --quick: exit 0
|
||||
(pipeline unchanged since not wired into trainer yet)
|
||||
* Pre-commit hook: PASS (0 memcpy_dtoh, 0 atomicAdd)
|
||||
|
||||
Next: Phase 2A-B (mixture backward + aux prior loss kernels).
|
||||
|
||||
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Self-review (controller, before dispatch)
|
||||
|
||||
### Spec coverage
|
||||
- ✅ Spec §1.1 (architecture): K=3 heads + gating
|
||||
- ✅ Spec §1.2 (reward unchanged): Phase 2A doesn't touch reward kernel
|
||||
- ✅ Spec §1.3 (specialization): 4 mechanisms — init bias + per-head entropy floor + gate entropy floor + aux KL prior
|
||||
- ✅ Spec §1.4 (aux loss): β=0.05 bootstrap
|
||||
- ✅ Spec §2 (ISV slots): 4 slots at 761-764
|
||||
- ✅ ADDENDUM §R.2 (Q-distill flow): documented in plan
|
||||
- ✅ ADDENDUM §R.3 (gating head input Option B): direct regime read
|
||||
- ✅ ADDENDUM §R.5 (falsification): G1 TVD on vol axis, G6 per-quartile popart Q1
|
||||
- ✅ ADDENDUM §R.6 (gradient flow): aux prior merged AFTER Q-distill
|
||||
|
||||
### Placeholders
|
||||
- Task bodies for 2A-B, 2A-C, 2A-D, 2A-E are OUTLINES. Expand each when its prerequisite phase passes.
|
||||
- N_ACTIONS=11 layout assumed; bias init indices need verification (probably FLAT=0, HOLD_LONG=5, HOLD_SHORT=6, etc., but should be cross-referenced against `crates/ml-backtesting/src/order.rs` or similar before dispatch).
|
||||
|
||||
### Type consistency
|
||||
- `MultiHeadPolicy` uses `_d` suffix for device slices throughout
|
||||
- `K` capital, `b_size` lowercase, consistent with existing trainer struct naming
|
||||
|
||||
## Execution handoff
|
||||
|
||||
**Plan complete and saved to `docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md`.**
|
||||
|
||||
Per the methodology in `feedback_investigation_first_falsification_methodology`:
|
||||
|
||||
1. Dispatch Phase 2A-A coder NOW (foundation: ISV slots + struct + kernels + tests).
|
||||
2. After 2A-A commits + tests pass: expand 2A-B task bodies, dispatch 2A-B coder.
|
||||
3. After 2A-B commits + gradient finite-diff tests pass: expand 2A-C, dispatch coder.
|
||||
4. After 2A-C smoke passes G1-G5 (3-seed): dispatch 2A-D cluster verification.
|
||||
5. After 2A-D cluster validates G6: dispatch 2A-E productionization.
|
||||
|
||||
Each phase is one commit. Each phase's verdict feeds into the next phase's scope.
|
||||
|
||||
**NOT dispatching automatically.** Await explicit authorization from controller given the multi-week scope. The plan is the proper artifact; the coder dispatch is the next decision point.
|
||||
@@ -1,8 +1,8 @@
|
||||
# Multi-Head Policy with R-Multiple Reward (Intervention #4-clean)
|
||||
|
||||
**Date**: 2026-06-02
|
||||
**Date**: 2026-06-02 (ADDENDUM 2026-06-03 — Q-distill + empirical regime-attenuation findings)
|
||||
**Status**: Spec draft (post-sp-critical-reviewer BLOCK on prior v6 attempt)
|
||||
**Prerequisite**: Intervention #1 (R-multiple reward) has shipped and the eval verdict is in
|
||||
**Prerequisite**: ~~Intervention #1 (R-multiple reward) has shipped and the eval verdict is in~~ **SUPERSEDED — see ADDENDUM § R below.** Empirical regime-stratification of Phase 0 baseline shows the multi-head architectural intervention is directly motivated and doesn't require R-multiple as prerequisite.
|
||||
**Linked specs**: `2026-06-02-regime-invariance-four-interventions.md` (parent)
|
||||
**Linked pearls**: [[pearl_surfer_baseline_was_train_only_never_eval]], [[pearl_two_alpha_modes_surfer_vs_trend]], [[pearl_reward_signal_anti_aligned_with_pnl]], [[pearl_pi_actor_collapses_without_entropy_floor]]
|
||||
|
||||
@@ -199,6 +199,222 @@ Risk: PPO surrogate on a mixture has higher variance than on a single Gaussian-i
|
||||
|
||||
Mitigation: existing PPO ratio clamp from B-7 / B-9 work handles this. Per-batch advantage normalization (Phase 4.5) helps.
|
||||
|
||||
---
|
||||
|
||||
# ADDENDUM 2026-06-03 — Q-distill compatibility + Empirical regime-attenuation findings
|
||||
|
||||
## §R.0 — Why this addendum exists
|
||||
|
||||
The original spec was written assuming PPO trains π. Per [[pearl_foxhunt_pi_trained_by_q_distillation_not_ppo]] (Phase 1B-C STOP discovery): **foxhunt's π is trained by Q→π distillation, NOT by PPO surrogate.** `crates/ml-alpha/src/trainer/integrated.rs:5577` zeros `ss_pi_grad_logits_d` BEFORE the Q-distill kernel writes its KL gradient. The PPO surrogate runs as a diagnostic-only loss.
|
||||
|
||||
This addendum revises the spec to:
|
||||
1. Document Q-distill gradient flow for multi-head (it still works — just via a different channel)
|
||||
2. Incorporate empirical findings on VSN regime-attenuation
|
||||
3. Update falsification gates with concrete numerical anchors from baseline data
|
||||
4. Make multi-head a STANDALONE intervention (not contingent on R-multiple)
|
||||
|
||||
## §R.1 — Empirical motivation refined (Phase 0 baseline data)
|
||||
|
||||
Phase 0 baseline (`cfaa420bd`, b=128 mid-smoke, eval_pnl = −$4.46M) action distribution stratified by regime quartiles:
|
||||
|
||||
| Regime marker | TVD(Q1, Q4 action dist) | Pnl Q1 vs Q4 | Interpretation |
|
||||
|---|---|---|---|
|
||||
| `popart_envelope_max` (agent-driven via reward signal) | **0.169** | −$1.99M vs −$32k | Policy IS partially regime-aware via reward-signal channel |
|
||||
| `session_pnl_variance_ema` (market-driven vol) | **0.015** | −$1.40M vs −$2.05M | Policy is NOT regime-aware on market vol — despite vol context in encoder input at `regime[3]` (snap_feature out[23]) |
|
||||
|
||||
**The asymmetry is diagnostic.** The policy can condition on popart_envelope (which modifies reward magnitude directly) but cannot condition on market vol (which exists in encoder input as `regime[3] = log1p(ema_rv_med × 1e4)`).
|
||||
|
||||
### The mechanism
|
||||
|
||||
`crates/ml-alpha/cuda/variable_selection.cu` implements TFT-style softmax-gated feature selection over the first 40 features of the 56-dim encoder input. **Softmax gates sum to 1.0 across 40 features — average gate ≈ 0.025.** For `regime[3]` to influence Mamba2 meaningfully, it would need an above-average gate, but it competes with 39 other features (microstructure z-scores, OFI deltas, trade activity, spread, Δt Fourier features) for that scarce attention budget. The training signal (BCE labels + V regression + Q-distill targets) doesn't reward attending to `regime[3]` specifically, so VSN learns to allocate gates to whatever helps immediate prediction targets — not regime-conditional policy.
|
||||
|
||||
**Multi-head policy with regime-gated routing bypasses VSN's softmax bottleneck for regime context.** The gating head reads regime features directly (a parallel channel that doesn't compete via VSN's 40-way softmax). K heads receive a fully-dimensioned regime signal via the gating output (K-way softmax over a small K, not 40-way) and can specialize per regime.
|
||||
|
||||
### Policy collapse signature (additional)
|
||||
|
||||
Phase 0 baseline action_hist shows **action 2 (probably FLAT_SHORT-class) dominating 40-54% of action mass.** This is effectively a 1-action-dominated policy with noise — the encoder's representation has collapsed to "mostly emit action 2." Multi-head + per-head aux KL priors directly attacks this by REQUIRING per-head action distributions to retain mass across distinct action groups.
|
||||
|
||||
## §R.2 — Q-distill compatibility (the architectural correction)
|
||||
|
||||
Original §1.1 assumed PPO surrogate writes the policy gradient. Corrected flow:
|
||||
|
||||
```
|
||||
Current single-head:
|
||||
encoder → h_t → policy_head Linear(128→11) → π_logits
|
||||
↓
|
||||
ss_pi_grad_logits_d ← (zeroed line 5577) ← Q-distill KL gradient ← Q-Thompson sample from PER
|
||||
↓
|
||||
policy_head backward + Adam
|
||||
|
||||
Proposed multi-head:
|
||||
encoder → h_t (HIDDEN_DIM=128)
|
||||
↓
|
||||
K policy heads × Linear(128→11) → π_logits_k for k ∈ {0,1,2}
|
||||
↓
|
||||
Gating head Linear(GATE_INPUT_DIM → K) → gate_logits → softmax → gate[k]
|
||||
↓
|
||||
π_logits = log Σ_k gate[k] · softmax(π_logits_k) ← mixture distribution
|
||||
↓
|
||||
ss_pi_grad_logits_d ← (zeroed) ← Q-distill KL gradient writes here
|
||||
↓
|
||||
π_logits.backward → K head backward + gating head backward
|
||||
↓
|
||||
Adam updates on K head weights + gating head weights ONLY
|
||||
```
|
||||
|
||||
**Critical: Q-distill kernel doesn't care about multi-head structure.** It writes a gradient on `π_logits` regardless of how `π_logits` was computed. The backward pass through the mixture computes per-head + gate gradients automatically via chain rule.
|
||||
|
||||
The PPO surrogate gradient (still computed but still zeroed by line 5577) remains diagnostic-only. **No PPO behavior change required.** The aux KL prior loss (§1.4 in original spec) is added DIRECTLY to the policy backward pass, NOT zeroed by line 5577 — it lives in a separate gradient buffer that's added to `ss_pi_grad_logits_d` AFTER the Q-distill write. See §R.6 below for the gradient-merge sketch.
|
||||
|
||||
## §R.3 — Gating head input (the load-bearing design choice)
|
||||
|
||||
This is the single most important design choice in the addendum. Three options ordered by empirical motivation:
|
||||
|
||||
### Option A — Gate from encoder output `h_t` only (original spec §1.1 default)
|
||||
```
|
||||
gate_logits = Linear_HIDDEN_DIM_to_K(h_t)
|
||||
```
|
||||
**Problem**: h_t is the same representation that the empirical evidence shows under-encodes regime context. VSN's softmax bottleneck has already attenuated regime[3] before h_t is produced. Gating from h_t propagates the same bottleneck into the gating head.
|
||||
|
||||
### Option B — Gate from regime features DIRECTLY (recommended)
|
||||
```
|
||||
// Read the 6 raw regime features per batch element (computed by loader,
|
||||
// already in trainer's step_regime_d device buffer).
|
||||
regime_h = step_regime_d[b][0..6] // 6-dim raw regime context
|
||||
gate_logits = Linear_6_to_K(regime_h)
|
||||
```
|
||||
|
||||
**Why this works**: bypasses VSN entirely. The gating head's softmax is K-way (K=3) so each regime feature gets ≈ 1/3 attention budget, not 1/40. The gating head can fully condition on `regime[3]` (vol) without VSN's gating constraint.
|
||||
|
||||
**Where to source**: `step_regime_d` is allocated in `crates/ml-alpha/src/trainer/perception.rs:1795` and populated per-batch from the loader. It's already on device, already accessible per (b, k). For multi-head gating we want one regime vector per batch element (use t=K-1 of the BPTT window for "current regime") or aggregate across K (mean).
|
||||
|
||||
### Option C — Gate from concat(h_t, regime_features) (compromise)
|
||||
```
|
||||
gate_input = concat(h_t, regime_h) // dim 128 + 6 = 134
|
||||
gate_logits = Linear_134_to_K(gate_input)
|
||||
```
|
||||
Gives gate access to both the encoder's processed state AND raw regime. More parameters, more risk of overfit but more expressiveness.
|
||||
|
||||
**Decision**: Default to **Option B**. The empirical evidence (TVD=0.015 on session_pnl_var) shows h_t doesn't encode market vol; routing the gate through h_t inherits that limitation. Option B is the cleanest bypass of the diagnosed bottleneck.
|
||||
|
||||
## §R.4 — Per-head aux KL priors (refined)
|
||||
|
||||
Original §1.4 prior assignments:
|
||||
- Head 0 (trend): biased toward HOLD actions
|
||||
- Head 1 (reversion): biased toward FLAT actions
|
||||
- Head 2 (flat): biased toward FLAT_FLAT
|
||||
|
||||
**Refined per empirical Phase 0 data**:
|
||||
|
||||
Looking at the action distribution in popart Q1 (worst regime, −$1.99M loss):
|
||||
- Action 0+1 (FLAT-ish): 0.10
|
||||
- Action 2 (currently dominant): 0.40
|
||||
- Action 5+6 (HOLDs): 0.22
|
||||
- Action 7+8: 0.24
|
||||
|
||||
The Q4 distribution (best, −$32k) is heavily skewed toward action 2 (0.54) and away from HOLDs (0.11). The empirically-correct Q1 policy would likely DECREASE the HOLDs and increase decisive transitions. So the priors should target:
|
||||
|
||||
- **Head 0 (defensive — for high vol / drawdown recovery regime)**: bias toward FLAT_FLAT + decisive transition actions. Prior weights ~0.40 on action 0, 0.10 each on 1-3, 0.05 each on others.
|
||||
- **Head 1 (active — for low-vol stable regime)**: bias toward HOLDs and transitions. Prior weights ~0.18 each on actions 5-8, 0.05 on others.
|
||||
- **Head 2 (transition — for regime change events)**: bias toward action 2 (the currently-dominant one). Prior weights ~0.35 on action 2, 0.10 each on 1+3, 0.05 elsewhere.
|
||||
|
||||
These priors are SCAFFOLDING — `β=0.05` so the priors don't dominate training. Heads can diverge from priors if Q-distill signal pushes them elsewhere.
|
||||
|
||||
## §R.5 — Falsification gates (concrete numerical anchors from baseline)
|
||||
|
||||
Original §4.2 gates revised with measurable anchors:
|
||||
|
||||
### G1 — Regime conditioning EMERGES via the multi-head architecture
|
||||
- TVD(Q1 action dist, Q4 action dist) on **session_pnl_variance_ema axis** ≥ **0.10** by step 1500
|
||||
- (Baseline Phase 0: 0.015 — multi-head should lift this 6× if regime routing works)
|
||||
- This is the **direct test that multi-head bypasses VSN's regime bottleneck**.
|
||||
|
||||
### G2 — Per-head specialization holds
|
||||
- Pairwise KL between any two heads ≥ 0.20 by step 1500
|
||||
- Gate entropy ≥ log(K) × 0.45 = 0.49 (heads actually used, not collapsed to one)
|
||||
- No `gate_probs[k] < 0.10` for any k at step 5000 (mode-collapse kill)
|
||||
|
||||
### G3 — Action diversity recovers (policy collapse signature)
|
||||
- Argmax action mass ≤ **0.30** (down from baseline 0.40-0.54 dominated by action 2)
|
||||
- This tests whether the per-head priors + Q-distill differentiation prevent collapse
|
||||
|
||||
### G4 — Eval pnl improvement on POPART Q1 (lower variance than aggregate)
|
||||
- Sum of Δpnl in popart_envelope Q1 quartile improves by **≥ $1M** vs Phase 0 baseline (Q1: −$1.99M)
|
||||
- Per-quartile measurement has lower variance than aggregate (since quartile selection itself smooths some noise)
|
||||
- **This is the load-bearing empirical gate** — if multi-head moves the worst regime quartile, the chain works
|
||||
|
||||
### G5 — Aggregate eval pnl (informational, given noise floor)
|
||||
- `eval_summary.total_pnl_usd` better than seeds 42-44 mean by 2σ (σ ≈ $1.9M empirical)
|
||||
- Required improvement: ~$3.8M (e.g., from −$5.6M mean to better than −$1.8M)
|
||||
- This is a HIGH BAR. G4 is the primary gate; G5 is corroboration.
|
||||
|
||||
### G_kill — Per-step kill criteria
|
||||
- Action entropy < log(11) × 0.30 = 0.72 at step 5000 (policy collapse)
|
||||
- Eval realized_pnl_cum_usd < −$10M at step 10000 (catastrophic abandon)
|
||||
- Cross-source consistency drift > $50 (Phase 0 contract violation)
|
||||
- Determinism check fails at any seed
|
||||
|
||||
### Stratified verdict methodology
|
||||
Per the noise floor finding (σ ≈ $1.9M at b=128 single-seed), use multi-seed AND per-quartile measurement:
|
||||
- Run smokes at seeds 42, 43, 44 (3-seed minimum)
|
||||
- For each, compute per-quartile Δpnl on popart_envelope
|
||||
- Compare mean per-quartile improvement vs the 3-seed Phase 0 baseline distribution
|
||||
|
||||
## §R.6 — Gradient flow sketch (implementation guidance)
|
||||
|
||||
```
|
||||
// Forward (multi-head):
|
||||
for k in 0..K:
|
||||
pi_logits_k[b, a] = w_k[a, :] @ h_t[b, :] + b_k[a] // K linear heads
|
||||
gate_logits[b, k] = w_gate[k, :] @ regime_h[b, :] + b_gate[k] // gating head reads regime_h DIRECTLY
|
||||
gate_probs[b, k] = softmax(gate_logits[b])
|
||||
pi_probs_k[b, a] = softmax(pi_logits_k[b])
|
||||
pi_probs[b, a] = Σ_k gate_probs[b, k] * pi_probs_k[b, a] // mixture in probability space
|
||||
pi_logits[b, a] = log(pi_probs[b, a] + 1e-12) // back to logit space for downstream
|
||||
|
||||
// Backward (Q-distill + aux):
|
||||
ss_pi_grad_logits_d[b, a] := 0 // existing zero at line 5577
|
||||
rl_q_pi_distill_grad writes into ss_pi_grad_logits_d // existing path, unchanged
|
||||
ss_pi_aux_grad_logits_d[b, a] := Σ_k β · gate_probs[b, k] · ∂KL(pi_probs_k || pi_prior_k) / ∂pi_logits_k[b, a] // NEW
|
||||
ss_pi_grad_logits_d[b, a] += ss_pi_aux_grad_logits_d[b, a] // merge
|
||||
|
||||
// Then policy_head.backward (existing path) propagates ss_pi_grad_logits_d
|
||||
// back through the mixture (∂pi_logits/∂pi_logits_k via mixture jacobian)
|
||||
// and through the gating head (∂pi_logits/∂gate_logits via mixture jacobian)
|
||||
// Adam updates: w_k, b_k for each head + w_gate, b_gate
|
||||
```
|
||||
|
||||
The mixture's backward jacobian is standard (Shazeer 2017 MoE; standard implementation in production frameworks). ~100 LOC for two backward kernels (mixture_backward + aux_prior_backward).
|
||||
|
||||
## §R.7 — Scope estimate (revised)
|
||||
|
||||
| Component | Original spec | Revised |
|
||||
|---|---|---|
|
||||
| Multi-head policy struct + Adam states | ~80 LOC | ~80 LOC |
|
||||
| Mixture forward kernel | ~80 LOC | ~80 LOC |
|
||||
| Mixture backward kernel | ~120 LOC | ~140 LOC (handles K=3 with gate jacobian) |
|
||||
| Aux prior loss + backward | ~60 LOC | ~80 LOC |
|
||||
| Gating head: regime-direct input wiring | not in original | ~60 LOC (read step_regime_d, project to K) |
|
||||
| Bootstrap entries for new ISV slots | ~20 LOC | ~20 LOC |
|
||||
| Diag emission (per-head + gate + TVD diagnostics) | ~40 LOC | ~50 LOC (more diagnostics for regime conditioning) |
|
||||
| Tests (invariants for mixture, prior pull, gate entropy) | ~150 LOC | ~200 LOC |
|
||||
| **Total** | **~550 LOC** | **~710 LOC** |
|
||||
|
||||
Multi-week intervention. Single coder commit (per-component tasks would split it).
|
||||
|
||||
## §R.8 — Why this is now the right intervention (over R-multiple or further reward work)
|
||||
|
||||
Per ADDENDUM 2026-06-02e of `pearl_reward_signal_anti_aligned_with_pnl`: reward dimension has been tested twice (Phase 1A pure-pnl, Path A GAE-V). Both produced null/noise-level changes. Continuing in the reward dimension would be "going under the streetlight" (per `feedback_investigation_first_falsification_methodology`).
|
||||
|
||||
Multi-head policy targets a different mechanism:
|
||||
- **Architecturally addresses the diagnosed VSN bottleneck** (regime[3] attenuation)
|
||||
- **Empirically motivated**: TVD(0.015) on vol axis vs TVD(0.17) on popart axis proves the policy CAN condition on regime when the signal reaches it through the right channel
|
||||
- **Q-distill compatible**: gradient flows through mixture's standard backward jacobian; the existing Q-distill kernel is unchanged
|
||||
- **Targets the popart Q1 regime (worst $1.99M of $4.46M loss)** with measurable per-quartile gate
|
||||
- **Single ~700 LOC intervention** (multi-week but bounded)
|
||||
|
||||
If multi-head fails G1 (TVD on vol axis doesn't lift): the eval-collapse is genuinely not regime-mediated, and the architecture/data are the bottleneck. We'd then have empirical justification for the deeper interventions (regime-stratified training, encoder architecture changes, etc.).
|
||||
|
||||
### §5.4 Eval doesn't show specialization (heads pulled to average)
|
||||
Risk: even with all mitigations, eval might produce uniform wr across regime strata → no specialization actually achieved.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user