From 71096d092978f7e2b7466b4a98dbc06ef67aabbd Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sat, 18 Jul 2026 23:04:14 -0500 Subject: =?UTF-8?q?RESULT=2049:=20=E5=9B=9E=E8=B7=AF=E8=A7=A3=E5=89=96=3D2?= =?UTF-8?q?3=E6=AE=B5=E4=B9=98=E7=A7=AF=E8=BF=91=E4=B8=B4=E7=95=8C(?= =?UTF-8?q?=E4=B8=8B=E8=A1=8C=C3=9772/=E4=B8=8A=E8=A1=8C=C3=9713k/?= =?UTF-8?q?=E9=A1=B6=E8=AF=BB1e-6),plain=E5=8F=AA=E7=83=ADb8=E4=B8=80?= =?UTF-8?q?=E6=AE=B5(h6=20logit=2094=20vs=20cent=2075)=E8=BE=B9=E9=99=85?= =?UTF-8?q?=E8=B6=8A=E7=BA=BF,=E7=81=BE=E5=8F=98=3D=E7=A1=AC=E5=BA=95?= =?UTF-8?q?=E5=93=8D=E5=BA=94=E6=89=93=E4=BC=A4b1(=E5=85=A5=E5=8F=A3vjp=20?= =?UTF-8?q?1.3=E2=86=9214.9);=20batch=E5=99=AA=E5=A3=B0=E5=AF=B9=E7=85=A70?= =?UTF-8?q?.39/0.45=3DQ1=E5=B0=81=E5=8F=A3;=20--logit=5Fknee=2080=E5=AE=9E?= =?UTF-8?q?=E8=A3=85+knee=E8=87=82=E6=8E=92c190=E5=90=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) (limited to 'ep_run/casc_eq_train.py') 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): -- cgit v1.2.3