| """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() |
|
|
