gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: cross_branch_q_attention comment — document as gated projection not self-attention

Browse Source 
Sigmoid(Q·K/√d) × V is a per-head feature gate, not standard multi-head
attention (seq_len=1 makes softmax trivial). Comments now accurately
describe the learned gating mechanism and its purpose.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 07:50:02 +02:00
parent 3f4f0c7038
commit 64898b0cce
1 changed files with 21 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -2784,8 +2784,21 @@ extern "C" __global__ void glu_backward(
/* ================================================================== */
/**
* 2-head self-attention over 12 Q-values (branches × actions collapsed).
* One thread per sample.
* Cross-branch Q-value coordination via 2-head gated projection.
*
* NOT standard multi-head self-attention (seq_len=1 makes softmax trivial).
* Instead: each head computes a learned gating weight via sigmoid(Q·K/√d)
* that modulates the value projection V. This is a per-head feature gate
* that lets branches coordinate — e.g. direction Q-values can suppress or
* amplify magnitude Q-values through the learned W_Q/W_K/W_V projections.
*
* Forward:
* Q_proj = x @ W_Q + b_Q, K_proj = x @ W_K + b_K, V_proj = x @ W_V + b_V
* For each head h (dim 6):
* gate_h = sigmoid(dot(Q_h, K_h) / sqrt(6)) // learned scalar gate
* head_h = gate_h * V_h // gated value
* output = [head_0; head_1] @ W_O + b_O // project back
* Q_coord = LayerNorm(x + output) // residual + normalize
*
* Weight layout in attn_params (624 floats):
* W_Q[12,12], W_K[12,12], W_V[12,12] (3×144 = 432)
@@ -2794,10 +2807,7 @@ extern "C" __global__ void glu_backward(
* b_O[12] (12)
* Total: 432 + 36 + 144 + 12 = 624
*
* 2 heads, head_dim = 6.
* Residual + LayerNorm applied before writing Q_coord.
*
* Grid: ceil(B/256), Block: 256
* One thread per sample. Grid: ceil(B/256), Block: 256.
*/
extern "C" __global__ void cross_branch_q_attention(
const float* __restrict__ Q_raw, /* [B, 12] */
@@ -2845,16 +2855,16 @@ extern "C" __global__ void cross_branch_q_attention(
/* 2-head attention */
for (int h = 0; h < NH; h++) {
int off = h * HD;
/* Attention scores for this single "sequence" of 1 token:
* Self-attention on the 12-dim vector treated as D=12, seq_len=1.
* With seq_len=1, attention is trivially weight=1 → output = V.
* Instead we apply cross-head attention: each head attends to all D dims. */
/* Per-head gating: score = dot(Q_h, K_h) / sqrt(head_dim).
* Sigmoid gate ∈ (0,1) modulates the value projection.
* High Q·K agreement → gate opens → V passes through.
* Low Q·K agreement → gate closes → V suppressed. */
float score = 0.0f;
for (int d = 0; d < HD; d++) {
score += Q_proj[off + d] * K_proj[off + d];
}
score *= scale;
float attn_w = 1.0f / (1.0f + expf(-score)); /* sigmoid as single-token softmax */
float attn_w = 1.0f / (1.0f + expf(-score)); /* sigmoid gate */
for (int d = 0; d < HD; d++) {
attn_out[off + d] = attn_w * V_proj[off + d];
}