from __future__ import annotations import torch import torch.nn as nn from .blocks import DenseFFN, MoEFFN, TransformerBlock class StructuredBiHMoE(nn.Module): """ V2.2: fixes optic chiasm routing by introducing an explicit QUERY token type (last non-pad token), and implements a soft chiasm split (alpha-left) so query/mid/val flow primarily into the right (global) hemisphere. Key idea: xl = x * alpha_left + type_L(type) xr = x * (1 - alpha_left) + type_R(type) """ def __init__( self, vocab_size: int, d_model: int, n_heads: int, n_layers_stem: int, n_layers_hemi: int, d_ff_dense: int, d_ff_expert: int, n_experts: int, top_k: int, workspace_tokens: int = 4, reconcile_every: int = 2, dropout: float = 0.0, pad_id: int = 0, fuse: str = "mean", moe_warmup_steps: int = 0, # asymmetry left_local_window: int = 0, right_noise_std: float = 0.0, # callosum callosum_competitive: bool = True, callosum_tau: float = 1.0, # optic chiasm (type + soft split) chiasm_enabled: bool = False, chiasm_mode: str = "bias", # "bias" or "split" noise_vocab: int = 16, key_start: int = -1, key_count: int = 0, mid_start: int = -1, mid_count: int = 0, val_start: int = -1, val_count: int = 0, chiasm_alpha_key: float = 0.9, # fraction of x sent to LEFT for KEY chiasm_alpha_mid: float = 0.1, # fraction to LEFT for MID chiasm_alpha_val: float = 0.1, # fraction to LEFT for VAL chiasm_alpha_query: float = 0.1, # fraction to LEFT for QUERY (i.e., mostly RIGHT) chiasm_alpha_other: float = 0.5, output_gate: bool = True, output_tau: float = 1.0, ): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.workspace_tokens = workspace_tokens self.reconcile_every = max(1, reconcile_every) self.pad_id = int(pad_id) self.fuse = fuse self.moe_warmup_steps = int(moe_warmup_steps) self.left_local_window = int(left_local_window) self.right_noise_std = float(right_noise_std) self.callosum_competitive = bool(callosum_competitive) self.callosum_tau = float(callosum_tau) self.chiasm_enabled = bool(chiasm_enabled) self.chiasm_mode = str(chiasm_mode) self.noise_vocab = int(noise_vocab) self.key_start, self.key_count = int(key_start), int(key_count) self.mid_start, self.mid_count = int(mid_start), int(mid_count) self.val_start, self.val_count = int(val_start), int(val_count) self.chiasm_alpha_key = float(chiasm_alpha_key) self.chiasm_alpha_mid = float(chiasm_alpha_mid) self.chiasm_alpha_val = float(chiasm_alpha_val) self.chiasm_alpha_query = float(chiasm_alpha_query) self.chiasm_alpha_other = float(chiasm_alpha_other) self.output_gate = bool(output_gate) self.output_tau = float(output_tau) self.tok = nn.Embedding(vocab_size, d_model) # types: # 0=PAD, 1=SPECIAL(BOS/SEP), 2=NOISE, 3=KEY, 4=MID, 5=VAL, 6=QUERY, 7=OTHER self.n_types = 8 if self.chiasm_enabled: self.type_L = nn.Embedding(self.n_types, d_model) self.type_R = nn.Embedding(self.n_types, d_model) self.stem = nn.ModuleList([ TransformerBlock(d_model, n_heads, DenseFFN(d_model, d_ff_dense, dropout=dropout), dropout=dropout) for _ in range(n_layers_stem) ]) def hemi_blocks(): blocks = [] for i in range(n_layers_hemi): if (i % 2) == 1: ffn = MoEFFN(d_model, d_ff_expert, n_experts=n_experts, top_k=top_k, dropout=dropout, council_mlp=True) else: ffn = DenseFFN(d_model, d_ff_dense, dropout=dropout) blocks.append(TransformerBlock(d_model, n_heads, ffn, dropout=dropout)) return nn.ModuleList(blocks) self.left = hemi_blocks() self.right = hemi_blocks() self.workspace = nn.Parameter(torch.zeros(workspace_tokens, d_model)) self.w_to_l = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.w_to_r = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.write_gate = nn.Linear(d_model, 1) self.l_to_w = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.r_to_w = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.ln = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) def _type_ids(self, input_ids: torch.Tensor) -> torch.Tensor: B, T = input_ids.shape tid = torch.full_like(input_ids, 7) # OTHER # PAD tid = torch.where(input_ids == self.pad_id, torch.zeros_like(tid), tid) # SPECIAL (BOS=1, SEP=2) tid = torch.where((input_ids == 1) | (input_ids == 2), torch.full_like(tid, 1), tid) # NOISE range: [3, 3+noise_vocab) n0 = 3 n1 = n0 + max(0, self.noise_vocab) tid = torch.where((input_ids >= n0) & (input_ids < n1), torch.full_like(tid, 2), tid) # KEY if self.key_count > 0 and self.key_start >= 0: ks, ke = self.key_start, self.key_start + self.key_count tid = torch.where((input_ids >= ks) & (input_ids < ke), torch.full_like(tid, 3), tid) # MID if self.mid_count > 0 and self.mid_start >= 0: ms, me = self.mid_start, self.mid_start + self.mid_count tid = torch.where((input_ids >= ms) & (input_ids < me), torch.full_like(tid, 4), tid) # VAL if self.val_count > 0 and self.val_start >= 0: vs, ve = self.val_start, self.val_start + self.val_count tid = torch.where((input_ids >= vs) & (input_ids < ve), torch.full_like(tid, 5), tid) # QUERY = last non-pad token (override) nonpad = (input_ids != self.pad_id).to(torch.long) last = nonpad.sum(dim=1).clamp_min(1) - 1 # (B,) b = torch.arange(B, device=input_ids.device) tid[b, last] = 6 return tid def _last_nonpad(self, input_ids: torch.Tensor, h: torch.Tensor) -> torch.Tensor: nonpad = (input_ids != self.pad_id).to(torch.long) lengths = nonpad.sum(dim=1).clamp_min(1) - 1 idx = lengths.view(-1, 1, 1).expand(-1, 1, h.size(-1)) return h.gather(dim=1, index=idx).squeeze(1) def _band_mask(self, T: int, window: int, device) -> torch.Tensor: i = torch.arange(T, device=device).view(T, 1) j = torch.arange(T, device=device).view(1, T) return (j < (i - window)) | (j > (i + window)) # True = disallow def forward(self, input_ids: torch.Tensor, return_aux: bool = False, global_step: int | None = None): B, T = input_ids.shape key_padding_mask = (input_ids == self.pad_id) x = self.tok(input_ids) for b in self.stem: x = b(x, key_padding_mask=key_padding_mask) # Optic chiasm (bias or split) if self.chiasm_enabled: tid = self._type_ids(input_ids) if self.chiasm_mode == "bias": # Both hemispheres see the full token embedding; type embeddings bias specialization. xl = x + self.type_L(tid) xr = x + self.type_R(tid) else: # Soft split (legacy): attenuates content; can starve the global branch on long/noisy tasks. alpha = torch.full((B, T, 1), self.chiasm_alpha_other, device=input_ids.device, dtype=x.dtype) alpha = torch.where((tid == 3).unsqueeze(-1), torch.tensor(self.chiasm_alpha_key, device=input_ids.device, dtype=x.dtype), alpha) alpha = torch.where((tid == 4).unsqueeze(-1), torch.tensor(self.chiasm_alpha_mid, device=input_ids.device, dtype=x.dtype), alpha) alpha = torch.where((tid == 5).unsqueeze(-1), torch.tensor(self.chiasm_alpha_val, device=input_ids.device, dtype=x.dtype), alpha) alpha = torch.where((tid == 6).unsqueeze(-1), torch.tensor(self.chiasm_alpha_query, device=input_ids.device, dtype=x.dtype), alpha) xl = x * alpha + self.type_L(tid) xr = x * (1.0 - alpha) + self.type_R(tid) else: xl, xr = x, x if self.training and self.right_noise_std > 0.0: xr = xr + (self.right_noise_std * torch.randn_like(xr)) w = self.workspace.unsqueeze(0).expand(B, self.workspace_tokens, self.d_model) warmup_active = (global_step is not None and global_step < self.moe_warmup_steps) warmup_expert = int(global_step) if warmup_active else None left_attn_mask = self._band_mask(T, self.left_local_window, device=input_ids.device) if (self.left_local_window and self.left_local_window > 0) else None gate_L_mean = None gate_R_mean = None n = min(len(self.left), len(self.right)) for i in range(n): # Left local if isinstance(self.left[i].ffn, MoEFFN) and warmup_active: xl = self.left[i](xl, attn_mask=left_attn_mask, key_padding_mask=key_padding_mask, ffn_kwargs={"warmup_expert": warmup_expert}) else: xl = self.left[i](xl, attn_mask=left_attn_mask, key_padding_mask=key_padding_mask) # Right global if isinstance(self.right[i].ffn, MoEFFN) and warmup_active: xr = self.right[i](xr, key_padding_mask=key_padding_mask, ffn_kwargs={"warmup_expert": warmup_expert}) else: xr = self.right[i](xr, key_padding_mask=key_padding_mask) if ((i + 1) % self.reconcile_every) == 0: dwL, _ = self.w_to_l(w, xl, xl, need_weights=False, key_padding_mask=key_padding_mask) dwR, _ = self.w_to_r(w, xr, xr, need_weights=False, key_padding_mask=key_padding_mask) if self.callosum_competitive: sL = self.write_gate(dwL.float()) sR = self.write_gate(dwR.float()) scores = torch.cat([sL, sR], dim=-1) g = torch.softmax(scores / max(1e-6, self.callosum_tau), dim=-1) dw = (g[..., 0:1].to(dwL.dtype) * dwL) + (g[..., 1:2].to(dwR.dtype) * dwR) gate_L_mean = float(g[..., 0].mean().item()) gate_R_mean = float(g[..., 1].mean().item()) else: dw = dwL + dwR w = w + dw dl, _ = self.l_to_w(xl, w, w, need_weights=False) dr, _ = self.r_to_w(xr, w, w, need_weights=False) xl = xl + dl xr = xr + dr y_merge = self.ln(w.mean(dim=1)) logits_merge = self.head(y_merge) y_l_last = self.ln(self._last_nonpad(input_ids, xl)) y_r_last = self.ln(self._last_nonpad(input_ids, xr)) logits_l_last = self.head(y_l_last) logits_r_last = self.head(y_r_last) # Decision-level fusion if self.output_gate: # confidence = max softmax prob (computed in fp32) pL = torch.softmax(logits_l_last.float(), dim=-1) pR = torch.softmax(logits_r_last.float(), dim=-1) cL = pL.max(dim=-1).values # (B,) cR = pR.max(dim=-1).values # (B,) scores = torch.stack([cL, cR], dim=-1) # (B,2) w = torch.softmax(scores / max(1e-6, self.output_tau), dim=-1) # (B,2) wL = w[:, 0].to(logits_l_last.dtype).unsqueeze(-1) wR = w[:, 1].to(logits_r_last.dtype).unsqueeze(-1) logits_hemi = (wL * logits_l_last) + (wR * logits_r_last) logits = 0.5 * logits_merge + 0.5 * logits_hemi else: logits = (logits_merge + logits_l_last + logits_r_last) / 3.0 if self.fuse == "mean" else logits_merge if not return_aux: return logits pl = torch.softmax(logits_l_last.float(), dim=-1) pr = torch.softmax(logits_r_last.float(), dim=-1) kl_lr = torch.sum(pl * (pl.clamp_min(1e-9).log() - pr.clamp_min(1e-9).log()), dim=-1) kl_rl = torch.sum(pr * (pr.clamp_min(1e-9).log() - pl.clamp_min(1e-9).log()), dim=-1) sym_kl = 0.5 * (kl_lr + kl_rl) aux = { "sym_kl": sym_kl, "warmup_active": bool(warmup_active), "gate_L_mean": gate_L_mean, "gate_R_mean": gate_R_mean, "out_wL_mean": (float(w[:,0].mean().item()) if ('w' in locals()) else None), "out_wR_mean": (float(w[:,1].mean().item()) if ('w' in locals()) else None), } return logits, aux