news-impact-v1 / cross_attention_collapse_proof.md
AylinMaylinn's picture
Add cross_attention_collapse_proof.md (4,316 bytes)
00be985 verified
|
Raw
History Blame Contribute Delete
4.32 kB

Cross-Attention Collapse β€” Mathematical Proof

The buggy code (pipeline.py:1119-1124)

def forward(self, input_ids, attention_mask, chart_window):
    t_out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
    t_pool = t_out.last_hidden_state[:,0]          # (B, D)   - text CLS pool
    c_pool = self.chart_enc(chart_window)          # (B, D)   - chart pool
    fused, _ = self.fuse(t_pool.unsqueeze(1),       # query (B, 1, D)
                         c_pool.unsqueeze(1),       # key   (B, 1, D)
                         c_pool.unsqueeze(1))       # value (B, 1, D)
    f = fused.squeeze(1)                           # (B, D)

self.fuse is nn.MultiheadAttention(text_dim, num_heads=8, batch_first=True).

Why this collapses

Standard MHA forward, with Q, K, V each shaped (B, T, D):

  1. Linear projection split into H heads:
    • Q' = Q @ W_q, K' = K @ W_k, V' = V @ W_v β†’ reshape to (B, H, T, d_head)
  2. Scaled-dot scores: S = Q' @ K'^T / sqrt(d_head) β†’ (B, H, T_q, T_k)
  3. Attention weights: A = softmax(S, dim=-1) β†’ (B, H, T_q, T_k)
  4. Weighted values: O = A @ V' β†’ (B, H, T_q, d_head) β†’ concat to (B, T_q, D)
  5. Output: Y = O @ W_o + b_o

In our call, T_q = T_k = T_v = 1. Hence:

  • S has shape (B, H, 1, 1) β€” a single scalar per head, per batch.
  • softmax over a 1-element vector is identically 1.0. So A is identically 1.
  • O = A @ V' = V' exactly (per head).
  • Y = V' @ W_o + b_o.

Plug in: V' = V @ W_v = (c_pool.unsqueeze(1)) @ W_v. Squeeze:

f = c_pool @ W_v @ W_o + b_o

The text query t_pool and key t_pool never enter f. W_q, W_k, and the text encoder's entire 336M parameters are mathematically irrelevant to the forward pass output.

Independent sanity check

You can verify by toggling the text input β€” any input that produces the same c_pool will produce the same f. Concretely:

m.train(False)  # disable dropout/batchnorm side-effects
text_a = tokenizer("Apple beats Q3 earnings", return_tensors='pt', padding='max_length', max_length=128)
text_b = tokenizer("Bitcoin crashes 40%",   return_tensors='pt', padding='max_length', max_length=128)
chart = torch.randn(1, 128, 5)  # same chart for both

out_a = m(text_a.input_ids, text_a.attention_mask, chart)
out_b = m(text_b.input_ids, text_b.attention_mask, chart)
# Predicted: out_a['sig_logits'] == out_b['sig_logits'] (bit-exact)

If this is bit-exact, the bug is confirmed.

Why it matters

All four heads (head_sig, head_dir, head_mag, head_reason) consume only f. So:

  • The model cannot use news semantics for any prediction.
  • Iteration scoreboard's plateau at 2.495 is the chart-only loss floor, not a data/label noise floor.
  • The "encoder choice doesn't matter" observation is trivially true because no encoder choice is ever exercised.
  • sig_acc = 0.557 = majority class is what a chart-only model achieves when it cannot read the news that triggers the move.
  • The v6 "winner" claim (delta=0.0004 over v4) is comparing two run-to-run noise draws of effectively the same chart-only model.

The fix

Replace the cross-attention with a real fusion. Examples:

# Option A β€” concatenate then project (no attention, deterministic)
self.fuse = nn.Sequential(nn.Linear(2*text_dim, text_dim), nn.GELU())
def forward(...):
    f = self.fuse(torch.cat([t_pool, c_pool], dim=-1))
# Option B β€” proper cross-attention with chart used as a SEQUENCE
# Don't pool c_pool yet; let chart_enc return per-bar features (B, T, D),
# and use t_pool as query against all T chart bars:
class ChartEncoder(nn.Module):
    def forward(self, x):
        ...
        return self.proj_out(h)  # (B, T, D) β€” DO NOT mean(dim=1)

def forward(self, ...):
    c_seq = self.chart_enc(chart_window)        # (B, T, D)
    fused, _ = self.fuse(t_pool.unsqueeze(1),   # (B, 1, D) query
                         c_seq, c_seq)           # (B, T, D) key/value
    f = fused.squeeze(1)

Option B is the architecturally honest version of the original intent.

Severity

CRITICAL. Until this is fixed, the project is, at best, a chart-only price movement classifier conditioned on which symbol the news happened to mention. No news semantics are used.