File size: 1,102 Bytes
fc3971e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """TENP expert-neuron projection plus optional energy control. No OLS."""
from __future__ import annotations
import torch
PROJ_EPS = 1e-6
def gated_signed_projection(
z: torch.Tensor,
y: torch.Tensor,
gate: torch.Tensor,
W_down: torch.Tensor,
eps: float = PROJ_EPS,
) -> torch.Tensor:
"""Per-token score `q = g * p` with signed projection onto expert output y.
z: [n, d_int] SwiGLU mid activations
y: [n, d] expert output (not yet multiplied by router gate)
gate: [n] router weight (top-1 after renormalize is ~1)
W_down: [d, d_int] down_proj.weight
returns q: [n, d_int]
"""
w = W_down.detach().to(device=y.device, dtype=torch.float32)
y_f = y.float()
z_f = z.float()
g = gate.float().reshape(-1, 1)
y_norm = y_f.norm(dim=-1, keepdim=True).clamp_min(eps)
p = z_f * (y_f @ w) / y_norm
return p * g
def gated_energy(z: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
"""Control score `g^2 z^2` (top-1 gate is ~1, so this is mostly z^2)."""
g = gate.float().reshape(-1, 1)
return (g * z.float()).square()
|