File size: 2,088 Bytes
9463e5c | 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 63 64 65 66 67 68 69 | """
GLADIUS v2.0 — Nexus Router
Routes hidden states to specialists. Top-k routing with load balancing.
Specialists run ON the kernel — they are not separate models.
STUB: Routes exist but only one specialist (reasoning) is wired.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from .config import KernelConfig
class NexusRouter(nn.Module):
"""
Learned router that activates top-k specialists per input.
argmax_specialist S(specialist | hidden_state)
"""
def __init__(self, config: KernelConfig):
super().__init__()
self.config = config
# Router: hidden_dim → num_specialists logits
self.gate = nn.Linear(config.hidden_dim, config.num_specialists, bias=False)
# Load balancing auxiliary loss coefficient
self.balance_coeff = 0.01
def forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Args:
hidden: (batch, hidden_dim) — pooled representation
Returns:
indices: (batch, top_k) — which specialists to activate
weights: (batch, top_k) — how much to weight each
"""
logits = self.gate(hidden) # (B, num_specialists)
probs = F.softmax(logits, dim=-1)
# Top-k selection
weights, indices = probs.topk(self.config.router_top_k, dim=-1)
# Renormalize weights
weights = weights / weights.sum(dim=-1, keepdim=True)
return indices, weights
def balance_loss(self, hidden: torch.Tensor) -> torch.Tensor:
"""
Auxiliary load-balancing loss.
Encourages equal specialist usage across a batch.
"""
logits = self.gate(hidden)
probs = F.softmax(logits, dim=-1)
# Mean probability per specialist across batch
mean_probs = probs.mean(dim=0)
# Ideal: uniform = 1/num_specialists
uniform = torch.ones_like(mean_probs) / self.config.num_specialists
# L2 distance from uniform
return self.balance_coeff * F.mse_loss(mean_probs, uniform)
|