File size: 2,346 Bytes
d266df1 | 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """
exit_gate.py — ExitGate: per-token adaptive depth routing.
Novel vs standard: every token in a standard transformer passes through all N
layers unconditionally. ExitGate computes a confidence scalar after each layer
norm. If confidence > exit_threshold the token's representation is frozen and
it skips all remaining layers, halving average compute on easy tokens.
The auxiliary training loss encourages the gate to be decisive (push toward 0
or 1) without forcing early exits — the model learns when it is confident.
"""
from __future__ import annotations
from typing import Tuple
import torch
import torch.nn as nn
from torch import Tensor
from .config import TMTConfig
class ExitGate(nn.Module):
"""Single linear → sigmoid confidence gate per token."""
def __init__(self, cfg: TMTConfig) -> None:
super().__init__()
self.threshold = cfg.exit_threshold
# Single scalar projection: d_model → 1
self.gate_proj = nn.Linear(cfg.d_model, 1)
nn.init.zeros_(self.gate_proj.weight)
nn.init.constant_(self.gate_proj.bias, -2.0) # start pessimistic
def forward(self, x: Tensor, exit_mask: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
"""
Args:
x: (B, S, D) current token representations
exit_mask: (B, S) bool — True where token has already exited
Returns:
x: (B, S, D) unchanged (gating is applied in TMTLayer)
new_mask: (B, S) bool — updated exit mask
confidence: (B, S) float confidence scores for auxiliary loss
"""
confidence = torch.sigmoid(self.gate_proj(x)).squeeze(-1) # (B, S)
# New exits: not already exited AND confidence above threshold
newly_exited = (~exit_mask) & (confidence > self.threshold)
new_mask = exit_mask | newly_exited
return x, new_mask, confidence
def auxiliary_loss(self, confidence: Tensor) -> Tensor:
"""
Encourage decisive gates — push confidence toward 0 or 1.
Loss = -E[|conf - 0.5|] so the model is penalised for being uncertain.
"""
return -(confidence - 0.5).abs().mean()
def __repr__(self) -> str:
p = sum(p.numel() for p in self.parameters())
return f"ExitGate(threshold={self.threshold}, params={p:,})"
|