diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 16 |
1 files changed, 15 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index cca37b5..c1593f3 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -68,6 +68,11 @@ ap.add_argument('--drift_adapt', type=float, default=0.0) # >0: adaptive drift ap.add_argument('--beta_cap_rho', type=float, default=0.0) # >0: LOOP-GAIN CAP on beta — if per-sweep residual # ratio rho^ exceeds this, bscale *= 0.8 (beta backs off # under the wall-2 ceiling); recovers x1.02 when rho^ low +ap.add_argument('--logit_knee', type=float, default=0.0) # >0: piecewise clamp of attn logits above the + # knee (slope knee_slope) — cuts the switching-regime + # Jacobian of runaway sharp heads (b8-class carriers) + # without touching healthy logits below the knee +ap.add_argument('--knee_slope', type=float, default=0.2) ap.add_argument('--cap_floor', type=float, default=0.05) # hard bottom of the rho-cap; 0 = pure ceiling-tracking # (cap follows the measured ceiling all the way down; a # pinned bottom above the true ceiling = disguised wall-2) @@ -211,6 +216,8 @@ class Olmo2Attn(nn.Module): fr = torch.outer(torch.arange(T).float(), inv) self.register_buffer('rc', fr.cos(), persistent=False) self.register_buffer('rs', fr.sin(), persistent=False) + if args.logit_knee > 0: + self.register_buffer('cmask', torch.ones(T, T, dtype=torch.bool).tril(), persistent=False) def rope(self, x): x1, x2 = x[..., ::2], x[..., 1::2] c, s = self.rc[None, None], self.rs[None, None] @@ -222,7 +229,14 @@ class Olmo2Attn(nn.Module): q = self.rope(q.view(B, T, self.H, self.hd).transpose(1, 2)) k = self.rope(k.view(B, T, self.H, self.hd).transpose(1, 2)) v = v.view(B, T, self.H, self.hd).transpose(1, 2) - y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + if args.logit_knee > 0: + kn = args.logit_knee + lg = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5) + lg = torch.where(lg > kn, kn + args.knee_slope * (lg - kn), lg) + lg = lg.masked_fill(~self.cmask[:T, :T], float('-inf')) + y = lg.softmax(-1) @ v + else: + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) return self.proj(y.transpose(1, 2).contiguous().view(B, T, C)) class Olmo2Block(nn.Module): |