#!/usr/bin/env python3 """ CEDL: A Hippocampal-Inspired Architecture for Language Modeling =============================================================== Reference implementation for the CEDL architecture at ~100M parameter scale, trained on WikiText-103 and evaluated via perplexity and LAMBADA accuracy. Architecture: C-Stage (EC) — Grid-cell periodic attention with per-layer AdaLN E-Stage (DG) — Expansion + top-k with AHSD and CSR pattern separation D-Stage (CA3) — Dual memory: attractor refinement + 256-slot memory bank L-Stage (CA1) — Two-channel comparator with learned fusion gate Feedback — Loop 1 (L→C via AdaLN), Loop 2 (D→E additive) NeuromodGate — Output-level gain modulation Baselines (parameter-matched): Transformer, Transformer-XL, RetNet, Mamba, LSTM Usage: python cedl_release.py --model CEDL --batch-size 4 --max-steps 30000 python cedl_release.py --model all --batch-size 8 --max-steps 30000 python cedl_release.py --model all --eval-only Paper: "CEDL: A Hippocampal-Inspired Architecture for Advancing LLMs" Author: Dian Jiao (University of Pennsylvania) License: MIT """ from __future__ import annotations import argparse import math import os import random import sys import time from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset LOGIT_SCALAR_MODES = frozenset({ "v3", "v3m", "v3m2", "v3m2_nll", "v4c", "v4c_gate", "v4d", }) @dataclass class Config: """Hyperparameters for 100M-scale training.""" dataset: str = "wikitext103" vocab_size: int = 50257 max_seq: int = 1024 batch_size: int = 4 grad_accum: int = 4 lr: float = 3e-4 min_lr: float = 1e-5 warmup_steps: int = 1000 max_steps: int = 30_000 weight_decay: float = 0.1 grad_clip: float = 1.0 bfloat16: bool = True eval_interval: int = 1000 eval_steps: int = 200 save_interval: int = 5000 save_dir: str = "checkpoints_100m" tpu: bool = False seed: int = 42 resume_checkpoint: Optional[str] = None resume_start_step: int = 0 save_trainstate: bool = False run_until_step: Optional[int] = None v5_grad_isolation: bool = False v5_trunk_aux_frac: float = 0.0 v5_aux_scale_early: float = 0.50 v5_aux_scale_mid: float = 0.25 v5_aux_scale_late: float = 0.10 v5_cadence_early: int = 8 v5_cadence_mid: int = 16 v5_cadence_late: int = 32 v5_family_reweight: bool = False v6_mixture: bool = False v6_lambda_init: float = -4.0 v6_lambda_a_init: float = 1.0 v6_aux_weight: float = 0.0 v6_margin_target: float = 1.0 v6_mix_weight: float = 0.5 v6_gate_weight: float = 0.01 v6_lambda_floor: float = 0.05 v6_lambda_head: bool = False v6_lambda_head_hidden: int = 160 v6_lambda_head_bias_init: float = -7.0 v6_bg_weight: float = 1.0 v6_bg_target: float = 0.01 v6_bce_objective: bool = False v6_sel_weight: float = 1.0 v6_lambda_head_w_init_std: float = 1e-3 v6_wt_sparsity_weight: float = 0.0 v6_wt_sparsity_target: float = 0.0 v6_mem_head_bank: bool = False v6_bank_ce_weight: float = 0.0 v6_bank_pair_ce_weight: float = 0.0 v6_bank_head_lr: float = 0.0 v6_bank_query_source: str = "h_d" v6_bank_readout_mode: str = "bank" v6_source_adapter: bool = False v6_context_adapter: bool = False v6_specialist_from_b0: bool = False v6_specialist_freeze: str = "none" v6_specialist_noinject: bool = False @dataclass class CEDLConfig: d_model: int = 640 n_heads: int = 10 c_layers: int = 8 ffn_dim: int = 2816 e_expand: int = 4 e_sparsity: float = 0.10 d_refine: int = 3 d_slots: int = 256 dropout: float = 0.1 n_feedback_iters: int = 1 feedback_decay: float = 0.5 feedback_warmup_start: int = 2000 feedback_warmup_end: int = 5000 sparsity_final: float = 0.05 sparsity_anneal_frac: float = 0.80 use_salience: bool = True salience_mode: str = "v0" lex_anchor_weight: float = 0.0 v4c_variant: str = "base" v4c_z_dim: int = 64 v4c_temperature: float = 0.05 v4c_proj_hidden: int = 256 v4c_nce_weight: float = 0.05 v4c_bce_weight: float = 0.02 v4c_norm_cap: float = 0.3 v4c_every: int = 4 v4c_batch_size: int = 32 v4c_max_seq: int = 128 v4c_warm_start_sigma: float = 0.05 v4c_margin_weight: float = 0.0 v4c_margin_target: float = 1.0 v4d_w_sal_sigma: float = 0.02 v4d_logit_cap: float = 4.0 v4d_causal_weight: float = 0.0 v4d_causal_gap: float = 0.25 v4d_causal_z_weight: float = 0.0 v4d_causal_z_gap: float = 0.15 v4d_w_sal_rank: int = 0 v4d_swap_consistency_weight: float = 0.0 v4d_role_sep_weight: float = 0.0 v4d_noinject: bool = False v6_mixture: bool = False v6_lambda_init: float = -4.0 v6_lambda_a_init: float = 1.0 v6_aux_weight: float = 0.0 v6_margin_target: float = 1.0 v6_mix_weight: float = 0.5 v6_gate_weight: float = 0.01 v6_lambda_floor: float = 0.05 v6_lambda_head: bool = False v6_lambda_head_hidden: int = 160 v6_lambda_head_bias_init: float = -7.0 v6_bg_weight: float = 1.0 v6_bg_target: float = 0.01 v6_bce_objective: bool = False v6_sel_weight: float = 1.0 v6_lambda_head_w_init_std: float = 1e-3 v6_wt_sparsity_weight: float = 0.0 v6_wt_sparsity_target: float = 0.0 v6_mem_head_bank: bool = False v6_bank_ce_weight: float = 0.0 v6_bank_pair_ce_weight: float = 0.0 v6_bank_query_source: str = "h_d" v6_bank_readout_mode: str = "bank" v6_source_adapter: bool = False v6_context_adapter: bool = False v4c_neg_weights: Dict[str, float] = field(default_factory=lambda: { "stale": 2.0, "distractor": 1.0, "neutral": 1.0}) @dataclass class TransformerConfig: d_model: int = 768 n_heads: int = 12 n_layers: int = 9 ffn_dim: int = 3072 dropout: float = 0.1 @dataclass class TransformerXLConfig: d_model: int = 768 n_heads: int = 12 n_layers: int = 9 ffn_dim: int = 3072 mem_len: int = 256 dropout: float = 0.1 @dataclass class RetNetConfig: d_model: int = 640 n_heads: int = 10 n_layers: int = 12 ffn_dim: int = 2880 dropout: float = 0.1 @dataclass class MambaConfig: d_model: int = 768 n_layers: int = 18 d_state: int = 16 d_conv: int = 4 expand: int = 2 dropout: float = 0.0 @dataclass class LSTMConfig: d_model: int = 640 hidden_size: int = 1760 n_layers: int = 3 dropout: float = 0.1 class MultiScaleRetention(nn.Module): """Parallel multi-scale retention (RetNet-style).""" def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads assert d_model % n_heads == 0 self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_out = nn.Linear(d_model, d_model) self.out_gate_proj = nn.Linear(d_model, d_model) self.gn = nn.GroupNorm(n_heads, d_model) gammas = torch.linspace(0.80, 0.99, n_heads) self.gamma_log = nn.Parameter( torch.log(gammas / (1.0 - gammas)) ) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, _ = x.shape K = self.n_heads dh = self.d_head Q = self.w_q(x).view(B, T, K, dh).transpose(1, 2) Kmat = self.w_k(x).view(B, T, K, dh).transpose(1, 2) V = self.w_v(x).view(B, T, K, dh).transpose(1, 2) gamma = torch.sigmoid(self.gamma_log) log_gamma = torch.log(gamma) positions = torch.arange(T, device=x.device, dtype=x.dtype) diff = (positions.unsqueeze(1) - positions.unsqueeze(0)).clamp(min=0) D = torch.exp(log_gamma.view(K, 1, 1) * diff.unsqueeze(0)) causal = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype)) D = D * causal.unsqueeze(0) retention = (Q @ Kmat.transpose(-1, -2)) * D.unsqueeze(0) row_sum = retention.sum(dim=-1, keepdim=True).abs().clamp(min=1.0) retention = retention / row_sum out = retention @ V out = out.transpose(1, 2).contiguous().view(B, T, -1) out = self.gn(out.transpose(1, 2)).transpose(1, 2) gate = F.silu(self.out_gate_proj(x)) out = self.w_out(self.dropout(out * gate)) return out class MultiScalePeriodicRetention(nn.Module): """Grid-cell-inspired periodic retention for CEDL C-stage. Replaces RetNet's exponential decay D[i,j] = γ^(i-j) with a damped periodic kernel D[i,j] = γ^(i-j) · cos(2π(i-j)/λ_k + φ_k), inspired by entorhinal cortex grid cells that fire at regular spatial intervals at multiple scales. No existing sequence model uses periodic attention weights — this is the key novelty of the C-stage. """ def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads assert d_model % n_heads == 0 self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_out = nn.Linear(d_model, d_model) self.out_gate_proj = nn.Linear(d_model, d_model) self.gn = nn.GroupNorm(n_heads, d_model) gammas = torch.linspace(0.85, 0.995, n_heads) self.gamma_log = nn.Parameter( torch.log(gammas / (1.0 - gammas))) self.log_lambda = nn.Parameter( torch.linspace(math.log(4.0), math.log(256.0), n_heads)) self.phi = nn.Parameter(torch.zeros(n_heads)) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, _ = x.shape K = self.n_heads dh = self.d_head Q = self.w_q(x).view(B, T, K, dh).transpose(1, 2) Kmat = self.w_k(x).view(B, T, K, dh).transpose(1, 2) V = self.w_v(x).view(B, T, K, dh).transpose(1, 2) gamma = torch.sigmoid(self.gamma_log) log_gamma = torch.log(gamma) lam = torch.exp(self.log_lambda) positions = torch.arange(T, device=x.device, dtype=x.dtype) diff = positions.unsqueeze(1) - positions.unsqueeze(0) diff_pos = diff.clamp(min=0) decay = torch.exp(log_gamma.view(K, 1, 1) * diff_pos.unsqueeze(0)) periodic = torch.cos( 2 * math.pi * diff.unsqueeze(0) / lam.view(K, 1, 1) + self.phi.view(K, 1, 1)) D = decay * periodic causal = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype)) D = D * causal.unsqueeze(0) retention = (Q @ Kmat.transpose(-1, -2)) * D.unsqueeze(0) row_sum = retention.abs().sum(dim=-1, keepdim=True).clamp(min=1.0) retention = retention / row_sum out = retention @ V out = out.transpose(1, 2).contiguous().view(B, T, -1) out = self.gn(out.transpose(1, 2)).transpose(1, 2) gate = F.silu(self.out_gate_proj(x)) out = self.w_out(self.dropout(out * gate)) return out class PeriodicRetentionLayer(nn.Module): """Pre-norm periodic retention layer with hierarchical neuromodulatory gating. Each layer's LayerNorm is modulated by feedback from the L-stage mismatch detector via learned scale+shift (Adaptive LayerNorm). This distributes neuromodulation across ALL C-stage layers, matching how cholinergic projections from the medial septum innervate every level of the hippocampal circuit — not just a single output stage. Zero-initialized so modulation starts as identity (no warmup needed). The model learns to use feedback gradually through training. """ def __init__(self, d_model: int, n_heads: int, ffn_dim: int, dropout: float = 0.1): super().__init__() self.ln1 = nn.LayerNorm(d_model, elementwise_affine=False) self.retention = MultiScalePeriodicRetention(d_model, n_heads, dropout) self.ln2 = nn.LayerNorm(d_model, elementwise_affine=False) self.ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, d_model), nn.Dropout(dropout), ) bottleneck = d_model // 4 self.neuro_proj = nn.Sequential( nn.Linear(d_model, bottleneck), nn.GELU(), nn.Linear(bottleneck, 4 * d_model), ) nn.init.zeros_(self.neuro_proj[2].weight) nn.init.zeros_(self.neuro_proj[2].bias) def forward(self, x: torch.Tensor, feedback: Optional[torch.Tensor] = None) -> torch.Tensor: if feedback is not None: params = self.neuro_proj(feedback) s1, b1, s2, b2 = params.chunk(4, dim=-1) s1 = s1.clamp(-0.5, 0.5) b1 = b1.clamp(-0.5, 0.5) s2 = s2.clamp(-0.5, 0.5) b2 = b2.clamp(-0.5, 0.5) h = (1 + s1) * self.ln1(x) + b1 x = x + self.retention(h) h = (1 + s2) * self.ln2(x) + b2 x = x + self.ffn(h) else: x = x + self.retention(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x class RetentionLayer(nn.Module): """Pre-norm retention layer with FFN.""" def __init__(self, d_model: int, n_heads: int, ffn_dim: int, dropout: float = 0.1): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.retention = MultiScaleRetention(d_model, n_heads, dropout) self.ln2 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, d_model), nn.Dropout(dropout), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.retention(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x class CStageRetention(nn.Module): """C-Stage: grid-cell-inspired periodic retention encoder.""" def __init__(self, vocab: int, max_seq: int, d_model: int, n_heads: int, n_layers: int, ffn_dim: int, dropout: float = 0.1): super().__init__() self.tok_emb = nn.Embedding(vocab, d_model) self.pos_emb = nn.Embedding(max_seq, d_model) self.layers = nn.ModuleList([ PeriodicRetentionLayer(d_model, n_heads, ffn_dim, dropout) for _ in range(n_layers) ]) self.ln = nn.LayerNorm(d_model) self.drop = nn.Dropout(dropout) def forward(self, ids: torch.Tensor, feedback: Optional[torch.Tensor] = None) -> torch.Tensor: B, T = ids.shape pos = torch.arange(T, device=ids.device).unsqueeze(0) x = self.drop(self.tok_emb(ids) + self.pos_emb(pos)) for layer in self.layers: x = layer(x, feedback=feedback) return self.ln(x) class EStage(nn.Module): """DG: expansion + pattern separation via AHSD + CSR. Anti-Hebbian Support Drift (AHSD): Tracks per-neuron activation frequencies and suppresses overused neurons before top-k selection. Prevents mode collapse and ensures full use of the sparse code space. Inspired by homeostatic plasticity and intrinsic excitability regulation that maintain the DG's characteristically low population activity (~2-5% of granule cells active; Chawla et al. 2005). Contrastive Support Repulsion (CSR): Auxiliary loss that penalizes support overlap between similar inputs. The E-stage doesn't just sparsify — it actively maximizes code dissimilarity for similar inputs, which is the computational definition of pattern separation. """ def __init__(self, d_model: int, expansion: int = 4, sparsity: float = 0.10, ema_decay: float = 0.99, inhibition_strength: float = 0.5): super().__init__() self.d_expand = d_model * expansion self.k = max(1, int(self.d_expand * sparsity)) self.expand = nn.Linear(d_model, self.d_expand) self.contract = nn.Linear(self.d_expand, d_model) self.ln = nn.LayerNorm(d_model) self.register_buffer('neuron_freq', torch.zeros(self.d_expand)) self.ema_decay = ema_decay self.inhibition_strength = inhibition_strength self.inhibition_temp = nn.Parameter(torch.tensor(1.0)) def set_sparsity(self, sparsity: float): """Update sparsity level (for annealing during training).""" self.k = max(1, int(self.d_expand * sparsity)) def forward(self, h_C: torch.Tensor, feedback: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: if feedback is not None: h_C = h_C + feedback expanded = F.gelu(self.expand(h_C)) if self.training and self.neuron_freq.sum() > 0: penalty = (self.neuron_freq * self.inhibition_strength * F.softplus(self.inhibition_temp)) inhibited = expanded - penalty else: inhibited = expanded topk_vals, topk_idx = inhibited.topk(self.k, dim=-1) original_vals = expanded.gather(-1, topk_idx) sparse = torch.zeros_like(expanded) sparse.scatter_(-1, topk_idx, original_vals) if self.training: with torch.no_grad(): batch_freq = (sparse != 0).float().mean(dim=(0, 1)) self.neuron_freq.mul_(self.ema_decay).add_( batch_freq, alpha=1 - self.ema_decay) contracted = self.contract(sparse) return self.ln(contracted + h_C), sparse @staticmethod def csr_loss(h_C: torch.Tensor, sparse: torch.Tensor, n_sample: int = 128) -> torch.Tensor: """Contrastive Support Repulsion: penalize support overlap between similar inputs. This IS pattern separation as an explicit objective. Args: h_C: [B, T, d] pre-expansion input sparse: [B, T, d_expand] sparse codes n_sample: subsample tokens for O(n^2) efficiency """ B, T, _ = h_C.shape if T > n_sample: idx = torch.randperm(T, device=h_C.device)[:n_sample] h_sub = h_C[:, idx] s_sub = sparse[:, idx] else: h_sub, s_sub = h_C, sparse n_sample = T h_norm = F.normalize(h_sub, dim=-1) input_sim = F.relu(torch.bmm(h_norm, h_norm.transpose(1, 2))) mask = (s_sub != 0).float() k = mask.sum(dim=-1, keepdim=True).clamp(min=1) mask_norm = mask / k.sqrt() overlap = torch.bmm(mask_norm, mask_norm.transpose(1, 2)) eye = torch.eye(n_sample, device=h_C.device).unsqueeze(0) loss = ((input_sim * overlap) * (1 - eye)).sum() return loss / (B * n_sample * max(n_sample - 1, 1)) class LogitInjectingMHA(nn.Module): """Multi-head cross-attention with optional pre-softmax logit injection. Used by V4d to bypass the lossy query-perturbation geometry of V0/V3m/ V3m2/V4c/V4c-M (where v_vec was added to the query). With this module, v_vec is projected to a per-slot bias `sal_bias [B, T, S]` and added DIRECTLY to the pre-softmax attention logits. Softmax cannot flatten logit-space perturbations — it exponentiates them — so small biases produce measurable retrieval-distribution shifts. Parameter names + shapes MATCH `nn.MultiheadAttention`'s standard parameterization (`in_proj_weight`, `in_proj_bias`, `out_proj.weight`, `out_proj.bias`) so legacy strict-load of an nn.MultiheadAttention state_dict would round-trip into this module (although V4d's gated conditional replacement means we don't rely on this in practice). Forward output is `(out, None)` to mirror `nn.MultiheadAttention`'s `(output, attn_weights)` return shape; callers `_, _ = self.cross_attn( q, k, v)` work unchanged. Side-channel storage (`_last_attn`, `_last_sal_bias`, scalar logit-stat fields) is gated on `store=True` per-call to avoid wasting GPU memory on the [B, H, T, S] attn tensor every forward. Probes and the V4d periodic logger toggle this flag externally before calling. """ def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0, \ f"d_model {d_model} not divisible by n_heads {n_heads}" self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads self.scale = self.d_head ** -0.5 self.in_proj_weight = nn.Parameter(torch.empty(3 * d_model, d_model)) self.in_proj_bias = nn.Parameter(torch.empty(3 * d_model)) self.out_proj = nn.Linear(d_model, d_model, bias=True) nn.init.xavier_uniform_(self.in_proj_weight) nn.init.constant_(self.in_proj_bias, 0.0) self.dropout = nn.Dropout(dropout) self._last_attn = None self._last_sal_bias = None self._last_base_logits_std = None self._last_attn_logits_std = None self._last_sal_logits_std = None def forward(self, q, k, v, sal_bias=None, store: bool = False): B, T, _ = q.shape S = k.shape[1] Wq, Wk, Wv = self.in_proj_weight.chunk(3, dim=0) bq, bk, bv = self.in_proj_bias.chunk(3, dim=0) Q = F.linear(q, Wq, bq).view( B, T, self.n_heads, self.d_head).transpose(1, 2) K = F.linear(k, Wk, bk).view( B, S, self.n_heads, self.d_head).transpose(1, 2) V = F.linear(v, Wv, bv).view( B, S, self.n_heads, self.d_head).transpose(1, 2) attn_logits = (Q @ K.transpose(-1, -2)) * self.scale if sal_bias is not None: attn_logits = attn_logits + sal_bias.unsqueeze(1) attn = F.softmax(attn_logits, dim=-1) if store: self._last_attn = attn.detach() self._last_sal_bias = (sal_bias.detach() if sal_bias is not None else None) with torch.no_grad(): base_logits = (Q @ K.transpose(-1, -2)) * self.scale self._last_base_logits_std = float(base_logits.std().item()) self._last_attn_logits_std = float(attn_logits.std().item()) self._last_sal_logits_std = (float(sal_bias.std().item()) if sal_bias is not None else 0.0) attn = self.dropout(attn) out = (attn @ V).transpose(1, 2).contiguous().view(B, T, self.d_model) out = self.out_proj(out) return out, None class DStage(nn.Module): """CA3: dual-memory — attractor refinement + explicit memory bank + Loop 2. Combines two complementary memory mechanisms, matching real CA3: 1. Attractor refinement (implicit): patterns stored in shared QKV weights, retrieved via iterative self-attention settling to fixed points. Analog: CA3 recurrent collateral autoassociation. 2. Memory bank (explicit): learned key-value slots accessed via cross-attention after attractor settling. Analog: CA3 specific learned synaptic associations. The attractor completes patterns from partial cues (implicit recall). The memory bank stores specific associations (explicit lookup). A learned gate integrates both sources. """ def __init__(self, d_model: int, d_expand: int, n_heads: int, refine_steps: int = 3, num_slots: int = 256, dropout: float = 0.1, use_salience: bool = True, salience_mode: str = "v0", salience_bias_init_sigma: float = 0.0, v4c_norm_cap: float = 0.3, v4d_w_sal_sigma: float = 0.02, v4d_logit_cap: float = 4.0, v4d_noinject: bool = False, v4d_w_sal_rank: int = 0): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads self.refine_steps = refine_steps self.scale = self.d_head ** -0.5 self.use_salience = use_salience assert salience_mode in ("v0", "v2a", "v3", "v3m", "v3m2", "v3m2_nll", "v4c", "v4c_gate", "v4d"), \ f"salience_mode must be one of 'v0'/'v2a'/'v3'/'v3m'/'v3m2'/" \ f"'v3m2_nll'/'v4c'/'v4c_gate'/'v4d', got {salience_mode!r}" self.salience_mode = salience_mode self._salience_bias_init_sigma = salience_bias_init_sigma self.v4c_norm_cap = v4c_norm_cap self.num_slots = num_slots self.v4d_logit_cap = v4d_logit_cap self.v4d_noinject = v4d_noinject self.v4d_w_sal_rank = v4d_w_sal_rank self.store_v4d_attn = False self.store_v4d_code = False self.store_v6_bank_sources = False self._last_q_attractor_v6 = None self._last_q_mem_v6 = None self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_out = nn.Linear(d_model, d_model) self.step_sizes = nn.ParameterList([ nn.Parameter(torch.tensor(0.1)) for _ in range(refine_steps)]) self.step_lns = nn.ModuleList([ nn.LayerNorm(d_model) for _ in range(refine_steps)]) self.dropout = nn.Dropout(dropout) self.mem_keys = nn.Parameter(torch.randn(num_slots, d_model) * 0.02) self.mem_vals = nn.Parameter(torch.randn(num_slots, d_model) * 0.02) if salience_mode == "v4d": self.cross_attn = LogitInjectingMHA( d_model, n_heads, dropout=dropout) else: self.cross_attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout, batch_first=True) self.ln_mem = nn.LayerNorm(d_model) self.mem_gate_linear = nn.Linear(2 * d_model, d_model) self.loop2_q_proj = nn.Linear(d_model, d_model) self.loop2_k_proj = nn.Linear(d_expand, d_model) self.loop2_v_proj = nn.Linear(d_expand, d_model) self.loop2_attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout, batch_first=True) self.loop2_gate = nn.Sequential( nn.Linear(2 * d_model, d_model), nn.Sigmoid(), ) self.ln_loop2 = nn.LayerNorm(d_model) if use_salience: with torch.random.fork_rng(devices=[]): self.salience_bias_down = nn.Linear(d_model, 16, bias=False) self.salience_bias_up = nn.Linear(16, d_model, bias=False) if self._salience_bias_init_sigma > 0.0: nn.init.normal_(self.salience_bias_up.weight, std=self._salience_bias_init_sigma) else: nn.init.zeros_(self.salience_bias_up.weight) if salience_mode == "v4d": if v4d_w_sal_rank and v4d_w_sal_rank > 0: k = v4d_w_sal_rank self.w_sal_A = nn.Linear(d_model, k, bias=False) self.w_sal_B = nn.Linear(k, num_slots, bias=False) nn.init.normal_(self.w_sal_A.weight, std=(1.0 / d_model) ** 0.5) nn.init.normal_(self.w_sal_B.weight, std=v4d_w_sal_sigma * (d_model / k) ** 0.5) else: self.w_sal = nn.Linear(d_model, num_slots, bias=False) nn.init.normal_(self.w_sal.weight, std=v4d_w_sal_sigma) def attractor_settle(self, h: torch.Tensor, causal: torch.Tensor) -> torch.Tensor: """Iterative attractor settling via self-attention with learned step sizes.""" B, T, D = h.shape q_state = h for step in range(self.refine_steps): Q = self.w_q(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2) K = self.w_k(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2) V = self.w_v(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2) attn = (Q @ K.transpose(-1, -2)) * self.scale attn = attn.masked_fill(causal.unsqueeze(0).unsqueeze(0), -1e9) attn = self.dropout(F.softmax(attn, dim=-1)) attn_out = (attn @ V).transpose(1, 2).contiguous().view(B, T, D) attn_out = self.w_out(attn_out) eta = torch.sigmoid(self.step_sizes[step]) q_state = self.step_lns[step](q_state + eta * attn_out) return q_state def forward(self, h_E: torch.Tensor, h_E_sparse: torch.Tensor, v_vec: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: B, T = h_E.size(0), h_E.size(1) causal = torch.triu( torch.ones(T, T, device=h_E.device, dtype=torch.bool), 1) q_attractor = self.attractor_settle(h_E, causal) if self.store_v6_bank_sources: self._last_q_attractor_v6 = q_attractor mk = self.mem_keys.unsqueeze(0).expand(B, -1, -1) mv = self.mem_vals.unsqueeze(0).expand(B, -1, -1) if (v_vec is not None and self.use_salience and self.salience_mode in ("v2a", "v4c")): bias_raw = self.salience_bias_up(self.salience_bias_down(v_vec)) if self.salience_mode == "v4c": qa_norm = q_attractor.detach().norm(dim=-1, keepdim=True) bias_norm = bias_raw.norm(dim=-1, keepdim=True).clamp(min=1e-6) scale_factor = (qa_norm * self.v4c_norm_cap / bias_norm ).clamp(max=1.0) bias_raw = bias_raw * scale_factor q_for_retrieval = q_attractor + bias_raw retrieved, _ = self.cross_attn(q_for_retrieval, mk, mv) elif (v_vec is not None and self.use_salience and self.salience_mode == "v4d"): if not self.v4d_noinject: if getattr(self, "v4d_w_sal_rank", 0) and hasattr(self, "w_sal_A"): g = self.w_sal_A(v_vec) if self.store_v4d_code: self._last_sal_code = g sal_bias = self.w_sal_B(g) else: sal_bias = self.w_sal(v_vec) if self.store_v4d_code: self._last_sal_code = v_vec sal_bias = sal_bias - sal_bias.mean(dim=-1, keepdim=True) sal_bias = sal_bias.clamp(-self.v4d_logit_cap, +self.v4d_logit_cap) else: sal_bias = None retrieved, _ = self.cross_attn(q_attractor, mk, mv, sal_bias=sal_bias, store=self.store_v4d_attn) if self.store_v4d_attn: self._last_cross_attn = self.cross_attn._last_attn self._last_sal_bias = self.cross_attn._last_sal_bias else: retrieved, _ = self.cross_attn(q_attractor, mk, mv) q_mem = self.ln_mem(q_attractor + retrieved) if self.salience_mode == "v4d" and self.store_v4d_attn: self._last_q_mem = q_mem.detach() if self.store_v6_bank_sources: self._last_q_mem_v6 = q_mem gate_logits = self.mem_gate_linear(torch.cat([q_attractor, q_mem], dim=-1)) if (v_vec is not None and self.use_salience and self.salience_mode in ("v0", "v3", "v3m", "v3m2", "v3m2_nll", "v4c_gate")): gate_bias = self.salience_bias_up(self.salience_bias_down(v_vec)) if self.salience_mode == "v4c_gate": gl_norm = gate_logits.detach().norm(dim=-1, keepdim=True) gb_norm = gate_bias.norm(dim=-1, keepdim=True).clamp(min=1e-6) gate_bias = gate_bias * (gl_norm * self.v4c_norm_cap / gb_norm).clamp(max=1.0) gate_logits = gate_logits + gate_bias gate = torch.sigmoid(gate_logits) q = gate * q_mem + (1 - gate) * q_attractor if self.salience_mode == "v4d" and self.store_v4d_attn: self._last_gate_logits = gate_logits.detach() loop2_q = self.loop2_q_proj(q) loop2_k = self.loop2_k_proj(h_E_sparse) loop2_v = self.loop2_v_proj(h_E_sparse) loop2_out, _ = self.loop2_attn( loop2_q, loop2_k, loop2_v, attn_mask=causal) gate = self.loop2_gate(torch.cat([q, loop2_out], dim=-1)) h_D = self.ln_loop2(q + gate * loop2_out) return h_D, loop2_out class LStage(nn.Module): """CA1: two-channel comparator.""" def __init__(self, d_model: int, vocab: int, n_heads: int, dropout: float = 0.1): super().__init__() self.mem_head = nn.Linear(d_model, vocab) self.per_head = nn.Linear(d_model, vocab, bias=False) self.gate_net = nn.Sequential( nn.Linear(2 * d_model, 128), nn.GELU(), nn.Linear(128, 1), nn.Sigmoid(), ) self.fuse_gate = nn.Sequential( nn.Linear(2 * d_model, d_model), nn.Sigmoid(), ) self.loop1_attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout, batch_first=True) self.loop1_gate = nn.Sequential( nn.Linear(2 * d_model, d_model), nn.Sigmoid(), ) self.ln_loop1 = nn.LayerNorm(d_model) def forward(self, h_D: torch.Tensor, h_C: torch.Tensor): logits_mem = self.mem_head(h_D) g = self.fuse_gate(torch.cat([h_D, h_C], dim=-1)) fused = g * h_D + (1 - g) * h_C T = h_C.size(1) causal = torch.triu(torch.ones(T, T, device=h_C.device, dtype=torch.bool), 1) attn_out, _ = self.loop1_attn(fused, h_C, h_C, attn_mask=causal) gate = self.loop1_gate(torch.cat([fused, attn_out], dim=-1)) loop1_fb = self.ln_loop1(fused + gate * attn_out) return logits_mem, loop1_fb class NeuromodulatorGate(nn.Module): """Output-level gain modulation of C-stage by L and D signals. Distinct from the per-layer cholinergic AdaLN in PeriodicRetentionLayer: this gate models dopaminergic/noradrenergic modulation that gates the final hippocampal output (Lisman & Grace 2005), while AdaLN models cholinergic modulation during encoding (Hasselmo 1999). """ def __init__(self, d_model: int): super().__init__() self.gain_from_L = nn.Sequential(nn.Linear(d_model, d_model), nn.Sigmoid()) self.gain_from_D = nn.Sequential(nn.Linear(d_model, d_model), nn.Sigmoid()) self.bias_from_L = nn.Linear(d_model, d_model) self.ln = nn.LayerNorm(d_model) def forward(self, h_C, loop1_fb, h_D): g_L = self.gain_from_L(loop1_fb) g_D = self.gain_from_D(h_D) bias = self.bias_from_L(loop1_fb) return self.ln(g_L * g_D * h_C + bias) DISCOURSE_MARKERS = [ "however", "but", "although", "nevertheless", "despite", "instead", "whereas", "conversely", "previously", "currently", "originally", "initially", "formerly", "recently", "now", ] def build_marker_token_ids() -> torch.Tensor: """Tokenize the discourse-marker lexicon with GPT-2 BPE. Returns a sorted 1-D LongTensor of unique FIRST-token IDs (sub-word continuations are excluded; the lead BPE chunk carries the marker signal). Returns an empty tensor if transformers is unavailable so the marker BCE loss degrades gracefully.""" try: from transformers import GPT2TokenizerFast tok = GPT2TokenizerFast.from_pretrained("gpt2") except Exception: return torch.empty(0, dtype=torch.long) ids = set() for w in DISCOURSE_MARKERS: for form in (w, " " + w, w.capitalize(), " " + w.capitalize()): toks = tok.encode(form, add_special_tokens=False) if toks: ids.add(toks[0]) return torch.tensor(sorted(ids), dtype=torch.long) def build_marker_mask(ids: torch.Tensor, marker_ids: torch.Tensor) -> torch.Tensor: """[B, T] -> [B, T] float mask, 1 at any position within a +/-1-token window of a marker. Smooths supervision and matches phrase-boundary semantics of "however,", " but ", etc.""" if marker_ids.numel() == 0: return torch.zeros_like(ids, dtype=torch.float) in_set = torch.isin(ids, marker_ids).float().unsqueeze(1) kernel = torch.ones(1, 1, 3, device=ids.device, dtype=in_set.dtype) dilated = F.conv1d(in_set, kernel, padding=1).squeeze(1) return (dilated > 0).float() class SalienceLoop(nn.Module): """Causal, low-rank salience scorer + contrast detector + d->r->d gate. Computed from h_E (E-stage output) so it can bias D-stage memory fusion before final D output is committed. (SL injects v_vec into the mem_gate logits — the blend of attractor output and memory-bank retrieval — which runs AFTER attractor_settle(). It does not modify the attractor settling iterations themselves.) Modality-agnostic by construction. Outputs: v_vec [B, T, d] vector signal for D-stage mem-fusion bias v_scalar [B, T, 1] scalar signal for aux losses (NLL-surprise + marker BCE) """ def __init__(self, d: int = 640, H: int = 8, kernel: int = 5, gate_rank: int = 16, scalar_from_contrast: bool = False, contrastive_head: bool = False, z_dim: int = 64, proj_hidden: int = 256): super().__init__() self._d = d self.U = nn.Parameter(torch.randn(H, d) * (d ** -0.5)) self.b = nn.Parameter(torch.zeros(H)) self.contrast = nn.Conv1d(H, H, kernel_size=kernel, groups=H, padding=0, bias=False) nn.init.normal_(self.contrast.weight, std=0.1) self.pad = kernel - 1 self.proj = nn.Linear(2 * H, d, bias=True) self.ln = nn.LayerNorm(d) self.gate_down = nn.Linear(d, gate_rank, bias=False) self.gate_up = nn.Linear(gate_rank, d, bias=False) self.scale = nn.Parameter(torch.tensor(0.1)) self.scalar_from_contrast = scalar_from_contrast if scalar_from_contrast: self.scalar_head = nn.Linear(2 * H, 1, bias=True) nn.init.constant_(self.scalar_head.bias, -2.0) self.contrastive_head = contrastive_head if contrastive_head: self.sl_embed_head = nn.Sequential( nn.Linear(d, proj_hidden), nn.GELU(), nn.Linear(proj_hidden, z_dim), ) self._z_log_inv_temperature = nn.Parameter( torch.tensor(math.log(1.0 / 0.05))) def forward(self, h_E: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: s = F.softplus(h_E @ self.U.t() + self.b) s_c = F.pad(s.transpose(1, 2), (self.pad, 0)) c = self.contrast(s_c).transpose(1, 2) v_feat = torch.cat([s, c.abs()], dim=-1) v_raw = self.ln(self.proj(v_feat)) v_gate = self.gate_up(self.gate_down(v_raw)) v_vec = torch.tanh(self.scale) * v_gate if self.scalar_from_contrast: v_scalar = self.scalar_head(v_feat) else: v_scalar = s.mean(-1, keepdim=True) return v_vec, v_scalar def embed_z(self, h_D: torch.Tensor) -> torch.Tensor: """V4c-only: compute the L2-normalized contrastive embedding from POST-DStage state h_D. Used by `_v4c_contrastive_loss` to supervise a contrastive manifold whose gradient backprops through h_D -> v_vec -> salience_bias_up -> everything upstream. Args: h_D [B, T, d] Returns: z [B, T, z_dim], with ||z[..., :]||_2 = 1 per row. """ assert self.contrastive_head, "embed_z called on a SalienceLoop " \ "built without contrastive_head=True (V4c-only API)." return F.normalize(self.sl_embed_head(h_D), dim=-1) def anchor_ortho_reg(U: torch.Tensor) -> torch.Tensor: """Anchor orthogonality regularizer: keeps the H salience anchors pointing in distinct directions instead of collapsing to redundancy. """ G = U @ U.t() I = torch.eye(U.size(0), device=U.device, dtype=U.dtype) return ((G - I) ** 2).mean() def _ema_zscore(x: torch.Tensor, ema_mean_buf: torch.Tensor, ema_var_buf: torch.Tensor, init_buf: torch.Tensor, decay: float, training: bool) -> torch.Tensor: """Initialize-from-first-batch EMA z-score for V3/V3m. EMA updates ONLY when training=True. The first observed batch initializes the EMA buffers from its own mean/var (sentinel `init_buf` < 0.5) so we avoid the start-at-0 vs NLL-around-6 explosion when aux loss first activates. Subsequent batches update with `decay` (≈0.99 → ~100-batch effective window). Returns the z-scored tensor (not clamped — callers clamp to ±sl_z_clip before combining channels). """ if training: with torch.no_grad(): if init_buf.item() < 0.5: ema_mean_buf.copy_(x.mean()) ema_var_buf.copy_(x.var().clamp(min=1e-6)) init_buf.fill_(1.0) else: ema_mean_buf.mul_(decay).add_((1 - decay) * x.mean()) ema_var_buf.mul_(decay).add_((1 - decay) * x.var()) return (x - ema_mean_buf) / (ema_var_buf.sqrt() + 1e-6) def sigreg_loss(z: torch.Tensor) -> torch.Tensor: """Soft variance regularizer for D-stage outputs. Adapted from LeWorldModel's SIGReg (arXiv:2603.19312), but softened: only penalizes variance COLLAPSE (< 0.1), not high variance. This prevents representational collapse without fighting attractor basin structure (which is inherently multimodal/high-variance). """ if z.dim() == 3: z = z.reshape(-1, z.size(-1)) var = z.var(dim=0) return F.relu(0.1 - var).mean() def predictive_latent_loss(h_D: torch.Tensor, h_C: torch.Tensor) -> torch.Tensor: """JEPA-inspired predictive loss: D-stage output at position t should predict C-stage encoding at position t+1 in representation space. This makes the D-stage a forward model (not just a pattern completer), directly implementing the hippocampal predictive coding hypothesis: CA3 generates a prediction via Schaffer collaterals, CA1 compares it against actual EC input (Hasselmo 2005, Lisman & Redish 2009). Stop-gradient on target prevents representational collapse (JEPA principle). """ pred = h_D[:, :-1, :] target = h_C[:, 1:, :].detach() pred_n = F.normalize(pred, dim=-1) tgt_n = F.normalize(target, dim=-1) return F.smooth_l1_loss(pred_n, tgt_n) V3M2_ANCHOR_NEGATIVE_WORDS = [ "<|endoftext|>", ] V3M2_LABEL_NEGATIVE_PUNCT = [ ".", ",", ";", ":", "!", "?", "—", "-", '"', "'", "(", ")", "[", "]", "{", "}", "/", "\\", "`", "*", "_", "~", "|", "<", ">", "=", "+", "&", "%", "$", "#", "@", ] V3M2_LABEL_NEGATIVE_WORDS = [ "a", "an", "the", "is", "was", "were", "be", "been", "being", "am", "are", "have", "has", "had", "do", "does", "did", "will", "would", "shall", "should", "may", "might", "can", "could", "must", "ought", "he", "she", "it", "they", "we", "you", "I", "me", "him", "her", "us", "them", "his", "hers", "its", "their", "our", "your", "my", "of", "in", "on", "at", "to", "for", "with", "by", "from", "into", "onto", "upon", "about", "over", "under", "between", "through", ] def _bpe_token_ids_for_forms(words_or_punct, tokenizer): """For each word/punctuation entry, return the FIRST GPT-2 BPE token id for both bare and leading-space forms (and capitalization variants for words). Returns sorted unique list of int ids.""" ids = set() for entry in words_or_punct: forms = [entry, " " + entry] if entry and entry[0].isalpha(): forms.append(entry.capitalize()) forms.append(" " + entry.capitalize()) for f in forms: try: toks = tokenizer.encode(f, add_special_tokens=False) except Exception: continue if toks: ids.add(int(toks[0])) return sorted(ids) def build_v3m2_lexicons(): """Return (anchor_negative_ids_tensor, label_negative_ids_tensor) as sorted 1-D LongTensors. Built once at model __init__ from a hand-curated list (committed to the repo above); empty tensors if transformers is unavailable so the masks degrade to "allow everything".""" try: from transformers import GPT2TokenizerFast tok = GPT2TokenizerFast.from_pretrained("gpt2") except Exception: return (torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long)) anchor_ids = set() eot = tok.encode("<|endoftext|>", add_special_tokens=False) for t in eot: anchor_ids.add(int(t)) try: space_id = tok.encode(" ", add_special_tokens=False) for t in space_id: anchor_ids.add(int(t)) except Exception: pass label_ids = set() for tok_id in _bpe_token_ids_for_forms(V3M2_LABEL_NEGATIVE_PUNCT, tok): label_ids.add(tok_id) for tok_id in _bpe_token_ids_for_forms(V3M2_LABEL_NEGATIVE_WORDS, tok): label_ids.add(tok_id) anchor_t = torch.tensor(sorted(anchor_ids), dtype=torch.long) label_t = torch.tensor(sorted(label_ids), dtype=torch.long) return anchor_t, label_t def _curvature_torch(h): """[B, T, D] -> [B, T-1] per-token curvature 1 - cos(d1, d2), padded with one leading 0 to align with downstream [:, :-1] slicing. h is detached + float internally; caller passes Pass-0 tensors only.""" h_n = F.normalize(h.detach().float(), dim=-1) d1 = F.normalize(h_n[:, 1:-1] - h_n[:, :-2], dim=-1) d2 = F.normalize(h_n[:, 2:] - h_n[:, 1:-1], dim=-1) c = 1.0 - (d1 * d2).sum(-1) return F.pad(c, (1, 0)) def _topk_nms_with_threshold(scores, max_anchors, z_threshold, min_gap): """Per-row top-K with NMS min-gap and absolute threshold. Routine rows (no scores above threshold) get zero anchors. Returns bool [B, T]. Sync-conscious: pulls sorted scores + indices to CPU ONCE per call so the inner double loop runs as pure Python over numpy arrays rather than one GPU sync per `.item()`. Single CPU->GPU sync at the end to set the bool mask. NMS itself is inherently sequential (per-row top-K with a min-gap exclusion) so loops are unavoidable, but they are no longer per-iter GPU syncs.""" B, T = scores.shape sorted_scores, sorted_idx = scores.sort(dim=-1, descending=True) sorted_scores_cpu = sorted_scores.detach().cpu().numpy() sorted_idx_cpu = sorted_idx.detach().cpu().numpy() rows, cols = [], [] NEG_INF = float('-inf') for b in range(B): accepted = [] for j in range(T): s = float(sorted_scores_cpu[b, j]) if s == NEG_INF or s < z_threshold: break pos = int(sorted_idx_cpu[b, j]) ok = True for ap in accepted: if abs(pos - ap) < min_gap: ok = False break if ok: accepted.append(pos) rows.append(b); cols.append(pos) if len(accepted) >= max_anchors: break out = torch.zeros_like(scores, dtype=torch.bool) if rows: out[torch.as_tensor(rows, device=scores.device, dtype=torch.long), torch.as_tensor(cols, device=scores.device, dtype=torch.long)] = True return out def _build_decay_target_tensorized(anchors, strength, is_label_eligible, window=12, tau=2.0, k_max=6): """Build per-position graded labels via exponential decay over the content-window AFTER each anchor. Args: anchors: [B, T] bool, True at anchor positions strength: [B, T] float in [0, 1], anchor strength is_label_eligible: [B, T] bool, True for content tokens window: int, max position offset to consider after anchor tau: float, decay constant in content-rank units k_max: int, cap on content rank assigned a positive label Returns: target [B, T] float; target[t] = max over anchors of strength[t_a] * exp(-(content_rank - 1) / tau), for label- eligible positions within `window` of t_a and within k_max content-eligible positions of t_a (zero otherwise). """ B, T = anchors.shape device = anchors.device cum_le = is_label_eligible.long().cumsum(dim=-1) anchor_counts = anchors.sum(dim=-1) A = int(anchor_counts.max().item()) if anchor_counts.max() > 0 else 0 if A == 0: return torch.zeros(B, T, device=device, dtype=strength.dtype) anchor_pos = torch.full((B, A), -1, dtype=torch.long, device=device) for b in range(B): idx = anchors[b].nonzero(as_tuple=True)[0] anchor_pos[b, :idx.size(0)] = idx offsets = torch.arange(1, window + 1, device=device).view(1, 1, window) cand_pos = anchor_pos.unsqueeze(-1) + offsets cand_valid = (anchor_pos.unsqueeze(-1) >= 0) & (cand_pos < T) cand_pos_clamped = cand_pos.clamp(min=0, max=T - 1) le_at_cand = is_label_eligible.gather( 1, cand_pos_clamped.view(B, -1) ).view(B, A, window) cum_at_cand = cum_le.gather( 1, cand_pos_clamped.view(B, -1) ).view(B, A, window) anchor_pos_clamped = anchor_pos.clamp(min=0) cum_at_anchor = cum_le.gather(1, anchor_pos_clamped) content_rank = cum_at_cand - cum_at_anchor.unsqueeze(-1) valid_label = cand_valid & le_at_cand & (content_rank >= 1) & (content_rank <= k_max) rank_for_decay = (content_rank - 1).clamp(min=0).float() decay = torch.exp(-rank_for_decay / tau) label_vals = strength.gather(1, anchor_pos_clamped).unsqueeze(-1) * decay label_vals = label_vals * valid_label.float() target = torch.zeros(B, T, device=device, dtype=strength.dtype) flat_pos = cand_pos_clamped.view(B, -1) flat_val = label_vals.view(B, -1) target = target.scatter_reduce( dim=1, index=flat_pos, src=flat_val, reduce="amax", include_self=True, ) return target def _build_v3m2_targets( *, ids, h_C_p0, h_E_p0, logits_mem_p0, model, salience_mode, update_ema, with_nll, ): """Build the V3m2 supervision target. Callable from training (update_ema =True) and from --preview-only (update_ema=False). Returns (target, w, debug_info) where target is [B, T-1] in [0, 1], w is per-token weight [B, T-1] normalized to mean 1, and debug_info is a dict with intermediate tensors useful for inspection. """ decay = model.nll_ema_decay curve_C = _curvature_torch(h_C_p0) curve_E = _curvature_torch(h_E_p0) h_C_n = F.normalize(h_C_p0.detach().float(), dim=-1) h_E_n = F.normalize(h_E_p0.detach().float(), dim=-1) cosmis = (1.0 - (h_C_n * h_E_n).sum(-1))[:, :-1] def _z(x, m_buf, v_buf, i_buf): if update_ema: return _ema_zscore(x, m_buf, v_buf, i_buf, decay, training=True ).clamp(-4.0, 4.0) return ((x - m_buf) / (v_buf.sqrt() + 1e-6)).clamp(-4.0, 4.0) curve_C_z = _z(curve_C, model.curve_C_ema_mean, model.curve_C_ema_var, model.curve_C_ema_initialized) curve_E_z = _z(curve_E, model.curve_E_ema_mean, model.curve_E_ema_var, model.curve_E_ema_initialized) cosmis_z = _z(cosmis, model.cosmis_ema_mean, model.cosmis_ema_var, model.cosmis_ema_initialized) event_score = 0.5 * curve_C_z + 0.5 * curve_E_z + 0.7 * cosmis_z if with_nll and logits_mem_p0 is not None: flat_logits = logits_mem_p0[:, :-1].reshape( -1, logits_mem_p0.size(-1)).float() nll = F.cross_entropy( flat_logits, ids[:, 1:].reshape(-1), reduction='none' ).view(ids.size(0), -1) nll_z = _z(nll, model.nll_ema_mean, model.nll_ema_var, model.nll_ema_initialized) event_score = event_score + 0.7 * nll_z ANCHOR_Z = 1.5 TEMP = 0.75 strength = torch.sigmoid((event_score - ANCHOR_Z) / TEMP) ids_for_mask = ids[:, :-1] is_anchor_eligible = ~torch.isin(ids_for_mask, model.anchor_negative_ids) is_label_eligible = ~torch.isin(ids_for_mask, model.label_negative_ids) es_for_anchor = event_score.masked_fill( ~is_anchor_eligible, float('-inf')) max_anchors = min(8, max(1, ids.size(1) // 128)) anchors = _topk_nms_with_threshold( es_for_anchor, max_anchors=max_anchors, z_threshold=ANCHOR_Z, min_gap=6, ) target = _build_decay_target_tensorized( anchors=anchors, strength=strength, is_label_eligible=is_label_eligible, window=12, tau=2.0, k_max=6, ) w = 1.0 + 4.0 * target.detach() w = w / w.mean().clamp(min=1e-6) debug = { "event_score": event_score, "strength": strength, "is_anchor_eligible": is_anchor_eligible, "is_label_eligible": is_label_eligible, "anchors": anchors, "target": target, } return target, w, debug def _pool_span(z: torch.Tensor, span: Tuple[int, int]) -> torch.Tensor: """Mean-pool then L2-normalize z over (start, end) — used to extract a single span vector for InfoNCE. z is [T, z_dim]; returns [z_dim].""" s, e = span if e <= s: return None pooled = z[s:e].mean(dim=0) return F.normalize(pooled, dim=-1) def _v4c_contrastive_loss(*, z, v_scalar, spans, cfg, salience_module, bce_alpha, nce_alpha, randomize_positive_mapping: bool = False): """Compute the V4c contrastive aux loss. InfoNCE on `z` pooled over (query, current/positive, stale, distractor, neutral) spans with weighted negatives + in-batch false-negative masking. Span-level BCE on `v_scalar` (positive mask = current span; negative mask = stale ∪ distractor ∪ neutral; other tokens ignored). Args: z: [B, T, z_dim] L2-normalized contrastive embedding. v_scalar: [B, T, 1] LOGIT from SalienceLoop.scalar_head. spans: list of len-B V4cItem-like dicts (or objects). cfg: CEDLConfig. salience_module: the SalienceLoop instance. bce_alpha, nce_alpha: ramp buffers in [0, 1]. randomize_positive_mapping: V4c-randlabel control. When True, each anchor's positive is REPLACED by a different anchorable item's pooled `current` z (cross-item rotation). BCE is also DISABLED in this mode — the per-token BCE target is meaningless when the positive comes from another sequence. This is the proper random- label control: same prompt distribution, but the (anchor → positive) mapping is shuffled, destroying the contrastive signal-to-supervision link. Returns: scalar tensor. """ device = z.device B = z.size(0) inv_temp = salience_module._z_log_inv_temperature.exp().clamp(min=10.0, max=100.0) neg_weights = cfg.v4c_neg_weights anchors = [] positives = [] item_neg_pool = [] global_neg_pool = [] item_in_batch_currents = [] for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") ids_b = item.ids if hasattr(item, "ids") else item.get("ids") if family == "neutral_control": for sp in (item.neutral if hasattr(item, "neutral") else item.get("neutral", [])): v = _pool_span(z[b], sp) if v is not None: global_neg_pool.append((b, v, neg_weights.get("neutral", 1.0), "neutral")) continue q_spans = item.query if hasattr(item, "query") else item.get("query", []) c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", []) n_spans = item.neutral if hasattr(item, "neutral") else item.get("neutral", []) pp_spans = (item.paraphrase_positives if hasattr(item, "paraphrase_positives") else item.get("paraphrase_positives", [])) if not q_spans or not c_spans: continue a = _pool_span(z[b], q_spans[0]) if a is None: continue anchors.append((b, a)) pos_for_b = [] cv = _pool_span(z[b], c_spans[0]) if cv is not None: pos_for_b.append(cv) for pp in pp_spans: ppv = _pool_span(z[b], pp) if ppv is not None: pos_for_b.append(ppv) positives.append((b, pos_for_b)) for sp in s_spans: v = _pool_span(z[b], sp) if v is not None: item_neg_pool.append((b, v, neg_weights.get("stale", 2.0), "stale")) for sp in d_spans: v = _pool_span(z[b], sp) if v is not None: item_neg_pool.append((b, v, neg_weights.get("distractor", 1.0), "distractor")) for sp in n_spans: v = _pool_span(z[b], sp) if v is not None: item_neg_pool.append((b, v, neg_weights.get("neutral", 1.0), "neutral")) current_tok_ids = frozenset(ids_b[c_spans[0][0]:c_spans[0][1]]) query_tok_ids = frozenset(ids_b[q_spans[0][0]:q_spans[0][1]]) if pos_for_b: item_in_batch_currents.append( (b, pos_for_b[0], current_tok_ids, query_tok_ids)) if randomize_positive_mapping and len(positives) > 1: rotated = positives[1:] + positives[:1] positives_for_loss = [(b, plist) for (_, _), (b, plist) in zip(anchors, rotated)] else: positives_for_loss = positives L_nce = z.new_zeros(()) n_anchored = 0 for (b, a), (b_pos, pos_list) in zip(anchors, positives_for_loss): if not pos_list: continue candidates = [] for p in pos_list: candidates.append((p, 1.0, True)) for (b_n, v, w, kind) in item_neg_pool: if b_n == b: candidates.append((v, w, False)) for (b_origin, v, w, kind) in global_neg_pool: candidates.append((v, w, False)) own_current_ids = None own_query_ids = None for (bb, _v, cids, qids) in item_in_batch_currents: if bb == b: own_current_ids = cids own_query_ids = qids break for (bb, v_other, c_ids_other, q_ids_other) in item_in_batch_currents: if bb == b: continue if randomize_positive_mapping and bb == b_pos: continue if own_current_ids is not None: if c_ids_other & own_current_ids: continue if q_ids_other == own_query_ids: continue candidates.append((v_other, 1.0, False)) if len(candidates) < 2: continue cand_vecs = torch.stack([c[0] for c in candidates], dim=0) cand_weights = torch.tensor([c[1] for c in candidates], device=device, dtype=z.dtype) is_pos = torch.tensor([c[2] for c in candidates], device=device, dtype=torch.bool) logits = (cand_vecs @ a) * inv_temp max_l = logits.max().detach() log_denom = (cand_weights * (logits - max_l).exp()).sum().clamp(min=1e-12).log() + max_l pos_mask = is_pos.float() pos_count = pos_mask.sum().clamp(min=1.0) L_i = -((logits - log_denom) * pos_mask).sum() / pos_count L_nce = L_nce + L_i n_anchored += 1 if n_anchored > 0: L_nce = L_nce / n_anchored else: L_nce = z.sum() * 0.0 if randomize_positive_mapping: L_bce = z.sum() * 0.0 else: v_logit = v_scalar[:, :, 0] target_mask = torch.zeros_like(v_logit) target_value = torch.zeros_like(v_logit) for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") if family == "neutral_control": for sp in (item.neutral if hasattr(item, "neutral") else item.get("neutral", [])): s, e = sp target_mask[b, s:e] = 1.0 target_value[b, s:e] = 0.0 continue c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", []) pp_spans = (item.paraphrase_positives if hasattr(item, "paraphrase_positives") else item.get("paraphrase_positives", [])) for s, e in c_spans + pp_spans: target_mask[b, s:e] = 1.0 target_value[b, s:e] = 1.0 for s, e in s_spans + d_spans: target_mask[b, s:e] = 1.0 target_value[b, s:e] = 0.0 bce = F.binary_cross_entropy_with_logits( v_logit, target_value, reduction='none') L_bce = (bce * target_mask).sum() / target_mask.sum().clamp(min=1.0) nce_w = float(getattr(cfg, "v4c_nce_weight", 0.05)) bce_w = float(getattr(cfg, "v4c_bce_weight", 0.02)) return nce_alpha * nce_w * L_nce + bce_alpha * bce_w * L_bce def _v4d_compute_per_item_margins(logits_mem, ids, spans): """Compute per-item answer margins (log_p_cur − log_p_stale) at the last prompt position. Returns: margins: [N_anchorable] tensor — one margin per item that passed all validity checks (non-neutral, has spans, valid length, current_tok != stale_tok). May be empty if all skipped. item_idxs: list[int] — the batch index `b` of each kept item, so the caller can ALIGN margins across two forwards (normal + clamped) by checking that item_idxs match. Used by both _v4c_answer_margin_loss (hinge over each margin) and the V4d-Causal loss (per-item margin gap between normal/clamped forwards). """ B = logits_mem.size(0) margins, item_idxs = [], [] for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") if family == "neutral_control": continue c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) ids_b = item.ids if hasattr(item, "ids") else item.get("ids") if not c_spans or not s_spans: continue orig_len = (item.original_length if hasattr(item, "original_length") else item.get("original_length", 0)) if orig_len <= 0 or orig_len > logits_mem.size(1): continue last_pos = orig_len - 1 cur_tok = ids_b[c_spans[0][0]] stale_tok = ids_b[s_spans[0][0]] if cur_tok == stale_tok: continue logp = F.log_softmax(logits_mem[b, last_pos].float(), dim=-1) margins.append(logp[cur_tok] - logp[stale_tok]) item_idxs.append(b) if margins: return torch.stack(margins), item_idxs return logits_mem.new_zeros((0,)), [] def _v4d_compute_per_item_z_margins(z, spans): """Compute per-item z-cosine margins matching the held-out probe's formula: m_z = cos(z_q, z_current) − max(cos(z_q, z_s), cos(z_q, z_d), cos(z_q, z_n)) where the neutral span is the first 8 tokens of the prompt (probe convention). z is [B, T, z_dim], already L2-normalized per row (SalienceLoop.embed_z output). For each anchorable item, pool z over spans, re-normalize, take cosines, return the margin tensor + item indices. Returns: margins: [N_anchorable] tensor of m_z per item item_idxs: list[int] of batch indices aligned with margins (so two calls under normal/clamped forwards can be aligned). """ B = z.size(0) margins, item_idxs = [], [] def _pool(z_seq, span): s, e = span if e <= s: return None return F.normalize(z_seq[s:e].mean(dim=0), dim=-1) for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") if family == "neutral_control": continue q_spans = item.query if hasattr(item, "query") else item.get("query", []) c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", []) if not q_spans or not c_spans: continue z_seq = z[b] T = z_seq.size(0) z_q = _pool(z_seq, q_spans[0]) z_c = _pool(z_seq, c_spans[0]) if z_q is None or z_c is None: continue cos_c = (z_q * z_c).sum() cos_others = [] if s_spans: z_s = _pool(z_seq, s_spans[0]) if z_s is not None: cos_others.append((z_q * z_s).sum()) if d_spans: z_d = _pool(z_seq, d_spans[0]) if z_d is not None: cos_others.append((z_q * z_d).sum()) z_n = _pool(z_seq, (0, min(8, T))) if z_n is not None: cos_others.append((z_q * z_n).sum()) if not cos_others: continue max_other = torch.stack(cos_others).max() margins.append(cos_c - max_other) item_idxs.append(b) if margins: return torch.stack(margins), item_idxs return z.new_zeros((0,)), [] def _v4d_causal_dependence_loss(margins_normal, margins_clamped, causal_gap: float = 0.25): """Hinge loss that forces margin_normal − margin_clamped ≥ causal_gap. Both inputs are per-item margin tensors of equal length (caller must align item indices across normal/clamped forwards). The loss is mean of relu(causal_gap − (m_n − m_c)). Penalty pushes the head to MAKE its margin depend on v_vec: if clamping v_vec doesn't drop the margin by at least `causal_gap`, the head is bypassing the SL pathway. """ if margins_normal.numel() == 0 or margins_clamped.numel() == 0: return margins_normal.sum() * 0.0 gap_satisfied = margins_normal - margins_clamped deficit = causal_gap - gap_satisfied return F.relu(deficit).mean() def _v4d_pool_code_by_role(code, spans): """Mean-pool the salience code g ([B, T, k]) over each item's current-span and stale-span. Returns {batch_pos: (g_current[k], g_stale[k])} for every anchorable item. Used by the V4e swap-consistency + role-separation losses (pools by ROLE span — own offsets per item, review Blocker 3).""" out = {} if code is None: return out T = code.size(1) for b, it in enumerate(spans): if (getattr(it, "family", "") == "neutral_control" or not it.current or not it.stale): continue cs, ce = it.current[0] ss, se = it.stale[0] if ce <= cs or se <= ss or ce > T or se > T: continue out[b] = (code[b, cs:ce].mean(0), code[b, ss:se].mean(0)) return out def _v4c_answer_margin_loss(logits_mem, ids, spans, target_margin: float = 1.0): """Hinge loss on the LM's preference for `current` over `stale` at the last prompt position. For each anchorable item: - last_pos = item.original_length − 1 (the position whose logits predict the NEXT token, i.e., what follows the prompt's "Answer:"). Uses the stored original_length so the function is pad-token-id-agnostic. - get log p(first_token_of_current) and log p(first_token_of_stale) from softmax(logits_mem[last_pos]) - margin = log_p_current − log_p_stale - loss_i = max(0, target_margin − margin) Returns mean over anchorable items. Note: this aux loss does NOT require multi-token generation — only the FIRST token of current/stale is scored. Since the generator's current/stale words are typically 1-2 BPE tokens that differ at the first token (e.g., "car" vs "hat", "Berlin" vs "Bonn", "C7" vs "A12"), first-token margin is a clean signal. The gradient path: log_p depends on logits_mem ← L-stage ← h_D ← DStage(v_vec=v_vec) ← v_vec ← salience_bias_up. So the answer-margin loss directly trains the SL pathway, which is exactly the V4c+M hypothesis: output-side pressure makes the LM use v_vec. """ B = logits_mem.size(0) losses = [] n_skipped_neutral = 0 n_skipped_no_spans = 0 n_skipped_no_length = 0 n_skipped_equal_token = 0 for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") if family == "neutral_control": n_skipped_neutral += 1 continue c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) ids_b = item.ids if hasattr(item, "ids") else item.get("ids") if not c_spans or not s_spans: n_skipped_no_spans += 1 continue orig_len = (item.original_length if hasattr(item, "original_length") else item.get("original_length", 0)) if orig_len <= 0 or orig_len > logits_mem.size(1): n_skipped_no_length += 1 continue last_pos = orig_len - 1 cur_tok = ids_b[c_spans[0][0]] stale_tok = ids_b[s_spans[0][0]] if cur_tok == stale_tok: n_skipped_equal_token += 1 continue logp = F.log_softmax(logits_mem[b, last_pos].float(), dim=-1) log_p_cur = logp[cur_tok] log_p_stale = logp[stale_tok] margin = log_p_cur - log_p_stale loss_i = F.relu(target_margin - margin) losses.append(loss_i) if not losses: out = logits_mem.sum() * 0.0 else: out = torch.stack(losses).mean() out._v4c_margin_n_used = len(losses) out._v4c_margin_n_skip_neu = n_skipped_neutral out._v4c_margin_n_skip_spn = n_skipped_no_spans out._v4c_margin_n_skip_len = n_skipped_no_length out._v4c_margin_n_skip_eq = n_skipped_equal_token return out def _v4c_collect_answer_quads(ids: torch.Tensor, spans): """V6.1 helper. Mirrors _v4c_answer_margin_loss's per-item filtering (skip neutral, skip no-span, skip no-length, skip equal-token), returning LongTensors `(b_idx, p_idx, cur_tok, stale_tok)` of length N on `ids.device`. The V6.1 hook uses these as fancy-indices into [B, T, ...] activations to avoid materializing the full [B, T, V=50257] mixture log-prob tensor. Empty case (zero valid items): returns four empty `[0]` LongTensors on `ids.device`, NOT None — caller's `if b_idx.numel() > 0:` is the gate. """ B = ids.size(0) T = ids.size(1) b_list, p_list, cur_list, stale_list = [], [], [], [] for b in range(B): item = spans[b] family = item.family if hasattr(item, "family") else item.get("family") if family == "neutral_control": continue c_spans = item.current if hasattr(item, "current") else item.get("current", []) s_spans = item.stale if hasattr(item, "stale") else item.get("stale", []) ids_b = item.ids if hasattr(item, "ids") else item.get("ids") if not c_spans or not s_spans: continue orig_len = (item.original_length if hasattr(item, "original_length") else item.get("original_length", 0)) if orig_len <= 0 or orig_len > T: continue cur_t = int(ids_b[c_spans[0][0]]) stale_t = int(ids_b[s_spans[0][0]]) if cur_t == stale_t: continue b_list.append(b) p_list.append(orig_len - 1) cur_list.append(cur_t) stale_list.append(stale_t) dev = ids.device return (torch.tensor(b_list, dtype=torch.long, device=dev), torch.tensor(p_list, dtype=torch.long, device=dev), torch.tensor(cur_list, dtype=torch.long, device=dev), torch.tensor(stale_list, dtype=torch.long, device=dev)) def _v4c_collect_background_quads(ids: torch.Tensor, spans, *, ans_b_idx: torch.Tensor, ans_p_idx: torch.Tensor, k_per_item: int = 8): """Stage 2a — deterministic background-position collector for L_bg. For each V4c item: - Valid range = [0, original_length-1); excludes the answer p_idx. - INCLUDES entity-span positions (current/stale/distractor) — the head must learn to keep λ low at "X owns a Y" entity-mention rows and lift only at the structural "Answer:" answer row. That's the whole selectivity hypothesis; do NOT exclude entity spans here. - DETERMINISTIC: evenly-spaced via per-item offset = b % step. Avoids random background rows so fixed seeds remain reproducible. - If valid_len < k_per_item: cycle deterministically with replacement. - Skip items with valid_len == 0. Returns (b_bg, p_bg) LongTensors on `ids.device`, dtype torch.long. Empty case: two `[0]` LongTensors on `ids.device` (caller's `.numel()` gate is the activation check). """ B = ids.size(0) ans_by_b = {int(_b): int(_p) for _b, _p in zip( ans_b_idx.tolist(), ans_p_idx.tolist())} b_list, p_list = [], [] for b in range(B): item = spans[b] ol = int(getattr(item, "original_length", 0) if hasattr(item, "original_length") else item.get("original_length", 0)) if ol <= 1: continue skip = ans_by_b.get(b, -1) valid = [p for p in range(ol - 1) if p != skip] if not valid: continue if len(valid) >= k_per_item: step = max(1, len(valid) // k_per_item) offset = b % step picks = valid[offset::step][:k_per_item] while len(picks) < k_per_item: picks.append(valid[(offset + len(picks)) % len(valid)]) else: picks = [valid[i % len(valid)] for i in range(k_per_item)] b_list.extend([b] * len(picks)) p_list.extend(picks) dev = ids.device return (torch.tensor(b_list, dtype=torch.long, device=dev), torch.tensor(p_list, dtype=torch.long, device=dev)) def _pad_v4c_items(items, T_max): """Right-pad each item's ids to T_max (pad=0), clipping spans, in place. Returns the [N, T_max] LongTensor.""" ids_padded = [] for it in items: ids = it.ids if len(ids) > T_max: ids = ids[:T_max] for fname in ("query", "current", "stale", "distractor", "neutral", "paraphrase_positives"): spans = getattr(it, fname) spans[:] = [(s, min(e, T_max)) for (s, e) in spans if s < T_max] pad_len = T_max - len(ids) if pad_len > 0: ids = ids + [0] * pad_len ids_padded.append(ids) it.ids = ids return torch.tensor(ids_padded, dtype=torch.long) def _build_v4c_contrastive_batch(generator_fn, tokenizer, B: int, T_max: int, *, randomize_positive_mapping: bool = False, seed: int = 0, split: str = "all", hard_collision_frac: float = 0.0, swap_fn=None, family_weights=None): """Build one V4c synthetic mini-batch. Args: generator_fn: callable like data_v4c_pairs.generate(tokenizer, n, seed) tokenizer: GPT2TokenizerFast (or compatible) B: batch size T_max: pad/truncate sequences to this length randomize_positive_mapping: if True, rotate anchor→positive mapping across items (V4c-randlabel control). For each anchor `a_i`, swap its current span's positive role with item a_(i+1 mod B)'s current span (with FN masking still applied at loss time). seed: rng seed for this batch Additional kwargs: split: vocabulary partition ('train'/'test'/'all') passed to the generator (V4e novel-entity training). hard_collision_frac: fraction of but_update items with stale pinned to a recent current object (V4e role-by-position pressure). swap_fn: if given (e.g. data_v4c_pairs.make_entity_swap), also build an entity-swapped twin per item and return (ids, items, swap_ids, swap_items) for the V4e swap-consistency / role-separation losses. Otherwise returns (ids, items, None, None). Returns: (ids_tensor, items, swap_ids_or_None, swap_items_or_None) """ try: items = generator_fn(tokenizer, B, seed=seed, split=split, hard_collision_frac=hard_collision_frac, family_weights=family_weights) except TypeError: try: items = generator_fn(tokenizer, B, seed=seed, split=split, hard_collision_frac=hard_collision_frac) except TypeError: items = generator_fn(tokenizer, B, seed=seed) ids_tensor = _pad_v4c_items(items, T_max) if swap_fn is None: return ids_tensor, items, None, None import random as _random swap_rng = _random.Random(seed + 10007) twins = [swap_fn(it, tokenizer, swap_rng, split) for it in items] swap_tensor = _pad_v4c_items(twins, T_max) return ids_tensor, items, swap_tensor, twins class V6SourceAdapter(nn.Module): """Small residual MLP used only by the V6 direct specialist readout.""" def __init__(self, d_model: int): super().__init__() self.ln = nn.LayerNorm(d_model) self.fc1 = nn.Linear(d_model, d_model) self.fc2 = nn.Linear(d_model, d_model) nn.init.normal_(self.fc2.weight, std=1e-3) nn.init.zeros_(self.fc2.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return x + self.fc2(F.gelu(self.fc1(self.ln(x)))) class V6ContextAdapter(nn.Module): """Causal prefix adapter for the V6 direct specialist readout.""" def __init__(self, d_model: int, n_heads: int): super().__init__() self.ln = nn.LayerNorm(d_model) self.attn = nn.MultiheadAttention( d_model, n_heads, dropout=0.0, batch_first=True) def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dim() != 3: raise RuntimeError( "V6ContextAdapter expects [B,T,d] source states") T = x.size(1) mask = torch.ones(T, T, device=x.device, dtype=torch.bool).triu(1) h = self.ln(x) ctx, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False) return x + ctx class CEDLTwoLoop100M(nn.Module): """CEDL Two-Loop at 100M scale with coupled feedback loops. Forward pass: Pass 0 (feedforward): C -> E -> D -> L (no feedback) Pass 1..N (feedback): L->C and D->E signals enrich re-encoding Final: NeuromodulatorGate applies supplementary gain modulation Auxiliary losses (training only): - CSR: contrastive support repulsion for E-stage pattern separation - SIGReg: Gaussian regularizer on D-stage outputs (replaces attractor contrastive) - Predictive latent: D(t) predicts C(t+1) — JEPA-inspired forward model """ def __init__(self, cfg: CEDLConfig, vocab: int, max_seq: int): super().__init__() d = cfg.d_model d_expand = d * cfg.e_expand self.n_feedback_iters = cfg.n_feedback_iters self.feedback_decay = cfg.feedback_decay self.use_salience = cfg.use_salience self.salience_mode = cfg.salience_mode self.lex_anchor_weight = cfg.lex_anchor_weight self.v4c_variant = getattr(cfg, "v4c_variant", "base") self._v4c_cfg = cfg self.nll_ema_decay = 0.99 self.sl_window_k = 3 self.sl_z_clip = 4.0 self.register_buffer('feedback_alpha', torch.tensor(0.0)) _sb_sigma = 0.0 if (cfg.salience_mode in ("v4c", "v4c_gate") and self.v4c_variant != "cold"): _sb_sigma = float(getattr(cfg, "v4c_warm_start_sigma", 0.05)) _v4c_norm_cap = float(getattr(cfg, "v4c_norm_cap", 0.3)) _v4d_w_sal_sigma = float(getattr(cfg, "v4d_w_sal_sigma", 0.02)) _v4d_logit_cap = float(getattr(cfg, "v4d_logit_cap", 4.0)) _v4d_noinject = bool(getattr(cfg, "v4d_noinject", False)) _v4d_w_sal_rank = int(getattr(cfg, "v4d_w_sal_rank", 0)) self.c_stage = CStageRetention( vocab, max_seq, d, cfg.n_heads, cfg.c_layers, cfg.ffn_dim, cfg.dropout) self.e_stage = EStage(d, cfg.e_expand, cfg.e_sparsity) self.d_stage = DStage(d, d_expand, cfg.n_heads, cfg.d_refine, cfg.d_slots, cfg.dropout, use_salience=cfg.use_salience, salience_mode=cfg.salience_mode, salience_bias_init_sigma=_sb_sigma, v4c_norm_cap=_v4c_norm_cap, v4d_w_sal_sigma=_v4d_w_sal_sigma, v4d_logit_cap=_v4d_logit_cap, v4d_noinject=_v4d_noinject, v4d_w_sal_rank=_v4d_w_sal_rank) _role_w = float(getattr(cfg, "v4d_role_sep_weight", 0.0)) if cfg.use_salience and cfg.salience_mode == "v4d" and _role_w > 0: _code_dim = _v4d_w_sal_rank if _v4d_w_sal_rank > 0 else d self.v4d_role_head = nn.Linear(_code_dim, 2) self.l_stage = LStage(d, vocab, cfg.n_heads, cfg.dropout) self.modulator = NeuromodulatorGate(d) if self.use_salience: with torch.random.fork_rng(devices=[]): self.register_buffer('sl_alpha', torch.tensor(0.0)) self.register_buffer('v4c_bce_alpha', torch.tensor(0.0), persistent=False) self.register_buffer('v4c_nce_alpha', torch.tensor(0.0), persistent=False) self.salience = SalienceLoop( d, H=8, kernel=5, gate_rank=16, scalar_from_contrast=(cfg.salience_mode in ( "v3", "v3m", "v3m2", "v3m2_nll", "v4c", "v4c_gate", "v4d")), contrastive_head=(cfg.salience_mode in ( "v4c", "v4c_gate", "v4d")), z_dim=int(getattr(cfg, "v4c_z_dim", 64)), proj_hidden=int(getattr(cfg, "v4c_proj_hidden", 256)), ) self.val_a = nn.Parameter(torch.tensor(1.0)) self.val_b = nn.Parameter(torch.tensor(0.0)) self.register_buffer('marker_token_ids', build_marker_token_ids(), persistent=False) _need_nll_ema = cfg.salience_mode in ("v3", "v3m", "v3m2_nll") if _need_nll_ema: self.register_buffer('nll_ema_mean', torch.tensor(0.0)) self.register_buffer('nll_ema_var', torch.tensor(1.0)) self.register_buffer('nll_ema_initialized', torch.tensor(0.0)) if cfg.salience_mode in ("v3m", "v3m2", "v3m2_nll"): for _name in ("curve_C", "curve_E", "cosmis"): self.register_buffer(f"{_name}_ema_mean", torch.tensor(0.0)) self.register_buffer(f"{_name}_ema_var", torch.tensor(1.0)) self.register_buffer(f"{_name}_ema_initialized", torch.tensor(0.0)) if cfg.salience_mode in ("v3m2", "v3m2_nll"): _anchor_neg, _label_neg = build_v3m2_lexicons() self.register_buffer('anchor_negative_ids', _anchor_neg, persistent=False) self.register_buffer('label_negative_ids', _label_neg, persistent=False) self.l_stage.per_head.weight = self.c_stage.tok_emb.weight self.l_stage.mem_head.weight = self.c_stage.tok_emb.weight self.d_model = d self.v6_mixture = bool(getattr(cfg, "v6_mixture", False)) self.v6_lambda_init = float(getattr(cfg, "v6_lambda_init", -4.0)) self.v6_lambda_a_init = float(getattr(cfg, "v6_lambda_a_init", 1.0)) self.v6_bank_query_source = str( getattr(cfg, "v6_bank_query_source", "h_d")) _valid_bank_sources = {"h_d", "h_e", "q_attractor", "q_mem"} if self.v6_bank_query_source not in _valid_bank_sources: raise ValueError( f"unknown v6_bank_query_source={self.v6_bank_query_source!r}; " f"expected one of {sorted(_valid_bank_sources)}") self.v6_bank_readout_mode = str( getattr(cfg, "v6_bank_readout_mode", "bank")) _valid_readout_modes = {"bank", "direct"} if self.v6_bank_readout_mode not in _valid_readout_modes: raise ValueError( f"unknown v6_bank_readout_mode={self.v6_bank_readout_mode!r}; " f"expected one of {sorted(_valid_readout_modes)}") self.use_v6_source_adapter = bool( getattr(cfg, "v6_source_adapter", False)) self.use_v6_context_adapter = bool( getattr(cfg, "v6_context_adapter", False)) self.bank_off = False self.outputs_log_probs = self.v6_mixture if self.v6_mixture: self.bank_q_proj = nn.Linear(d, d) nn.init.eye_(self.bank_q_proj.weight) nn.init.zeros_(self.bank_q_proj.bias) self.v6_lambda_a = nn.Parameter(torch.tensor(self.v6_lambda_a_init)) self.v6_lambda_b = nn.Parameter(torch.tensor(self.v6_lambda_init)) if self.use_v6_source_adapter: self.v6_source_adapter = V6SourceAdapter(d) if self.use_v6_context_adapter: self.v6_context_adapter = V6ContextAdapter(d, cfg.n_heads) self.use_v6_lambda_head = bool( getattr(cfg, "v6_lambda_head", False)) if self.use_v6_lambda_head: _h = int(getattr(cfg, "v6_lambda_head_hidden", 160)) _b_init = float(getattr(cfg, "v6_lambda_head_bias_init", -7.0)) _w_std = float(getattr(cfg, "v6_lambda_head_w_init_std", 1e-3)) self.v6_lambda_head = nn.Sequential( nn.Linear(d, _h), nn.GELU(), nn.Linear(_h, 1), ) nn.init.normal_(self.v6_lambda_head[2].weight, std=_w_std) nn.init.constant_(self.v6_lambda_head[2].bias, _b_init) self.use_v6_mem_head_bank = bool( getattr(cfg, "v6_mem_head_bank", False)) if self.use_v6_mem_head_bank: self.mem_head_bank = nn.Linear(d, vocab, bias=True) with torch.no_grad(): self.mem_head_bank.weight.copy_( self.c_stage.tok_emb.weight) if self.l_stage.mem_head.bias is not None: self.mem_head_bank.bias.copy_( self.l_stage.mem_head.bias) else: nn.init.zeros_(self.mem_head_bank.bias) def _v6_needs_dstage_bank_sources(self) -> bool: return self.v6_bank_query_source in ("q_attractor", "q_mem") def _v6_bank_query_input(self, h_D: torch.Tensor, h_E: Optional[torch.Tensor] = None, b_idx: Optional[torch.Tensor] = None, p_idx: Optional[torch.Tensor] = None) -> torch.Tensor: """Select the source vector used by V6 bank_q_proj. The default h_d path preserves original V6 behavior. q_attractor/q_mem are populated by DStage when store_v6_bank_sources is enabled. """ src = self.v6_bank_query_source if src == "h_d": base = h_D elif src == "h_e": if h_E is None: raise RuntimeError("v6_bank_query_source='h_e' requires h_E") base = h_E elif src == "q_attractor": base = getattr(self.d_stage, "_last_q_attractor_v6", None) if base is None: raise RuntimeError( "q_attractor source requested but DStage did not store it") elif src == "q_mem": base = getattr(self.d_stage, "_last_q_mem_v6", None) if base is None: raise RuntimeError( "q_mem source requested but DStage did not store it") else: raise RuntimeError(f"unknown V6 bank query source {src!r}") if b_idx is None: return base return base[b_idx, p_idx] def _v6_direct_readout_source(self, bank_src: torch.Tensor) -> torch.Tensor: if getattr(self, "use_v6_context_adapter", False): bank_src = self.v6_context_adapter(bank_src) if getattr(self, "use_v6_source_adapter", False): return self.v6_source_adapter(bank_src) return bank_src def _v6_bank_logits_from_source(self, bank_src: torch.Tensor, adapt_direct: bool = True) -> torch.Tensor: """Project the selected V6 source to vocab logits. `bank` is the original V6 external readout. `direct` bypasses the 256-slot mem_keys/mem_vals compression after direct-source probing showed q_mem is decodable while the bank bottleneck is not. """ if self.v6_bank_readout_mode == "direct": if adapt_direct: bank_src = self._v6_direct_readout_source(bank_src) if getattr(self, "use_v6_mem_head_bank", False): return self.mem_head_bank(bank_src) return self.l_stage.mem_head(bank_src) q_bank = self.bank_q_proj(bank_src) alpha = F.softmax( q_bank @ self.d_stage.mem_keys.T / math.sqrt(self.d_model), dim=-1, ) v_bank = alpha @ self.d_stage.mem_vals if getattr(self, "use_v6_mem_head_bank", False): return self.mem_head_bank(v_bank) return self.l_stage.mem_head(v_bank) def forward(self, ids: torch.Tensor, v4c_spans=None, v4c_swap_ids=None, v4c_swap_spans=None): if v4c_spans is not None: assert self.use_salience and self.salience_mode in ( "v4c", "v4c_gate", "v4d"), \ "v4c_spans only meaningful for V4c-family models with SL enabled" _swap_w = float(getattr(self._v4c_cfg, "v4d_swap_consistency_weight", 0.0)) _role_w = float(getattr(self._v4c_cfg, "v4d_role_sep_weight", 0.0)) _v4e_active = (self.salience_mode == "v4d" and (_swap_w > 0 or _role_w > 0) and v4c_swap_ids is not None) if _v4e_active: self.d_stage.store_v4d_code = True h_C = self.c_stage(ids, feedback=None) h_E, h_E_sparse = self.e_stage(h_C, feedback=None) v_vec, v_scalar = self.salience(h_E) _prev_v6_store = getattr( self.d_stage, "store_v6_bank_sources", False) self.d_stage.store_v6_bank_sources = ( self._v6_needs_dstage_bank_sources()) try: h_D, _ = self.d_stage(h_E, h_E_sparse, v_vec=v_vec) finally: self.d_stage.store_v6_bank_sources = _prev_v6_store g_main = getattr(self.d_stage, "_last_sal_code", None) if _v4e_active else None if _v4e_active: self.d_stage.store_v4d_code = False z = self.salience.embed_z(h_D) _randlabel = (getattr(self, "v4c_variant", "base") == "randlabel") v4c_aux = _v4c_contrastive_loss( z=z, v_scalar=v_scalar, spans=v4c_spans, cfg=self._v4c_cfg, salience_module=self.salience, bce_alpha=getattr(self, "v4c_bce_alpha", torch.tensor(0.0, device=ids.device)), nce_alpha=getattr(self, "v4c_nce_alpha", torch.tensor(0.0, device=ids.device)), randomize_positive_mapping=_randlabel, ) margin_weight = float(getattr(self._v4c_cfg, "v4c_margin_weight", 0.0)) if margin_weight > 0 and not _randlabel: logits_mem, _loop1 = self.l_stage(h_D, h_C) margin_target = float(getattr(self._v4c_cfg, "v4c_margin_target", 1.0)) L_margin = _v4c_answer_margin_loss( logits_mem, ids, v4c_spans, target_margin=margin_target) self.salience._last_v4c_margin_stats = { "loss": float(L_margin.detach().item()), "n_used": int(getattr(L_margin, "_v4c_margin_n_used", 0)), "n_skip_neutral": int(getattr(L_margin, "_v4c_margin_n_skip_neu", 0)), "n_skip_no_spans": int(getattr(L_margin, "_v4c_margin_n_skip_spn", 0)), "n_skip_no_len": int(getattr(L_margin, "_v4c_margin_n_skip_len", 0)), "n_skip_eq_token": int(getattr(L_margin, "_v4c_margin_n_skip_eq", 0)), } bce_alpha = getattr(self, "v4c_bce_alpha", torch.tensor(0.0, device=ids.device)) v4c_aux = v4c_aux + bce_alpha * margin_weight * L_margin self.salience._last_v6_aux_stats = None _v6_w = float(getattr(self._v4c_cfg, "v6_aux_weight", 0.0)) if (self.v6_mixture and _v6_w > 0 and not _randlabel and self.salience_mode == "v4d"): if margin_weight <= 0: raise RuntimeError( "V6.1 (v6_aux_weight > 0) requires v4c_margin_weight " "> 0 so logits_mem can be reused from the V4c+M " "branch — otherwise the fallback materializes a " "[B,T,V] tensor (~800 MB per V4c batch). Set both " "knobs together.") b_idx, p_idx, cur_tok, stale_tok = _v4c_collect_answer_quads( ids, v4c_spans) if b_idx.numel() > 0: h_C_mod = self.modulator(h_C, _loop1, h_D) h_D_ans = h_D[b_idx, p_idx] h_C_mod_ans = h_C_mod[b_idx, p_idx] v_scalar_ans = v_scalar[b_idx, p_idx] logits_mem_a = logits_mem[b_idx, p_idx] logits_per_a = self.l_stage.per_head(h_C_mod_ans) gate_a = self.l_stage.gate_net( torch.cat([h_D_ans, h_C_mod_ans], dim=-1)) logits_trunk_a = (gate_a * logits_mem_a + (1 - gate_a) * logits_per_a) if getattr(self, "use_v6_lambda_head", False): lam_logit_a = self.v6_lambda_head(h_D_ans) else: lam_logit_a = (self.v6_lambda_a * v_scalar_ans.detach() + self.v6_lambda_b) lam_a = torch.sigmoid(lam_logit_a).clamp( 1e-4, 1.0 - 1e-4) if getattr(self, "use_v6_context_adapter", False): bank_src_seq = self._v6_bank_query_input(h_D, h_E) bank_src_seq = self._v6_direct_readout_source( bank_src_seq) bank_src_a = bank_src_seq[b_idx, p_idx] logits_bank_a = self._v6_bank_logits_from_source( bank_src_a, adapt_direct=False) else: bank_src_a = self._v6_bank_query_input( h_D, h_E, b_idx, p_idx) logits_bank_a = self._v6_bank_logits_from_source( bank_src_a) log_p_mix_a = torch.logaddexp( torch.log1p(-lam_a) + F.log_softmax( logits_trunk_a, dim=-1), torch.log(lam_a) + F.log_softmax( logits_bank_a, dim=-1)) _target = float(getattr(self._v4c_cfg, "v6_margin_target", 1.0)) log_p_bank_a = F.log_softmax(logits_bank_a, dim=-1) lp_bk_cur = log_p_bank_a.gather( 1, cur_tok.unsqueeze(1)).squeeze(1) lp_bk_stale = log_p_bank_a.gather( 1, stale_tok.unsqueeze(1)).squeeze(1) L_bank = F.relu(_target + lp_bk_stale - lp_bk_cur).mean() L_bank_ce = torch.zeros((), device=ids.device) _w_bank_ce = float(getattr(self._v4c_cfg, "v6_bank_ce_weight", 0.0)) if (_w_bank_ce > 0 and getattr(self, "use_v6_mem_head_bank", False)): L_bank_ce = F.cross_entropy(logits_bank_a, cur_tok) L_bank_pair_ce = torch.zeros((), device=ids.device) _w_bank_pair_ce = float(getattr( self._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)) if (_w_bank_pair_ce > 0 and getattr(self, "use_v6_mem_head_bank", False)): bk_cur_logit = logits_bank_a.gather( 1, cur_tok.unsqueeze(1)).squeeze(1) bk_stale_logit = logits_bank_a.gather( 1, stale_tok.unsqueeze(1)).squeeze(1) pair_logits = torch.stack( [bk_cur_logit, bk_stale_logit], dim=1) L_bank_pair_ce = F.cross_entropy( pair_logits, torch.zeros_like(cur_tok, dtype=torch.long)) lp_mix_cur = log_p_mix_a.gather( 1, cur_tok.unsqueeze(1)).squeeze(1) lp_mix_stale = log_p_mix_a.gather( 1, stale_tok.unsqueeze(1)).squeeze(1) L_mix = F.relu(_target + lp_mix_stale - lp_mix_cur).mean() _lam_floor = float(getattr(self._v4c_cfg, "v6_lambda_floor", 0.05)) _lam_floor_logit = math.log( _lam_floor / (1.0 - _lam_floor)) L_gate = F.relu(_lam_floor_logit - lam_logit_a.view(-1)).mean() L_bg = torch.zeros((), device=ids.device) _bg_mean_val = 0.0 _bg_std_val = 0.0 _bg_n_items = 0 if getattr(self, "use_v6_lambda_head", False): b_bg, p_bg = _v4c_collect_background_quads( ids, v4c_spans, ans_b_idx=b_idx, ans_p_idx=p_idx, k_per_item=8) if b_bg.numel() > 0: h_D_bg = h_D[b_bg, p_bg] lam_logit_bg = self.v6_lambda_head(h_D_bg).squeeze(-1) lam_bg = torch.sigmoid(lam_logit_bg).clamp( 1e-4, 1.0 - 1e-4) _bg_target = float(getattr(self._v4c_cfg, "v6_bg_target", 0.01)) L_bg = F.relu(lam_bg - _bg_target).mean() with torch.no_grad(): _lb = lam_bg.detach().float() _bg_mean_val = float(_lb.mean().item()) _bg_std_val = float(_lb.std( unbiased=False).item()) _bg_n_items = int(_lb.numel()) _use_bce = (getattr(self, "use_v6_lambda_head", False) and bool(getattr(self._v4c_cfg, "v6_bce_objective", False))) _w_mix = float(getattr(self._v4c_cfg, "v6_mix_weight", 0.5)) L_sel_ans = torch.zeros((), device=ids.device) L_sel_bg = torch.zeros((), device=ids.device) L_sel = torch.zeros((), device=ids.device) if _use_bce: L_sel_ans = -F.logsigmoid(lam_logit_a.view(-1)).mean() if (getattr(self, "use_v6_lambda_head", False) and b_bg.numel() > 0): L_sel_bg = -F.logsigmoid( -lam_logit_bg.view(-1)).mean() L_sel = 0.5 * L_sel_ans + 0.5 * L_sel_bg _w_sel = float(getattr(self._v4c_cfg, "v6_sel_weight", 1.0)) L_v6 = (L_bank + _w_mix * L_mix + _w_sel * L_sel + _w_bank_ce * L_bank_ce + _w_bank_pair_ce * L_bank_pair_ce) else: _w_gate = float(getattr(self._v4c_cfg, "v6_gate_weight", 0.01)) _w_bg = float(getattr(self._v4c_cfg, "v6_bg_weight", 1.0)) L_v6 = (L_bank + _w_mix * L_mix + _w_gate * L_gate + _w_bg * L_bg + _w_bank_ce * L_bank_ce + _w_bank_pair_ce * L_bank_pair_ce) bce_alpha = getattr(self, "v4c_bce_alpha", torch.tensor(0.0, device=ids.device)) v4c_aux = v4c_aux + bce_alpha * _v6_w * L_v6 with torch.no_grad(): _l = lam_a.detach().float().view(-1) _llg = lam_logit_a.detach().float().view(-1) _stats = { "loss": float(L_v6.detach().item()), "L_bank": float(L_bank.detach().item()), "L_bank_ce": float(L_bank_ce.detach().item()), "L_bank_pair_ce": float( L_bank_pair_ce.detach().item()), "L_mix": float(L_mix.detach().item()), "L_gate": float(L_gate.detach().item()), "lam_ans_mean": float(_l.mean().item()), "lam_ans_std": float(_l.std(unbiased=False).item()), "lam_ans_sat_low": float((_l < 1e-3).float().mean().item()), "lam_ans_sat_high": float((_l > 1.0 - 1e-3).float().mean().item()), "n_items": int(_l.numel()), "v6_lambda_a": float(self.v6_lambda_a.item()), "v6_lambda_b": float(self.v6_lambda_b.item()), "lam_logit_ans_mean": float(_llg.mean().item()), "lam_logit_ans_std": float(_llg.std(unbiased=False).item()), "L_bg": float(L_bg.detach().item()), "lam_bg_mean": _bg_mean_val, "lam_bg_std": _bg_std_val, "n_bg_items": _bg_n_items, "L_sel": float(L_sel.detach().item()), "L_sel_ans": float(L_sel_ans.detach().item()), "L_sel_bg": float(L_sel_bg.detach().item()), "bce_active": bool(_use_bce), } if getattr(self, "use_v6_lambda_head", False): _stats["v6_lambda_head_bias"] = float( self.v6_lambda_head[2].bias.item()) self.salience._last_v6_aux_stats = _stats causal_weight = float(getattr(self._v4c_cfg, "v4d_causal_weight", 0.0)) causal_z_weight = float(getattr(self._v4c_cfg, "v4d_causal_z_weight", 0.0)) run_clamped_fwd = ((causal_weight > 0 or causal_z_weight > 0) and not _randlabel and self.salience_mode == "v4d") if run_clamped_fwd: def _zero_v_vec_hook(_mod, _inputs, output): v_vec, v_scalar = output return torch.zeros_like(v_vec), v_scalar handle = self.salience.register_forward_hook(_zero_v_vec_hook) try: h_C2 = self.c_stage(ids, feedback=None) h_E2, h_E_sparse2 = self.e_stage(h_C2, feedback=None) v_vec2, v_scalar2 = self.salience(h_E2) h_D2, _ = self.d_stage(h_E2, h_E_sparse2, v_vec=v_vec2) logits_mem_clamped = None if causal_weight > 0 and margin_weight > 0: logits_mem_clamped, _ = self.l_stage(h_D2, h_C2) z_clamped = self.salience.embed_z(h_D2) finally: handle.remove() nce_alpha = getattr(self, "v4c_nce_alpha", torch.tensor(0.0, device=ids.device)) if (causal_weight > 0 and margin_weight > 0 and logits_mem_clamped is not None): causal_gap = float(getattr(self._v4c_cfg, "v4d_causal_gap", 0.25)) margins_n, item_idxs_n = _v4d_compute_per_item_margins( logits_mem, ids, v4c_spans) margins_c, item_idxs_c = _v4d_compute_per_item_margins( logits_mem_clamped, ids, v4c_spans) if item_idxs_n == item_idxs_c and margins_n.numel() > 0: L_causal = _v4d_causal_dependence_loss( margins_n, margins_c, causal_gap=causal_gap) v4c_aux = v4c_aux + nce_alpha * causal_weight * L_causal self.salience._last_v4d_causal_stats = { "loss": float(L_causal.detach().item()), "margin_normal_mean": float(margins_n.detach().mean().item()), "margin_clamped_mean": float(margins_c.detach().mean().item()), "margin_gap_mean": float((margins_n.detach() - margins_c.detach()).mean().item()), "n_items": int(margins_n.numel()), } if causal_z_weight > 0: causal_z_gap = float(getattr(self._v4c_cfg, "v4d_causal_z_gap", 0.15)) mz_n, item_idxs_zn = _v4d_compute_per_item_z_margins(z, v4c_spans) mz_c, item_idxs_zc = _v4d_compute_per_item_z_margins(z_clamped, v4c_spans) if item_idxs_zn == item_idxs_zc and mz_n.numel() > 0: L_causal_z = _v4d_causal_dependence_loss( mz_n, mz_c, causal_gap=causal_z_gap) v4c_aux = v4c_aux + nce_alpha * causal_z_weight * L_causal_z self.salience._last_v4d_causal_z_stats = { "loss": float(L_causal_z.detach().item()), "mz_normal_mean": float(mz_n.detach().mean().item()), "mz_clamped_mean": float(mz_c.detach().mean().item()), "mz_gap_mean": float((mz_n.detach() - mz_c.detach()).mean().item()), "n_items": int(mz_n.numel()), } if _v4e_active and g_main is not None: self.d_stage.store_v4d_code = True try: h_Cs = self.c_stage(v4c_swap_ids, feedback=None) h_Es, h_Es_sp = self.e_stage(h_Cs, feedback=None) v_vecs, _ = self.salience(h_Es) self.d_stage(h_Es, h_Es_sp, v_vec=v_vecs) g_swap = getattr(self.d_stage, "_last_sal_code", None) finally: self.d_stage.store_v4d_code = False nce_alpha = getattr(self, "v4c_nce_alpha", torch.tensor(0.0, device=ids.device)) m_pool = _v4d_pool_code_by_role(g_main, v4c_spans) s_pool = _v4d_pool_code_by_role(g_swap, v4c_swap_spans) common = sorted(set(m_pool) & set(s_pool)) if common: gc_m = torch.stack([m_pool[b][0] for b in common]) gs_m = torch.stack([m_pool[b][1] for b in common]) gc_s = torch.stack([s_pool[b][0] for b in common]) gs_s = torch.stack([s_pool[b][1] for b in common]) if _swap_w > 0: L_swap = ((gc_m - gc_s).pow(2).sum(-1) + (gs_m - gs_s).pow(2).sum(-1)).mean() v4c_aux = v4c_aux + nce_alpha * _swap_w * L_swap self.salience._last_v4d_swap_stats = { "loss": float(L_swap.detach().item()), "n_pairs": int(len(common)), } if _role_w > 0 and hasattr(self, "v4d_role_head"): feats = torch.cat([gc_m, gs_m, gc_s, gs_s], dim=0) n = gc_m.size(0) labels = torch.cat([ torch.ones(n, dtype=torch.long, device=ids.device), torch.zeros(n, dtype=torch.long, device=ids.device), torch.ones(n, dtype=torch.long, device=ids.device), torch.zeros(n, dtype=torch.long, device=ids.device)]) role_logits = self.v4d_role_head(feats) L_role = F.cross_entropy(role_logits, labels) v4c_aux = v4c_aux + nce_alpha * _role_w * L_role with torch.no_grad(): role_acc = (role_logits.argmax(-1) == labels).float().mean().item() self.salience._last_v4d_role_stats = { "loss": float(L_role.detach().item()), "acc": float(role_acc), "n": int(4 * n), } return None, v4c_aux h_C = self.c_stage(ids, feedback=None) h_E, h_E_sparse = self.e_stage(h_C, feedback=None) if self.use_salience: v_vec, v_scalar = self.salience(h_E) else: v_vec, v_scalar = None, None _prev_v6_store = getattr(self.d_stage, "store_v6_bank_sources", False) self.d_stage.store_v6_bank_sources = ( self._v6_needs_dstage_bank_sources()) try: h_D, loop2_signal = self.d_stage(h_E, h_E_sparse, v_vec=v_vec) logits_mem, loop1_fb = self.l_stage(h_D, h_C) v6_bank_src_p0 = ( self._v6_bank_query_input(h_D, h_E) if (self.v6_mixture and not self.bank_off) else None ) h_C_p0, h_E_p0, h_E_sparse_p0, h_D_p0 = h_C, h_E, h_E_sparse, h_D if self.use_salience: logits_mem_p0 = logits_mem v_scalar_p0 = v_scalar alpha = self.feedback_alpha.item() for i in range(self.n_feedback_iters): decay = self.feedback_decay ** i fb_scale = alpha * decay if fb_scale > 0: h_C = self.c_stage(ids, feedback=fb_scale * loop1_fb) h_E, h_E_sparse = self.e_stage(h_C, feedback=fb_scale * loop2_signal) h_D, loop2_signal = self.d_stage(h_E, h_E_sparse, v_vec=v_vec) logits_mem, loop1_fb = self.l_stage(h_D, h_C) finally: self.d_stage.store_v6_bank_sources = _prev_v6_store aux_loss = torch.tensor(0.0, device=ids.device) if self.training: aux_loss = aux_loss + 0.05 * self.e_stage.csr_loss(h_C_p0, h_E_sparse_p0) aux_loss = aux_loss + 0.1 * sigreg_loss(h_D_p0) aux_loss = aux_loss + 0.05 * predictive_latent_loss(h_D_p0, h_C_p0) if self.use_salience: aux_loss = aux_loss + 1e-3 * anchor_ortho_reg(self.salience.U) sl_a = self.sl_alpha if self.salience_mode in ("v0", "v2a"): if sl_a.item() > 0: with torch.no_grad(): flat_logits = logits_mem_p0[:, :-1].reshape( -1, logits_mem_p0.size(-1)) nll_flat = F.cross_entropy( flat_logits.float(), ids[:, 1:].reshape(-1), reduction='none') nll = nll_flat.view(ids.size(0), ids.size(1) - 1) nll_z = torch.sigmoid( (nll - nll.mean()) / (nll.std() + 1e-6)) pred = torch.sigmoid( self.val_a * v_scalar_p0[:, :-1, 0] + self.val_b) aux_loss = aux_loss + 0.03 * sl_a * F.mse_loss(pred, nll_z) if self.marker_token_ids.numel() > 0: mask = build_marker_mask(ids, self.marker_token_ids) logits_mark = (self.val_a * v_scalar_p0[:, :, 0] + self.val_b) aux_loss = aux_loss + 0.02 * sl_a * \ F.binary_cross_entropy_with_logits( logits_mark, mask) elif self.salience_mode in ("v3", "v3m"): with torch.no_grad(): flat_logits = logits_mem_p0[:, :-1].reshape( -1, logits_mem_p0.size(-1)).float() nll = F.cross_entropy( flat_logits, ids[:, 1:].reshape(-1), reduction='none').view(ids.size(0), -1) nll_z = _ema_zscore( nll, self.nll_ema_mean, self.nll_ema_var, self.nll_ema_initialized, self.nll_ema_decay, self.training ).clamp(-self.sl_z_clip, self.sl_z_clip) if self.salience_mode == "v3": raw_t = (nll_z >= 1.5).float() pad = F.pad(raw_t.unsqueeze(1), (0, self.sl_window_k)) target = F.max_pool1d( pad, kernel_size=self.sl_window_k + 1, stride=1).squeeze(1)[:, :nll_z.size(1)] else: def _curvature(h): h_n = F.normalize(h.detach().float(), dim=-1) d1 = F.normalize(h_n[:, 1:-1] - h_n[:, :-2], dim=-1) d2 = F.normalize(h_n[:, 2:] - h_n[:, 1:-1], dim=-1) cc = 1.0 - (d1 * d2).sum(-1) return F.pad(cc, (1, 0)) curve_C = _curvature(h_C_p0) curve_E = _curvature(h_E_p0) h_C_n = F.normalize(h_C_p0.detach().float(), dim=-1) h_E_n = F.normalize(h_E_p0.detach().float(), dim=-1) cosmis = (1.0 - (h_C_n * h_E_n).sum(-1))[:, :-1] curve_C_z = _ema_zscore( curve_C, self.curve_C_ema_mean, self.curve_C_ema_var, self.curve_C_ema_initialized, self.nll_ema_decay, self.training ).clamp(-self.sl_z_clip, self.sl_z_clip) curve_E_z = _ema_zscore( curve_E, self.curve_E_ema_mean, self.curve_E_ema_var, self.curve_E_ema_initialized, self.nll_ema_decay, self.training ).clamp(-self.sl_z_clip, self.sl_z_clip) cosmis_z = _ema_zscore( cosmis, self.cosmis_ema_mean, self.cosmis_ema_var, self.cosmis_ema_initialized, self.nll_ema_decay, self.training ).clamp(-self.sl_z_clip, self.sl_z_clip) target_logit = (0.7 * nll_z + 0.5 * curve_C_z + 0.5 * curve_E_z + 0.7 * cosmis_z - 2.5) target_soft = torch.sigmoid(target_logit) pad = F.pad(target_soft.unsqueeze(1), (0, self.sl_window_k)) target = F.max_pool1d( pad, kernel_size=self.sl_window_k + 1, stride=1).squeeze(1)[:, :nll_z.size(1)] if self.training: mean_target = target.mean().item() frac_high = (target > 0.5).float().mean().item() pred_sig = torch.sigmoid( v_scalar_p0[:, :-1, 0]).detach() frac_v_high = (pred_sig > 0.5).float().mean().item() step_idx = getattr( self, '_v3m_density_log_step', 0) if step_idx % 100 == 0: print(f" [v3m_density step={step_idx:>5d}] " f"mean_target={mean_target:.3f} " f"frac_target>0.5={frac_high:.3f} " f"frac_sig_v_scalar>0.5={frac_v_high:.3f}") if 200 <= step_idx < 1000 and frac_v_high > 0.8: print(f" [v3m_density EARLY-ABORT] " f"frac_sig_v_scalar>0.5={frac_v_high:.3f} " f"> 0.8 before sl_alpha starts at " f"step 1000 — scalar_head init is " f"saturated, ABORT and check " f"scalar_head.bias initialization " f"(should be ≲ -2.0 for sparse " f"salience prior)") if step_idx >= 2000: if frac_high > 0.50: print(f" [v3m_density WARNING] " f"frac_target>0.5={frac_high:.3f} " f"> 0.50 — V3m target severely broad") if frac_v_high > 0.80: print(f" [v3m_density WARNING] " f"frac_sig_v_scalar>0.5={frac_v_high:.3f} " f"> 0.80 — v_scalar saturating") self._v3m_density_log_step = step_idx + 1 if self.salience_mode == "v3": w = 1.0 + 4.0 * torch.sigmoid(nll_z.detach()) else: w = 1.0 + 4.0 * target.detach() w = w / w.mean() if sl_a.item() > 0: v_logit = v_scalar_p0[:, :-1, 0] aux_loss = aux_loss + 0.05 * sl_a * ( F.binary_cross_entropy_with_logits( v_logit, target, reduction='none') * w ).mean() if (self.lex_anchor_weight > 0 and self.marker_token_ids.numel() > 0): mask = build_marker_mask(ids, self.marker_token_ids) v_logit_full = v_scalar_p0[:, :, 0] aux_loss = aux_loss + self.lex_anchor_weight * sl_a * \ F.binary_cross_entropy_with_logits( v_logit_full, mask) elif self.salience_mode in ("v3m2", "v3m2_nll"): with torch.no_grad(): target, w, _dbg = _build_v3m2_targets( ids=ids, h_C_p0=h_C_p0, h_E_p0=h_E_p0, logits_mem_p0=logits_mem_p0, model=self, salience_mode=self.salience_mode, update_ema=True, with_nll=(self.salience_mode == "v3m2_nll"), ) mean_target = target.mean().item() frac_pos = (target > 0.1).float().mean().item() pred_sig = torch.sigmoid( v_scalar_p0[:, :-1, 0]).detach() frac_v_high = (pred_sig > 0.5).float().mean().item() step_idx = getattr(self, '_v3m2_density_log_step', 0) if step_idx % 100 == 0: print(f" [v3m2_density step={step_idx:>5d}] " f"mean_target={mean_target:.3f} " f"frac_pos={frac_pos:.3f} " f"frac_sig_v_scalar>0.5={frac_v_high:.3f}") if 200 <= step_idx < 1000 and frac_v_high > 0.8: print(f" [v3m2_density EARLY-ABORT] " f"frac_sig_v_scalar>0.5={frac_v_high:.3f} " f"> 0.8 before sl_alpha starts at step 1000 " f"— scalar_head init saturated, ABORT") if step_idx >= 3000: if frac_pos > 0.30: print(f" [v3m2_density WARNING] " f"frac_pos={frac_pos:.3f} > 0.30 " f"— V3m2 target too broad (anchor " f"selection should reject routine text)") if frac_v_high > 0.50: print(f" [v3m2_density WARNING] " f"frac_sig_v_scalar>0.5={frac_v_high:.3f} " f"> 0.50 — v_scalar saturating") self._v3m2_density_log_step = step_idx + 1 if sl_a.item() > 0: v_logit = v_scalar_p0[:, :-1, 0] aux_loss = aux_loss + 0.05 * sl_a * ( F.binary_cross_entropy_with_logits( v_logit, target, reduction='none') * w ).mean() del h_C_p0, h_E_p0, h_E_sparse_p0, h_D_p0 lam = None if (self.v6_mixture and not self.bank_off and self.use_salience and v_scalar is not None): if getattr(self, "use_v6_lambda_head", False): _wt_w = float(getattr(self._v4c_cfg, "v6_wt_sparsity_weight", 0.0)) if self.training and _wt_w > 0: lam_logit_main = self.v6_lambda_head(h_D) lam_with_grad = torch.sigmoid( lam_logit_main).clamp(1e-4, 1.0 - 1e-4) _wt_target = float(getattr(self._v4c_cfg, "v6_wt_sparsity_target", 0.0)) if _wt_target > 0: wt_term = F.relu( lam_with_grad - _wt_target).mean() else: wt_term = lam_with_grad.mean() aux_loss = aux_loss + _wt_w * wt_term lam = lam_with_grad.detach() else: with torch.no_grad(): lam_logit_main = self.v6_lambda_head(h_D) lam = torch.sigmoid(lam_logit_main).clamp(1e-4, 1.0 - 1e-4) else: lam = torch.sigmoid(self.v6_lambda_a * v_scalar.detach() + self.v6_lambda_b) lam = lam.clamp(1e-4, 1.0 - 1e-4) if not self.training: with torch.no_grad(): _l = lam.detach().float().view(-1) _l_cpu = _l.cpu() _gt = (_l > 0.7) _stats = { "lam_mean": float(_l.mean().item()), "lam_std": float(_l.std(unbiased=False).item()), "lam_sat_low": float((_l < 1e-3).float().mean().item()), "lam_sat_high": float((_l > 1.0 - 1e-3).float().mean().item()), "lam_max": float(_l.max().item()), "lam_gt_0_7_count": int(_gt.sum().item()), "lam_total_count": int(_l.numel()), "lam_tensor_cpu": _l_cpu, "v6_lambda_a": float(self.v6_lambda_a.item()), "v6_lambda_b": float(self.v6_lambda_b.item()), } self._last_lam_stats = _stats if self.use_salience: del logits_mem_p0, v_scalar_p0, v_vec, v_scalar h_E_bank = None del h_E_sparse del h_E h_C_mod = self.modulator(h_C, loop1_fb, h_D) del loop1_fb logits_per_mod = self.l_stage.per_head(h_C_mod) gate_alpha = self.l_stage.gate_net(torch.cat([h_D, h_C_mod], dim=-1)) del h_C, h_C_mod logits_per_mod.mul_(1 - gate_alpha) logits_mem.mul_(gate_alpha) logits_per_mod.add_(logits_mem) del logits_mem if lam is not None: bank_src = v6_bank_src_p0 if bank_src is None: raise RuntimeError( "V6 readout requested but pass-0 bank source was not captured") logits_bank = self._v6_bank_logits_from_source(bank_src) del h_D, bank_src if h_E_bank is not None: del h_E_bank log_p_trunk = F.log_softmax(logits_per_mod, dim=-1) del logits_per_mod log_p_bank = F.log_softmax(logits_bank, dim=-1) del logits_bank log_p_final = torch.logaddexp( torch.log1p(-lam) + log_p_trunk, torch.log(lam) + log_p_bank) del log_p_trunk, log_p_bank, lam return log_p_final, aux_loss del h_D if h_E_bank is not None: del h_E_bank return logits_per_mod, aux_loss class CausalSelfAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() self.attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout, batch_first=True) def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: out, _ = self.attn(x, x, x, attn_mask=mask) return out class TransformerBlock(nn.Module): def __init__(self, d_model: int, n_heads: int, ffn_dim: int, dropout: float = 0.1): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_heads, dropout) self.ln2 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, d_model), nn.Dropout(dropout), ) def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x), mask) x = x + self.ffn(self.ln2(x)) return x class TransformerLM100M(nn.Module): """GPT-2-style causal Transformer at 100M scale.""" def __init__(self, cfg: TransformerConfig, vocab: int, max_seq: int): super().__init__() self.tok_emb = nn.Embedding(vocab, cfg.d_model) self.pos_emb = nn.Embedding(max_seq, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList([ TransformerBlock(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout) for _ in range(cfg.n_layers) ]) self.ln_f = nn.LayerNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, vocab, bias=False) self.head.weight = self.tok_emb.weight def forward(self, ids: torch.Tensor) -> torch.Tensor: B, T = ids.shape pos = torch.arange(T, device=ids.device).unsqueeze(0) x = self.drop(self.tok_emb(ids) + self.pos_emb(pos)) mask = torch.triu(torch.ones(T, T, device=ids.device, dtype=torch.bool), 1) for block in self.blocks: x = block(x, mask) return self.head(self.ln_f(x)) class TransformerXLLM100M(nn.Module): """Transformer-XL at 100M scale with segment-level recurrence. Note: Uses learned absolute position embeddings (applied only to the current segment) rather than the relative position encodings of the original Dai et al. 2019 paper. This is a simplification — memory tokens carry positional information from their original segment. The attention mask gives memory tokens full visibility while maintaining causal masking within the current segment. """ def __init__(self, cfg: TransformerXLConfig, vocab: int, max_seq: int): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(vocab, cfg.d_model) self.pos_emb = nn.Embedding(max_seq, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList([ TransformerBlock(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout) for _ in range(cfg.n_layers) ]) self.ln_f = nn.LayerNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, vocab, bias=False) self.head.weight = self.tok_emb.weight self.mems: List[Optional[torch.Tensor]] = [None] * cfg.n_layers def reset_memory(self): self.mems = [None] * self.cfg.n_layers def forward(self, ids: torch.Tensor) -> torch.Tensor: B, T = ids.shape pos = torch.arange(T, device=ids.device).unsqueeze(0) x = self.drop(self.tok_emb(ids) + self.pos_emb(pos)) for i, block in enumerate(self.blocks): mem = self.mems[i] if mem is not None: M = mem.size(1) cat = torch.cat([mem, x], dim=1) else: M = 0 cat = x S = M + T mask = torch.zeros(S, S, device=ids.device, dtype=torch.bool) if T > 1: mask[M:, M:] = torch.triu( torch.ones(T, T, device=ids.device, dtype=torch.bool), 1) if self.training: out = torch.utils.checkpoint.checkpoint( block, cat, mask, use_reentrant=False) else: out = block(cat, mask) x = out[:, -T:] with torch.no_grad(): self.mems[i] = x[:, -self.cfg.mem_len:].detach() return self.head(self.ln_f(x)) class RetNetLM100M(nn.Module): """Multi-scale retention LM at 100M scale. Uses learned absolute position embeddings instead of the original RetNet paper's xPos. The decay matrix provides temporal weighting; absolute embeddings provide positional discrimination — a simpler but effective substitute for a benchmark comparison. """ def __init__(self, cfg: RetNetConfig, vocab: int, max_seq: int): super().__init__() self.tok_emb = nn.Embedding(vocab, cfg.d_model) self.pos_emb = nn.Embedding(max_seq, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.layers = nn.ModuleList([ RetentionLayer(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout) for _ in range(cfg.n_layers) ]) self.ln_f = nn.LayerNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, vocab, bias=False) self.head.weight = self.tok_emb.weight def forward(self, ids: torch.Tensor) -> torch.Tensor: B, T = ids.shape pos = torch.arange(T, device=ids.device).unsqueeze(0) x = self.drop(self.tok_emb(ids) + self.pos_emb(pos)) for layer in self.layers: x = layer(x) return self.head(self.ln_f(x)) class RMSNorm(nn.Module): """Root Mean Square Layer Normalization (used by Mamba).""" def __init__(self, d: int, eps: float = 1e-5): super().__init__() self.weight = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight class SelectiveSSM(nn.Module): """Selective state-space model (S6) — pure PyTorch implementation. Faithful to Gu & Dao 2023: per-channel delta via low-rank dt_proj, input-dependent B/C, S4D-Lin initialization for A, parallel scan. Core: h[t] = A_bar[t] * h[t-1] + B_bar[t] * x[t] y[t] = C[t] * h[t] """ def __init__(self, d_inner: int, d_state: int = 16, d_conv: int = 4, dt_rank: Optional[int] = None): super().__init__() self.d_inner = d_inner self.d_state = d_state if dt_rank is None: dt_rank = math.ceil(d_inner / 16) self.dt_rank = dt_rank self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv, padding=d_conv - 1, groups=d_inner) self.x_proj = nn.Linear(d_inner, dt_rank + d_state * 2, bias=False) self.dt_proj = nn.Linear(dt_rank, d_inner) dt = torch.exp( torch.rand(d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) ) inv_dt = dt + torch.log(-torch.expm1(-dt)) with torch.no_grad(): self.dt_proj.bias.copy_(inv_dt) self.A_log = nn.Parameter( torch.log(torch.arange(1, d_state + 1, dtype=torch.float32) .unsqueeze(0).expand(d_inner, -1).clone()) ) self.D = nn.Parameter(torch.ones(d_inner)) def forward(self, x: torch.Tensor) -> torch.Tensor: """x: [B, T, d_inner] -> y: [B, T, d_inner]""" B, T, D = x.shape N = self.d_state x_conv = self.conv1d(x.transpose(1, 2))[:, :, :T].transpose(1, 2) x_conv = F.silu(x_conv) proj = self.x_proj(x_conv) dt_input, B_input, C_input = proj.split( [self.dt_rank, N, N], dim=-1) delta = F.softplus(self.dt_proj(dt_input)) A = -torch.exp(self.A_log) A_bar = torch.exp( delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0) ) B_bar = delta.unsqueeze(-1) * B_input.unsqueeze(2) a = A_bar b = B_bar * x_conv.unsqueeze(-1) h_all = self._chunked_scan(a, b) y = (h_all * C_input[:, :, None, :]).sum(-1) y = y + x * self.D return y @staticmethod def _chunked_scan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Chunked sequential scan for linear recurrence h[t] = a[t]*h[t-1] + b[t]. Sequential within chunks of 64 steps, state propagated between chunks. O(T) total iterations with good cache locality. Args: a: [B, T, D, N] multiplicative coefficients b: [B, T, D, N] additive terms Returns: h: [B, T, D, N] all hidden states """ T = a.shape[1] CHUNK = 64 h_list = [] h_prev = torch.zeros_like(a[:, 0]) for start in range(0, T, CHUNK): end = min(start + CHUNK, T) a_chunk = a[:, start:end] b_chunk = b[:, start:end] h_chunk = [] h = h_prev for t in range(end - start): h = a_chunk[:, t] * h + b_chunk[:, t] h_chunk.append(h) h_prev = h h_list.append(torch.stack(h_chunk, dim=1)) return torch.cat(h_list, dim=1)[:, :T] class MambaBlock(nn.Module): """Single Mamba block: in_proj -> SSM -> out_proj with gating.""" def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4, expand: int = 2): super().__init__() d_inner = d_model * expand self.ln = RMSNorm(d_model) self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False) self.ssm = SelectiveSSM(d_inner, d_state, d_conv) self.out_proj = nn.Linear(d_inner, d_model, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: residual = x x = self.ln(x) xz = self.in_proj(x) x_ssm, z = xz.chunk(2, dim=-1) x_ssm = self.ssm(x_ssm) x_ssm = x_ssm * F.silu(z) return residual + self.out_proj(x_ssm) class MambaLM100M(nn.Module): """Mamba language model at 100M scale.""" def __init__(self, cfg: MambaConfig, vocab: int, max_seq: int): super().__init__() self.tok_emb = nn.Embedding(vocab, cfg.d_model) self.blocks = nn.ModuleList([ MambaBlock(cfg.d_model, cfg.d_state, cfg.d_conv, cfg.expand) for _ in range(cfg.n_layers) ]) self.ln_f = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, vocab, bias=False) self.head.weight = self.tok_emb.weight def forward(self, ids: torch.Tensor) -> torch.Tensor: x = self.tok_emb(ids) for block in self.blocks: x = block(x) return self.head(self.ln_f(x)) class LSTMLM100M(nn.Module): """LSTM language model at 100M scale.""" def __init__(self, cfg: LSTMConfig, vocab: int, max_seq: int): super().__init__() self.tok_emb = nn.Embedding(vocab, cfg.d_model) self.lstm = nn.LSTM( cfg.d_model, cfg.hidden_size, cfg.n_layers, batch_first=True, dropout=cfg.dropout if cfg.n_layers > 1 else 0.0) for name, param in self.lstm.named_parameters(): if 'bias' in name: n = param.size(0) // 4 param.data[n:2*n].fill_(1.0) self.ln_f = nn.LayerNorm(cfg.hidden_size) self.proj = nn.Linear(cfg.hidden_size, cfg.d_model, bias=False) self.head = nn.Linear(cfg.d_model, vocab, bias=False) self.head.weight = self.tok_emb.weight def forward(self, ids: torch.Tensor) -> torch.Tensor: x = self.tok_emb(ids) x, _ = self.lstm(x) x = self.proj(self.ln_f(x)) return self.head(x) ALL_MODELS = ["CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m", "CEDL-V3m2", "CEDL-V3m2-NLL", "CEDL-V3lex", "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold", "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M", "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject", "CEDL-V4d-strong", "CEDL-V4d-strong-noinject", "CEDL-V4d-causal", "CEDL-V4d-causal-strong", "CEDL-V4d-causalz", "CEDL-V4d-combo", "CEDL-V4e", "CEDL-V4e-noinject", "CEDL-V4e-k1", "CEDL-V4e-k2", "CEDL-V4e-k4", "CEDL-V4e-k16", "CEDL-V4e-k32", "CEDL-V5-LM", "CEDL-V6", "CEDL-B0", "Transformer", "Transformer-XL", "RetNet", "Mamba"] def build_model(tag: str, vocab: int, max_seq: int, *, v6_mixture: bool = False, v6_lambda_init: float = -4.0, v6_lambda_a_init: float = 1.0, v6_aux_weight: float = 0.0, v6_margin_target: float = 1.0, v6_mix_weight: float = 0.5, v6_gate_weight: float = 0.01, v6_lambda_floor: float = 0.05, v6_lambda_head: bool = False, v6_lambda_head_hidden: int = 160, v6_lambda_head_bias_init: float = -7.0, v6_bg_weight: float = 1.0, v6_bg_target: float = 0.01, v6_bce_objective: bool = False, v6_sel_weight: float = 1.0, v6_lambda_head_w_init_std: float = 1e-3, v6_wt_sparsity_weight: float = 0.0, v6_wt_sparsity_target: float = 0.0, v6_mem_head_bank: bool = False, v6_bank_ce_weight: float = 0.0, v6_bank_pair_ce_weight: float = 0.0, v6_bank_query_source: str = "h_d", v6_bank_readout_mode: str = "bank", v6_source_adapter: bool = False, v6_context_adapter: bool = False, v6_specialist_noinject: bool = False, lambda_init: float | None = None, lambda_a_init: float | None = None, aux_weight: float | None = None, margin_target: float | None = None, mix_weight: float | None = None, gate_weight: float | None = None, lambda_floor: float | None = None, lambda_head: bool | None = None, lambda_head_hidden: int | None = None, lambda_head_bias_init: float | None = None, bg_weight: float | None = None, bg_target: float | None = None, bce_objective: bool | None = None, sel_weight: float | None = None, lambda_head_w_init_std: float | None = None, wt_sparsity_weight: float | None = None, wt_sparsity_target: float | None = None, memory_head_enabled: bool | None = None, memory_ce_weight: float | None = None, memory_pair_ce_weight: float | None = None, memory_query_source: str | None = None, memory_readout_mode: str | None = None, source_adapter: bool | None = None, context_adapter: bool | None = None, specialist_noinject: bool | None = None) -> nn.Module: """Build a CEDL model by tag.""" if lambda_init is not None: v6_lambda_init = lambda_init if lambda_a_init is not None: v6_lambda_a_init = lambda_a_init if aux_weight is not None: v6_aux_weight = aux_weight if margin_target is not None: v6_margin_target = margin_target if mix_weight is not None: v6_mix_weight = mix_weight if gate_weight is not None: v6_gate_weight = gate_weight if lambda_floor is not None: v6_lambda_floor = lambda_floor if lambda_head is not None: v6_lambda_head = lambda_head if lambda_head_hidden is not None: v6_lambda_head_hidden = lambda_head_hidden if lambda_head_bias_init is not None: v6_lambda_head_bias_init = lambda_head_bias_init if bg_weight is not None: v6_bg_weight = bg_weight if bg_target is not None: v6_bg_target = bg_target if bce_objective is not None: v6_bce_objective = bce_objective if sel_weight is not None: v6_sel_weight = sel_weight if lambda_head_w_init_std is not None: v6_lambda_head_w_init_std = lambda_head_w_init_std if wt_sparsity_weight is not None: v6_wt_sparsity_weight = wt_sparsity_weight if wt_sparsity_target is not None: v6_wt_sparsity_target = wt_sparsity_target if memory_head_enabled is not None: v6_mem_head_bank = memory_head_enabled if memory_ce_weight is not None: v6_bank_ce_weight = memory_ce_weight if memory_pair_ce_weight is not None: v6_bank_pair_ce_weight = memory_pair_ce_weight if memory_query_source is not None: v6_bank_query_source = memory_query_source if memory_readout_mode is not None: v6_bank_readout_mode = memory_readout_mode if source_adapter is not None: v6_source_adapter = source_adapter if context_adapter is not None: v6_context_adapter = context_adapter if specialist_noinject is not None: v6_specialist_noinject = specialist_noinject if tag == "CEDL": tag = "CEDL-V6" if tag == "CEDL-legacy": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v0"), vocab, max_seq) elif tag == "CEDL-V2a": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v2a"), vocab, max_seq) elif tag == "CEDL-V3": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v3", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V3m": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v3m", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V3m2": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v3m2", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V3m2-NLL": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v3m2_nll", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4c": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c", v4c_variant="base", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4c-G": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c_gate", v4c_variant="base", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4c-cold": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c", v4c_variant="cold", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4c-randlabel": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c", v4c_variant="randlabel", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4c-M": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0), vocab, max_seq) elif tag == "CEDL-V4c-frozen": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4c", v4c_variant="frozen", lex_anchor_weight=0.0), vocab, max_seq) elif tag == "CEDL-V4d": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.02, v4d_logit_cap=4.0, v4d_noinject=False), vocab, max_seq) elif tag == "CEDL-V4d-strong": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False), vocab, max_seq) elif tag == "CEDL-V4d-causalz": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False, v4d_causal_weight=0.0, v4d_causal_z_weight=0.10, v4d_causal_z_gap=0.15), vocab, max_seq) elif tag == "CEDL-V4d-combo": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False, v4d_causal_weight=0.10, v4d_causal_gap=0.25, v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15), vocab, max_seq) elif tag == "CEDL-V4e" or tag == "CEDL-V4e-noinject" \ or tag.startswith("CEDL-V4e-k"): _noinj = (tag == "CEDL-V4e-noinject") _rank = 8 if tag.startswith("CEDL-V4e-k"): try: _rank = int(tag.rsplit("k", 1)[1]) except ValueError: _rank = 8 return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=_noinj, v4d_causal_weight=0.10, v4d_causal_gap=0.25, v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15, v4d_w_sal_rank=_rank, v4d_swap_consistency_weight=0.05, v4d_role_sep_weight=0.05), vocab, max_seq) elif tag == "CEDL-V5-LM": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False, v4d_causal_weight=0.10, v4d_causal_gap=0.25, v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15, v4d_w_sal_rank=8, v4d_swap_consistency_weight=0.05, v4d_role_sep_weight=0.05), vocab, max_seq) elif tag == "CEDL-V6": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=bool(v6_specialist_noinject), v4d_causal_weight=0.10, v4d_causal_gap=0.25, v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15, v4d_w_sal_rank=8, v4d_swap_consistency_weight=0.05, v4d_role_sep_weight=0.05, v6_mixture=True, v6_lambda_init=float(v6_lambda_init), v6_lambda_a_init=float(v6_lambda_a_init), v6_aux_weight=float(v6_aux_weight), v6_margin_target=float(v6_margin_target), v6_mix_weight=float(v6_mix_weight), v6_gate_weight=float(v6_gate_weight), v6_lambda_floor=float(v6_lambda_floor), v6_lambda_head=bool(v6_lambda_head), v6_lambda_head_hidden=int(v6_lambda_head_hidden), v6_lambda_head_bias_init=float(v6_lambda_head_bias_init), v6_bg_weight=float(v6_bg_weight), v6_bg_target=float(v6_bg_target), v6_bce_objective=bool(v6_bce_objective), v6_sel_weight=float(v6_sel_weight), v6_lambda_head_w_init_std=float(v6_lambda_head_w_init_std), v6_wt_sparsity_weight=float(v6_wt_sparsity_weight), v6_wt_sparsity_target=float(v6_wt_sparsity_target), v6_mem_head_bank=bool(v6_mem_head_bank), v6_bank_ce_weight=float(v6_bank_ce_weight), v6_bank_pair_ce_weight=float(v6_bank_pair_ce_weight), v6_bank_query_source=str(v6_bank_query_source), v6_bank_readout_mode=str(v6_bank_readout_mode), v6_source_adapter=bool(v6_source_adapter), v6_context_adapter=bool(v6_context_adapter)), vocab, max_seq) elif tag == "CEDL-V4d-causal-strong": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False, v4d_causal_weight=0.10, v4d_causal_gap=0.25), vocab, max_seq) elif tag == "CEDL-V4d-causal": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=False, v4d_causal_weight=0.02, v4d_causal_gap=0.25), vocab, max_seq) elif tag == "CEDL-V4d-strong-noinject": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0, v4d_noinject=True), vocab, max_seq) elif tag == "CEDL-V4d-cold": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.0, v4d_logit_cap=4.0, v4d_noinject=False), vocab, max_seq) elif tag == "CEDL-V4d-noinject": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v4d", v4c_variant="base", lex_anchor_weight=0.0, v4c_margin_weight=0.05, v4c_margin_target=1.0, v4d_w_sal_sigma=0.02, v4d_logit_cap=4.0, v4d_noinject=True), vocab, max_seq) elif tag == "CEDL-V3lex": return CEDLTwoLoop100M( CEDLConfig(use_salience=True, salience_mode="v3", lex_anchor_weight=0.005), vocab, max_seq) elif tag == "CEDL-B0": return CEDLTwoLoop100M( CEDLConfig(use_salience=False), vocab, max_seq) elif tag == "Transformer": return TransformerLM100M(TransformerConfig(), vocab, max_seq) elif tag == "Transformer-XL": return TransformerXLLM100M(TransformerXLConfig(), vocab, max_seq) elif tag == "RetNet": return RetNetLM100M(RetNetConfig(), vocab, max_seq) elif tag == "Mamba": return MambaLM100M(MambaConfig(), vocab, max_seq) elif tag == "LSTM": return LSTMLM100M(LSTMConfig(), vocab, max_seq) else: raise ValueError(f"Unknown model: {tag}") def count_params(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters()) def verify_all_params(): """Print parameter counts for all models.""" print("=" * 60) print("Parameter Verification (target: ~100M)") print("=" * 60) for tag in ALL_MODELS: model = build_model(tag, 50257, 1024) n = count_params(model) print(f" {tag:20s}: {n:>12,d} ({n/1e6:.1f}M)") del model print("=" * 60) class TextChunkDataset(Dataset): """Pre-tokenized text split into fixed-length chunks.""" def __init__(self, token_ids: torch.Tensor, chunk_size: int): self.chunk_size = chunk_size n_chunks = len(token_ids) // chunk_size self.data = token_ids[:n_chunks * chunk_size].view(n_chunks, chunk_size) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def load_data(cfg: Config): """Load and tokenize dataset. Returns train/val/test DataLoaders.""" from transformers import GPT2TokenizerFast tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") def tokenize_rows(rows): """Tokenize row-by-row to avoid OOM from giant string join.""" all_ids = [] for text in rows: if text.strip(): ids = tokenizer.encode(text, add_special_tokens=False) all_ids.extend(ids) return torch.tensor(all_ids, dtype=torch.long) if cfg.dataset == "wikitext103": from datasets import load_dataset ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1") print("Tokenizing train split...") train_ids = tokenize_rows(ds["train"]["text"]) print("Tokenizing val split...") val_ids = tokenize_rows(ds["validation"]["text"]) print("Tokenizing test split...") test_ids = tokenize_rows(ds["test"]["text"]) del ds elif cfg.dataset == "c4": from datasets import load_dataset ds_train = load_dataset("allenai/c4", "en", split="train", streaming=True) ds_val = load_dataset("allenai/c4", "en", split="validation", streaming=True) all_ids = [] total_tokens = 0 for example in ds_train: ids = tokenizer.encode(example["text"], add_special_tokens=False) all_ids.extend(ids) total_tokens += len(ids) if total_tokens > 400_000_000: break train_ids = torch.tensor(all_ids, dtype=torch.long) del all_ids val_ids_list = [] val_tokens = 0 for example in ds_val: ids = tokenizer.encode(example["text"], add_special_tokens=False) val_ids_list.extend(ids) val_tokens += len(ids) if val_tokens > 2_000_000: break val_ids = torch.tensor(val_ids_list, dtype=torch.long) del val_ids_list test_ids = val_ids else: raise ValueError(f"Unknown dataset: {cfg.dataset}") chunk = cfg.max_seq + 1 train_ds = TextChunkDataset(train_ids, chunk) val_ds = TextChunkDataset(val_ids, chunk) test_ds = TextChunkDataset(test_ids, chunk) pin = not cfg.tpu train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=4, pin_memory=pin, drop_last=True, persistent_workers=True) val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=pin, drop_last=True, persistent_workers=True) test_loader = DataLoader(test_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=pin, drop_last=True, persistent_workers=True) print(f"Dataset: {cfg.dataset}") print(f" Train chunks: {len(train_ds):,}") print(f" Val chunks: {len(val_ds):,}") print(f" Test chunks: {len(test_ds):,}") return train_loader, val_loader, test_loader, tokenizer def make_loaders(train_ds, val_ds, test_ds, batch_size: int, tpu: bool = False, shuffle_seed: Optional[int] = None): """Build DataLoaders from pre-tokenized datasets (fast, no re-tokenization). If shuffle_seed is provided, the train_loader uses a seeded torch.Generator so batch ordering is identical across model tags in a paired run. """ pin = not tpu gen = None if shuffle_seed is not None: gen = torch.Generator() gen.manual_seed(shuffle_seed) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=pin, drop_last=True, persistent_workers=True, generator=gen) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=pin, drop_last=True, persistent_workers=True) test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=pin, drop_last=True, persistent_workers=True) return train_loader, val_loader, test_loader def get_lr(step: int, cfg: Config) -> float: """Cosine schedule with linear warmup.""" if step < cfg.warmup_steps: return cfg.lr * step / cfg.warmup_steps if step >= cfg.max_steps: return cfg.min_lr progress = (step - cfg.warmup_steps) / (cfg.max_steps - cfg.warmup_steps) return cfg.min_lr + 0.5 * (cfg.lr - cfg.min_lr) * (1 + math.cos(math.pi * progress)) def _v5_effective_v4c_every(gstep: int, cfg=None) -> int: """Synthetic-batch cadence by GLOBAL step (early/mid/late).""" e = int(getattr(cfg, "v5_cadence_early", 8)) if cfg else 8 m = int(getattr(cfg, "v5_cadence_mid", 16)) if cfg else 16 l_ = int(getattr(cfg, "v5_cadence_late", 32)) if cfg else 32 if gstep < 10000: return e if gstep < 20000: return m return l_ def _v5_aux_scale(gstep: int, cfg=None) -> float: """Global aux REPLAY strength applied to the TOTAL aux loss (early/mid/late).""" e = float(getattr(cfg, "v5_aux_scale_early", 0.50)) if cfg else 0.50 m = float(getattr(cfg, "v5_aux_scale_mid", 0.25)) if cfg else 0.25 l_ = float(getattr(cfg, "v5_aux_scale_late", 0.10)) if cfg else 0.10 if gstep < 10000: return e if gstep < 20000: return m return l_ _V5_SAL_PREFIXES = ("salience.", "d_stage.salience_bias_", "d_stage.w_sal", "v4d_role_head.", "bank_q_proj.", "v6_lambda_", "mem_head_bank.", "v6_source_adapter.", "v6_context_adapter.") def _v5_salience_param_ids(raw_model): """ids + names of AUX-ELIGIBLE params (V5 gradient isolation). Matches salience.*, d_stage.salience_bias_*, d_stage.w_sal* (incl. w_sal_A/B), v4d_role_head.* (parent-level), AND (V6.1) the bank-readout params bank_q_proj.* + v6_lambda_* — the latter are bank-readout semantically (not salience), but routed through the same isolation allowlist so the V4c-aux gradient survives the restore loop. d_stage.cross_attn.* is NOT matched → trunk.""" ids, names = set(), [] for n, p in raw_model.named_parameters(): if any(n.startswith(pre) for pre in _V5_SAL_PREFIXES): ids.add(id(p)) names.append(n) return ids, names def _unwrap_model_state_dict(obj, src_label: str = "checkpoint"): """Return a plain model state_dict from bare, train-state, or wrapper ckpts.""" if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): wrapper_keys = sorted(k for k in obj.keys() if k != "model") wrapper_msg = f" wrapper_keys={wrapper_keys}" if wrapper_keys else "" print(f" [load] unwrapping {src_label}['model'] -> model weights{wrapper_msg}") obj = obj["model"] if not isinstance(obj, dict): raise RuntimeError(f"{src_label} did not contain a state_dict-like object") if any(k.startswith("_orig_mod.") for k in obj): obj = {k.replace("_orig_mod.", ""): v for k, v in obj.items()} return obj def _load_model_state_from_path(path: str, src_label: str = "checkpoint"): state = torch.load(path, map_location="cpu", weights_only=True) return _unwrap_model_state_dict(state, src_label=src_label) def _is_v6_specialist_b0_missing_key(name: str) -> bool: """Keys present in a CEDL-V6 specialist shell but absent from B0.""" exact = { "sl_alpha", "val_a", "val_b", "v6_lambda_a", "v6_lambda_b", } prefixes = ( "salience.", "d_stage.salience_bias_", "d_stage.w_sal", "v4d_role_head.", "bank_q_proj.", "v6_lambda_head.", "mem_head_bank.", "v6_source_adapter.", "v6_context_adapter.", ) return name in exact or any(name.startswith(pre) for pre in prefixes) def _apply_v6_specialist_freeze(raw_model): """Freeze the B0 trunk and leave only the direct specialist trainable.""" trainable_prefixes = ["v6_lambda_head.", "mem_head_bank."] if not hasattr(raw_model, "v6_lambda_head"): raise RuntimeError("V6 specialist freeze requires v6_lambda_head.*") if not hasattr(raw_model, "mem_head_bank"): raise RuntimeError("V6 specialist freeze requires mem_head_bank.*") if getattr(raw_model, "use_v6_source_adapter", False): if not hasattr(raw_model, "v6_source_adapter"): raise RuntimeError( "V6 specialist freeze expected v6_source_adapter.*") trainable_prefixes.append("v6_source_adapter.") if getattr(raw_model, "use_v6_context_adapter", False): if not hasattr(raw_model, "v6_context_adapter"): raise RuntimeError( "V6 specialist freeze expected v6_context_adapter.*") trainable_prefixes.append("v6_context_adapter.") trainable_prefixes = tuple(trainable_prefixes) trainable_names = [] frozen_names = [] for n, p in raw_model.named_parameters(): keep_trainable = any(n.startswith(pre) for pre in trainable_prefixes) p.requires_grad = keep_trainable if keep_trainable: trainable_names.append(n) else: frozen_names.append(n) raw_model._v6_specialist_trainable_prefixes = list(trainable_prefixes) raw_model._v6_specialist_trainable_names = trainable_names n_trainable = sum(p.numel() for p in raw_model.parameters() if p.requires_grad) n_frozen = sum(p.numel() for p in raw_model.parameters() if not p.requires_grad) print(f" [V6 specialist freeze] trainable params={n_trainable:,} " f"frozen params={n_frozen:,} " f"trainable_prefixes={list(trainable_prefixes)}") print(f" [V6 specialist freeze] trainable names={trainable_names}") @torch.no_grad() def evaluate(model: nn.Module, loader: DataLoader, device: torch.device, max_steps: int = 200, tpu: bool = False) -> float: """Compute cross-entropy loss → perplexity.""" _xm = None if tpu: import torch_xla.core.xla_model as _xm model.eval() if hasattr(model, 'reset_memory'): model.reset_memory() total_loss = 0.0 count = 0 _raw = model._orig_mod if hasattr(model, '_orig_mod') else model _use_logprobs = bool(getattr(_raw, "outputs_log_probs", False) and not getattr(_raw, "bank_off", False)) _lam_acc = {"lam_mean": 0.0, "lam_std": 0.0, "lam_sat_low": 0.0, "lam_sat_high": 0.0} _lam_n = 0 _lam_max_global = -1.0 _lam_gt_0_7_sum = 0 _lam_total_sum = 0 _lam_tensors = [] for i, batch in enumerate(loader): if i >= max_steps: break batch = batch.to(device) input_ids = batch[:, :-1] targets = batch[:, 1:] logits = model(input_ids) if isinstance(logits, tuple): logits = logits[0] if _use_logprobs: loss = F.nll_loss(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) _ls = getattr(_raw, "_last_lam_stats", None) if _ls is not None: for _k in _lam_acc: _lam_acc[_k] += float(_ls[_k]) _lam_n += 1 _lam_max_global = max(_lam_max_global, float(_ls["lam_max"])) _lam_gt_0_7_sum += int(_ls["lam_gt_0_7_count"]) _lam_total_sum += int(_ls["lam_total_count"]) _lam_tensors.append(_ls["lam_tensor_cpu"]) else: loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) total_loss += loss.item() del logits, loss count += 1 if _xm is not None: _xm.mark_step() if _lam_n > 0: _raw._last_eval_lam_avg = {k: v / _lam_n for k, v in _lam_acc.items()} _raw._last_eval_lam_avg["n_batches"] = _lam_n _raw._last_eval_lam_avg["lam_max"] = _lam_max_global _raw._last_eval_lam_avg["lam_frac_gt_0_7"] = ( _lam_gt_0_7_sum / _lam_total_sum if _lam_total_sum > 0 else 0.0) try: _pooled = torch.cat(_lam_tensors).numpy() _raw._last_eval_lam_avg["lam_p99"] = float( np.percentile(_pooled, 99)) except Exception as _pe: print(f" [V6 λ] WARN: lam_p99 computation failed ({_pe})") _raw._last_eval_lam_avg["lam_p99"] = float("nan") _raw._last_eval_lam_avg["v6_lambda_a"] = float(_raw.v6_lambda_a.item()) _raw._last_eval_lam_avg["v6_lambda_b"] = float(_raw.v6_lambda_b.item()) _use_head = bool(getattr(_raw, "use_v6_lambda_head", False)) if _use_head: _raw._last_eval_lam_avg["v6_lambda_head_bias"] = float( _raw.v6_lambda_head[2].bias.item()) _avg = _raw._last_eval_lam_avg print(f" [V6 λ] mean={_avg['lam_mean']:.4f} " f"std={_avg['lam_std']:.4f} " f"sat_low={_avg['lam_sat_low']:.3f} " f"sat_high={_avg['lam_sat_high']:.3f} " f"(over {_lam_n} eval batches)") if _use_head: print(f" [V6 λ tail] max={_avg['lam_max']:.4f} " f"p99={_avg['lam_p99']:.4f} " f"frac>0.7={_avg['lam_frac_gt_0_7']:.4f} " f"(head_bias={_avg['v6_lambda_head_bias']:+.3f})") else: print(f" [V6 λ tail] max={_avg['lam_max']:.4f} " f"p99={_avg['lam_p99']:.4f} " f"frac>0.7={_avg['lam_frac_gt_0_7']:.4f} " f"(a={_avg['v6_lambda_a']:+.3f}, b={_avg['v6_lambda_b']:+.3f})") elif _use_logprobs: print(" [V6 λ] WARN: outputs_log_probs=True but no _last_lam_stats " "recorded — forward block may not have executed the λ path.") model.train() avg_loss = total_loss / max(count, 1) return math.exp(avg_loss) def train(model: nn.Module, tag: str, cfg: Config, device: torch.device, train_loader: DataLoader, val_loader: DataLoader): """Main training loop with gradient accumulation and cosine LR.""" xm = None if cfg.tpu: import torch_xla.core.xla_model as xm model.train() raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model is_cedl = isinstance(raw_model, CEDLTwoLoop100M) _resume_ckpt = getattr(cfg, "resume_checkpoint", None) if is_cedl and getattr(raw_model, "v4c_variant", "base") == "frozen": for n, p in raw_model.named_parameters(): in_sl_mod = (n.startswith("salience.") or n.startswith("d_stage.salience_bias_")) p.requires_grad = bool(in_sl_mod) n_trainable = sum(1 for p in raw_model.parameters() if p.requires_grad) n_frozen = sum(1 for p in raw_model.parameters() if not p.requires_grad) print(f" [V4c-frozen] {n_trainable} trainable params, " f"{n_frozen} frozen (LM backbone)") _specialist_from_b0 = ( is_cedl and bool(getattr(cfg, "v6_specialist_from_b0", False))) if _specialist_from_b0: if not _resume_ckpt: raise RuntimeError( "v6_specialist_from_b0=True but cfg.resume_checkpoint is empty") if bool(getattr(cfg, "v6_specialist_noinject", False)): _noinj_live = bool(getattr(raw_model._v4c_cfg, "v4d_noinject", False)) if not _noinj_live: raise RuntimeError( "v6_specialist_noinject=True but model was built with " "v4d_noinject=False") print(f" [V6 specialist load] source B0 checkpoint: {_resume_ckpt}") _b0_sd = _load_model_state_from_path(_resume_ckpt, src_label="B0") _res = raw_model.load_state_dict(_b0_sd, strict=False) _miss = set(_res.missing_keys) _unexp = set(_res.unexpected_keys) _bad_miss = sorted(k for k in _miss if not _is_v6_specialist_b0_missing_key(k)) if _bad_miss or _unexp: raise RuntimeError( "[V6 specialist load] B0->V6 audit failed: " f"unexpected_missing={_bad_miss}, " f"unexpected={sorted(_unexp)}") _fresh = sorted(_miss) raw_model._v6_specialist_loaded_missing = _fresh raw_model._v6_specialist_base_checkpoint = _resume_ckpt print(f" [V6 specialist load] matched B0 trunk; freshly initialized " f"{len(_fresh)} specialist keys") if _fresh: print(f" [V6 specialist load] fresh keys={_fresh}") if int(getattr(cfg, "resume_start_step", 0) or 0) != 0: print(" [V6 specialist load] WARN: ignoring --start-step for " "fresh B0-specialist bootstrap") cfg.resume_start_step = 0 _resume_ckpt = None if _specialist_from_b0 and getattr(cfg, "v6_specialist_freeze", "none") == "trunk": _apply_v6_specialist_freeze(raw_model) if is_cedl: c_param_ids = {id(p) for p in raw_model.c_stage.parameters() if p.requires_grad} c_params = [p for p in raw_model.c_stage.parameters() if p.requires_grad] _bank_head_lr = float(getattr(cfg, "v6_bank_head_lr", 0.0)) _use_bank_lr_group = (_bank_head_lr > 0 and bool(getattr(cfg, "v6_mem_head_bank", False))) bank_head_param_ids = set() if _use_bank_lr_group and hasattr(raw_model, "mem_head_bank"): bank_head_param_ids = { id(p) for p in raw_model.mem_head_bank.parameters() if p.requires_grad} if hasattr(raw_model, "v6_source_adapter"): bank_head_param_ids.update( id(p) for p in raw_model.v6_source_adapter.parameters() if p.requires_grad) if hasattr(raw_model, "v6_context_adapter"): bank_head_param_ids.update( id(p) for p in raw_model.v6_context_adapter.parameters() if p.requires_grad) edl_params = [p for p in raw_model.parameters() if p.requires_grad and id(p) not in c_param_ids and id(p) not in bank_head_param_ids] groups = [] if c_params: groups.append({"params": c_params, "lr": cfg.lr}) if edl_params: groups.append({"params": edl_params, "lr": cfg.lr * 1.5}) if _use_bank_lr_group: bank_head_params = [p for p in raw_model.mem_head_bank.parameters() if p.requires_grad] if hasattr(raw_model, "v6_source_adapter"): bank_head_params.extend( p for p in raw_model.v6_source_adapter.parameters() if p.requires_grad) if hasattr(raw_model, "v6_context_adapter"): bank_head_params.extend( p for p in raw_model.v6_context_adapter.parameters() if p.requires_grad) if bank_head_params: groups.append({ "params": bank_head_params, "lr": _bank_head_lr, "bank_head_lr_constant": True, }) print(f" [Stage2b-lr] direct readout params in dedicated AdamW group " f"with CONSTANT lr={_bank_head_lr} " f"({len(bank_head_params)} params, " f"{sum(p.numel() for p in bank_head_params):,} elements) " f"— cosine schedule applies to trunk/gate only.") optimizer = torch.optim.AdamW( groups, weight_decay=cfg.weight_decay, betas=(0.9, 0.95)) else: optimizer = torch.optim.AdamW( model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay, betas=(0.9, 0.95)) scaler = None step = int(getattr(cfg, "resume_start_step", 0) or 0) _resumed_best_ppl = None _resumed_run_id = None _arch_fork = False _v6_allowed_missing = {"bank_q_proj.weight", "bank_q_proj.bias", "v6_lambda_a", "v6_lambda_b"} if bool(getattr(cfg, "v6_lambda_head", False)) and is_cedl: _v6_allowed_missing |= { "v6_lambda_head.0.weight", "v6_lambda_head.0.bias", "v6_lambda_head.2.weight", "v6_lambda_head.2.bias", } if bool(getattr(cfg, "v6_mem_head_bank", False)) and is_cedl: _v6_allowed_missing |= { "mem_head_bank.weight", "mem_head_bank.bias", } _v6_mixture_on = bool(getattr(cfg, "v6_mixture", False)) and is_cedl def _audited_load(sd, src_label): result = raw_model.load_state_dict(sd, strict=False) miss = set(result.missing_keys) unexp = set(result.unexpected_keys) if _v6_mixture_on: extra_miss = miss - _v6_allowed_missing if extra_miss or unexp: raise RuntimeError( f"[V6 resume {src_label}] strict audit failed: " f"unexpected_missing={sorted(extra_miss)}, " f"unexpected={sorted(unexp)}") new_init = sorted(miss & _v6_allowed_missing) if new_init: print(f" [V6 resume {src_label}] freshly initialized: " f"{new_init}") else: if miss or unexp: raise RuntimeError( f"[resume {src_label}] strict load failed: " f"missing={sorted(miss)}, unexpected={sorted(unexp)}") if _resume_ckpt: _rs = torch.load(_resume_ckpt, map_location="cpu", weights_only=True) if isinstance(_rs, dict) and "model" in _rs and "optimizer" in _rs: _msd = _rs["model"] if any(k.startswith("_orig_mod.") for k in _msd): _msd = {k.replace("_orig_mod.", ""): v for k, v in _msd.items()} _audited_load(_msd, "trainstate") _src_tag = _rs.get("model_tag") _arch_fork = (_src_tag is not None and _src_tag != tag) try: optimizer.load_state_dict(_rs["optimizer"]) print(" [resume] restored optimizer state from train-state ckpt") except Exception as _e: print(f" [resume] WARN: optimizer state not restored ({_e}); " f"continuing with fresh AdamW") step = int(_rs.get("global_step", step)) if _arch_fork: print(f" [resume] ARCH-FORK detected: trainstate tag " f"{_src_tag!r} != current tag {tag!r}. Resetting " f"best_val_ppl=inf, best_step={step}, new run_started_at.") _resumed_best_ppl = None _resumed_run_id = None else: if _rs.get("best_val_ppl") is not None: _resumed_best_ppl = float(_rs["best_val_ppl"]) _resumed_run_id = _rs.get("run_started_at") print(f" [resume] restored best_val_ppl={_resumed_best_ppl} " f"run_id={_resumed_run_id}") else: if any(k.startswith("_orig_mod.") for k in _rs): _rs = {k.replace("_orig_mod.", ""): v for k, v in _rs.items()} _audited_load(_rs, "bare") print(" [resume] loaded bare state_dict; using FRESH AdamW " "(no optimizer state in source checkpoint)") print(f" [resume] {_resume_ckpt} global start_step={step} " f"target total_steps={cfg.max_steps}") _v5_iso = bool(getattr(cfg, "v5_grad_isolation", False)) and is_cedl _v5_trunk_frac = float(getattr(cfg, "v5_trunk_aux_frac", 0.0)) _sal_param_ids = set() if _v5_iso: _sal_param_ids, _sal_names = _v5_salience_param_ids(raw_model) _n_trunk = sum(1 for p in raw_model.parameters() if id(p) not in _sal_param_ids) _has_v4d = any(n.startswith("d_stage.w_sal") or n.startswith("v4d_role_head.") for n in _sal_names) if not _has_v4d: raise RuntimeError( "v5_grad_isolation=True but NO params matched d_stage.w_sal* / " "v4d_role_head.* — the salience set is wrong (a V5 model must " "have these). Aborting to avoid an invalid isolation run.") print(f" [V5 isolation] salience params={len(_sal_param_ids)} " f"trunk params={_n_trunk} trunk_aux_frac={_v5_trunk_frac} " f"(cross_attn stays trunk; aux-only trunk grads discarded)") if is_cedl and bool(getattr(cfg, "v6_mixture", False)): _v6_param_names = [n for n, _ in raw_model.named_parameters() if n.startswith("bank_q_proj.") or n in ("v6_lambda_a", "v6_lambda_b")] if len(_v6_param_names) == 0: raise RuntimeError( "cfg.v6_mixture=True but the model has NO bank_q_proj.* / " "v6_lambda_* parameters — V6 forward will be a silent no-op. " "Build the model with v6_mixture=True (use the CEDL-V6 tag, " "or pass v6_mixture=True to build_model for other CEDL tags).") _a_init = float(getattr(cfg, "v6_lambda_a_init", 1.0)) _a_live = float(raw_model.v6_lambda_a.item()) print(f" [V6 mixture] params={_v6_param_names} " f"v6_lambda_a_init={_a_init:+.3f} (live={_a_live:+.3f}) " f"λ_b_init=sigmoid({raw_model.v6_lambda_b.item():.3f})=" f"{torch.sigmoid(raw_model.v6_lambda_b).item():.4f} " f"(AUX-eligible via V5 allowlist; V4c bank-aux can reach)") for _k in ("v6_aux_weight", "v6_bg_weight", "v6_bg_target", "v6_lambda_head_hidden", "v6_lambda_head_bias_init", "v6_bce_objective", "v6_sel_weight", "v6_lambda_head_w_init_std", "v6_wt_sparsity_weight", "v6_wt_sparsity_target", "v6_mem_head_bank", "v6_bank_ce_weight", "v6_bank_pair_ce_weight", "v6_bank_query_source", "v6_bank_readout_mode", "v6_source_adapter", "v6_context_adapter"): _cv = getattr(cfg, _k, None) _mv = getattr(raw_model._v4c_cfg, _k, None) if isinstance(_cv, float) and isinstance(_mv, float): if abs(_cv - _mv) > 1e-9: raise RuntimeError( f"V6 propagation bug: cfg.{_k}={_cv} but " f"raw_model._v4c_cfg.{_k}={_mv}. The V4c hook reads " f"_v4c_cfg; without matching values the supervision " f"silently falls back to the default.") elif _cv != _mv: raise RuntimeError( f"V6 propagation bug: cfg.{_k}={_cv} but " f"raw_model._v4c_cfg.{_k}={_mv}.") _cfg_w = float(getattr(cfg, "v6_aux_weight", 0.0)) _mdl_w = float(getattr(raw_model._v4c_cfg, "v6_aux_weight", 0.0)) print(f" [V6.1] cfg.v6_aux_weight={_cfg_w} " f"_v4c_cfg.v6_aux_weight={_mdl_w} " f"v6_margin_target={getattr(raw_model._v4c_cfg, 'v6_margin_target', 1.0)} " f"v6_mix_weight={getattr(raw_model._v4c_cfg, 'v6_mix_weight', 0.5)} " f"v6_gate_weight={getattr(raw_model._v4c_cfg, 'v6_gate_weight', 0.01)} " f"v6_lambda_floor={getattr(raw_model._v4c_cfg, 'v6_lambda_floor', 0.05)} " f"(0 = V6 baseline; >0 = V6.1 aux-through-mixture active)") if getattr(raw_model, "use_v6_lambda_head", False): _head_bias = float(raw_model.v6_lambda_head[2].bias.item()) _bce_on = bool(getattr(raw_model._v4c_cfg, "v6_bce_objective", False)) _w_std = float(getattr(raw_model._v4c_cfg, "v6_lambda_head_w_init_std", 1e-3)) _w_live = float(raw_model.v6_lambda_head[2].weight.std().item()) print(f" [Stage2a] v6_lambda_head ON hidden=" f"{getattr(raw_model._v4c_cfg, 'v6_lambda_head_hidden', 160)} " f"head_bias={_head_bias:+.3f} " f"w_init_std={_w_std:.4f} (live_std={_w_live:.4f}) " f"bg_weight={getattr(raw_model._v4c_cfg, 'v6_bg_weight', 1.0)} " f"bg_target={getattr(raw_model._v4c_cfg, 'v6_bg_target', 0.01)} " f"(aux-only via torch.no_grad in main forward + V5 allowlist)") if _bce_on: print(f" [Stage2a-bce] BCE objective ON " f"sel_weight={getattr(raw_model._v4c_cfg, 'v6_sel_weight', 1.0)} " f"(L_gate + L_bg are telemetry-only; v4c_aux uses " f"L_bank + w_mix·L_mix + w_sel·L_sel where " f"L_sel = 0.5·BCE_ans + 0.5·BCE_bg)") _wt_w = float(getattr(raw_model._v4c_cfg, "v6_wt_sparsity_weight", 0.0)) if _wt_w > 0: _wt_t = float(getattr(raw_model._v4c_cfg, "v6_wt_sparsity_target", 0.0)) if _wt_t > 0: _wt_form = (f"hinge form aux_loss += {_wt_w}·relu(λ_wt" f" - {_wt_t}).mean() — free zone [0," f"{_wt_t}]") else: _wt_form = f"mean-λ form aux_loss += {_wt_w}·E[λ_wt]" print(f" [Stage2a-wts] WT sparsity loss ON " f"wt_sparsity_weight={_wt_w} " f"wt_sparsity_target={_wt_t} ({_wt_form}; head " f"receives gradient via parallel non-no_grad path; " f"mixture lam stays detached → LM-CE still cannot " f"drive head)") if getattr(raw_model, "use_v6_mem_head_bank", False): _bce_w = float(getattr(raw_model._v4c_cfg, "v6_bank_ce_weight", 0.0)) _pair_w = float(getattr(raw_model._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)) _mhb_w_std = float(raw_model.mem_head_bank.weight.std().item()) _mhb_b_mean = float(raw_model.mem_head_bank.bias.mean().item()) _tied = (raw_model.mem_head_bank.weight is raw_model.c_stage.tok_emb.weight) _bhlr = float(getattr(cfg, "v6_bank_head_lr", 0.0)) print(f" [Stage2b] mem_head_bank ON " f"weight.std={_mhb_w_std:.4f} " f"bias.mean={_mhb_b_mean:+.3f} " f"tied_to_tok_emb={_tied} (expect False) " f"bank_ce_weight={_bce_w} " f"bank_pair_ce_weight={_pair_w} " f"bank_query_source={getattr(raw_model, 'v6_bank_query_source', 'h_d')} " f"bank_readout_mode={getattr(raw_model, 'v6_bank_readout_mode', 'bank')} " f"source_adapter={getattr(raw_model, 'use_v6_source_adapter', False)} " f"context_adapter={getattr(raw_model, 'use_v6_context_adapter', False)} " f"bank_head_lr={_bhlr}{' (CONSTANT; separate AdamW group)' if _bhlr > 0 else ' (= main lr; cosine schedule)'} " f"(separate trainable bank vocab projection; aux-eligible " f"via V5 allowlist; L_v6 += {_bce_w}·CE(logits_bank_a, " f"cur_tok) at V4c answer rows)") assert not _tied, ( "Stage 2b: mem_head_bank.weight must be a separate " "nn.Parameter, not aliased to tok_emb.weight. " "Check __init__ — use .copy_() not assignment.") accum_loss = 0.0 best_val_ppl = _resumed_best_ppl if _resumed_best_ppl is not None else float('inf') best_step = step os.makedirs(cfg.save_dir, exist_ok=True) _run_started_at = _resumed_run_id or time.strftime("%Y%m%dT%H%M%S") _run_start_step = step def _save_trainstate(_gstep): """Resume-only train-state with best-so-far and run identity.""" torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "global_step": _gstep, "best_val_ppl": float(best_val_ppl), "best_step": int(best_step), "run_started_at": _run_started_at, "model_tag": tag}, os.path.join(cfg.save_dir, f"{tag}_trainstate_latest.pt")) print(f"\nTraining {tag} ({count_params(model)/1e6:.1f}M params)") print(f" Device: {device}") print(f" Effective batch: {cfg.batch_size * cfg.grad_accum} seqs " f"= {cfg.batch_size * cfg.grad_accum * cfg.max_seq / 1e3:.0f}K tokens") print(f" Schedule total (max_steps): {cfg.max_steps:,}" + (f" | run_until_step: {cfg.run_until_step:,}" if getattr(cfg, "run_until_step", None) else "")) t0 = time.time() data_iter = iter(train_loader) cedl_cfg = CEDLConfig() if is_cedl else None _v3m2_diagnostic_active = ( is_cedl and getattr(raw_model, 'use_salience', False) and getattr(raw_model, 'salience_mode', 'v0') in ("v3m2", "v3m2_nll") ) _v3m2_eval_batch = None if _v3m2_diagnostic_active: try: _v3m2_eval_batch = next(iter(val_loader)).to(device)[:16, :512] print(f" [v_vec_pathway] diagnostic ACTIVE " f"(eval batch: {tuple(_v3m2_eval_batch.shape)})") except Exception as _e: print(f" [v_vec_pathway] WARN: could not grab eval batch ({_e}); " f"diagnostic DISABLED") _v3m2_diagnostic_active = False _v4c_active = ( is_cedl and getattr(raw_model, 'use_salience', False) and getattr(raw_model, 'salience_mode', 'v0') in ( "v4c", "v4c_gate", "v4d") ) _v4c_pair_gen = None _v4c_tok = None _v4c_log_idx = 0 _v4c_last_lm_grad_norm = None _v4c_swap_fn = None _v4c_split = "all" _v4c_collision = 0.0 if _v4c_active: try: import data_v4c_pairs as _v4c_module from transformers import GPT2TokenizerFast as _V4cTok _v4c_pair_gen = _v4c_module.generate _v4c_tok = _V4cTok.from_pretrained("gpt2") v4c_variant = getattr(raw_model, "v4c_variant", "base") print(f" [V4c] ACTIVE variant={v4c_variant} " f"every={raw_model._v4c_cfg.v4c_every} steps " f"B={raw_model._v4c_cfg.v4c_batch_size} " f"T_max={raw_model._v4c_cfg.v4c_max_seq}") _v4e_on = (getattr(raw_model, "salience_mode", "v0") == "v4d" and (float(getattr(raw_model._v4c_cfg, "v4d_swap_consistency_weight", 0.0)) > 0 or float(getattr(raw_model._v4c_cfg, "v4d_role_sep_weight", 0.0)) > 0)) if _v4e_on: _v4c_swap_fn = _v4c_module.make_entity_swap _v4c_split = "train" _v4c_collision = 0.2 print(f" [V4e] swap/role losses ON split=train " f"hard_collision=0.2 " f"rank={getattr(raw_model._v4c_cfg, 'v4d_w_sal_rank', 0)}") except Exception as _e: raise RuntimeError( f"V4c setup failed for tag {tag!r} (model is " f"salience_mode={getattr(raw_model, 'salience_mode', '?')}): " f"{_e}. Check that data_v4c_pairs.py is on PYTHONPATH and " f"that the GPT2 tokenizer is available.") from _e _stop_step = int(cfg.run_until_step) if getattr(cfg, "run_until_step", None) \ else cfg.max_steps while step < _stop_step: optimizer.zero_grad() for micro in range(cfg.grad_accum): try: batch = next(data_iter) except StopIteration: if hasattr(model, 'reset_memory'): model.reset_memory() data_iter = iter(train_loader) batch = next(data_iter) batch = batch.to(device) input_ids = batch[:, :-1] targets = batch[:, 1:] autocast_device = 'xla' if cfg.tpu else device.type with torch.amp.autocast(autocast_device, dtype=torch.bfloat16, enabled=cfg.bfloat16): output = model(input_ids) if isinstance(output, tuple): logits, aux_loss = output[0], output[1] else: logits, aux_loss = output, torch.tensor(0.0, device=device) if bool(getattr(raw_model, "outputs_log_probs", False) and not getattr(raw_model, "bank_off", False)): loss = F.nll_loss(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) else: loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) loss = (loss + aux_loss) / cfg.grad_accum loss.backward() accum_loss += loss.item() del logits, loss if xm is not None: xm.mark_step() _v4c_eff_every = (_v5_effective_v4c_every(step, cfg) if _v5_iso else raw_model._v4c_cfg.v4c_every) if _v4c_active else 1 if _v4c_active and (step + 1) % _v4c_eff_every == 0: v4c_cfg_obj = raw_model._v4c_cfg _lm_grad_snapshot = { id(p): p.grad.detach().clone() for p in model.parameters() if p.grad is not None } _lm_grad_norm_sq = 0.0 for g in _lm_grad_snapshot.values(): _lm_grad_norm_sq += float((g * g).sum()) _v4c_last_lm_grad_norm = _lm_grad_norm_sq ** 0.5 try: _v4c_family_weights = ( {"but_update": 0.50, "however_revision": 0.20, "temporal_update": 0.15, "paraphrased_equiv": 0.05, "neutral_control": 0.10} if getattr(cfg, "v5_family_reweight", False) else None) _v4c_ids, _v4c_items, _v4c_swap_ids, _v4c_swap_items = \ _build_v4c_contrastive_batch( _v4c_pair_gen, _v4c_tok, B=v4c_cfg_obj.v4c_batch_size, T_max=v4c_cfg_obj.v4c_max_seq, randomize_positive_mapping=( getattr(raw_model, "v4c_variant", "base") == "randlabel"), seed=step * 7919 + 13, split=_v4c_split, hard_collision_frac=_v4c_collision, swap_fn=_v4c_swap_fn, family_weights=_v4c_family_weights, ) _v4c_ids = _v4c_ids.to(device) if _v4c_swap_ids is not None: _v4c_swap_ids = _v4c_swap_ids.to(device) _log_this_step = (_v4c_log_idx % 50 == 0) _v4d_mode = (getattr(raw_model, "salience_mode", "v0") == "v4d") if _v4d_mode and _log_this_step: raw_model.d_stage.store_v4d_attn = True _, _v4c_aux = model(_v4c_ids, v4c_spans=_v4c_items, v4c_swap_ids=_v4c_swap_ids, v4c_swap_spans=_v4c_swap_items) _v5_scale = _v5_aux_scale(step, cfg) if _v5_iso else 1.0 if _v5_iso: _v4c_aux = _v4c_aux * _v5_scale _v4c_aux.backward() _aux_grad_norm_sq = 0.0 for p in model.parameters(): if p.grad is None: continue g_lm = _lm_grad_snapshot.get(id(p)) if g_lm is None: _aux_grad_norm_sq += float((p.grad * p.grad).sum()) else: diff = p.grad - g_lm _aux_grad_norm_sq += float((diff * diff).sum()) _aux_gn_true = _aux_grad_norm_sq ** 0.5 if _v5_iso: for p in model.parameters(): if p.grad is None or id(p) in _sal_param_ids: continue snap = _lm_grad_snapshot.get(id(p)) if snap is None: p.grad = None elif _v5_trunk_frac > 0.0: p.grad = snap + _v5_trunk_frac * (p.grad - snap) else: p.grad = snap if _v4c_log_idx % 50 == 0: _ratio = (_aux_gn_true / _v4c_last_lm_grad_norm if _v4c_last_lm_grad_norm > 1e-6 else float('inf')) _msg = (f" [V4c step={step + 1:>5d}] " f"aux={_v4c_aux.item():.4f} " f"LM_grad={_v4c_last_lm_grad_norm:.3f} " f"V4c_grad={_aux_gn_true:.3f} " f"ratio={_ratio:.3f} " f"bce_α={raw_model.v4c_bce_alpha.item():.2f} " f"nce_α={raw_model.v4c_nce_alpha.item():.2f}") if _v5_iso: _msg += (f" [V5 iso every={_v4c_eff_every} " f"aux_scale={_v5_scale:.2f} " f"trunk_frac={_v5_trunk_frac:.2f}]") _mstats = getattr(raw_model.salience, "_last_v4c_margin_stats", None) if _mstats is not None: _msg += (f" margin={_mstats['loss']:.3f} " f"used={_mstats['n_used']} " f"skip={_mstats['n_skip_neutral']}/n," f"{_mstats['n_skip_eq_token']}/eq") _cstats = getattr(raw_model.salience, "_last_v4d_causal_stats", None) if _cstats is not None: _msg += (f" causal={_cstats['loss']:.3f} " f"gap_mean={_cstats['margin_gap_mean']:+.3f} " f"m_n={_cstats['margin_normal_mean']:+.2f} " f"m_c={_cstats['margin_clamped_mean']:+.2f}") _czstats = getattr(raw_model.salience, "_last_v4d_causal_z_stats", None) if _czstats is not None: _msg += (f" causalZ={_czstats['loss']:.4f} " f"mz_gap={_czstats['mz_gap_mean']:+.4f} " f"mz_n={_czstats['mz_normal_mean']:+.3f} " f"mz_c={_czstats['mz_clamped_mean']:+.3f}") _swstats = getattr(raw_model.salience, "_last_v4d_swap_stats", None) if _swstats is not None: _msg += (f" swap={_swstats['loss']:.4f}" f"(n={_swstats['n_pairs']})") _rlstats = getattr(raw_model.salience, "_last_v4d_role_stats", None) if _rlstats is not None: _msg += (f" role={_rlstats['loss']:.4f}" f"(acc={_rlstats['acc']:.2f})") _v6stats = getattr(raw_model.salience, "_last_v6_aux_stats", None) if _v6stats is not None: if bool(_v6stats.get("bce_active", False)): _msg += (f" v6={_v6stats['loss']:.3f}" f"(B={_v6stats['L_bank']:.3f}," f"P={_v6stats.get('L_bank_pair_ce', 0.0):.3f}," f"M={_v6stats['L_mix']:.3f}," f"Sel={_v6stats['L_sel']:.3f}" f"[ans={_v6stats['L_sel_ans']:.3f}," f"bg={_v6stats['L_sel_bg']:.3f}]) " f"lam_ans={_v6stats['lam_ans_mean']:.3f}" f"±{_v6stats['lam_ans_std']:.3f}" f"(n={_v6stats['n_items']})") else: _msg += (f" v6={_v6stats['loss']:.3f}" f"(B={_v6stats['L_bank']:.3f}," f"P={_v6stats.get('L_bank_pair_ce', 0.0):.3f}," f"M={_v6stats['L_mix']:.3f}," f"G={_v6stats['L_gate']:.3f}" f",Bg={_v6stats.get('L_bg', 0.0):.3f}) " f"lam_ans={_v6stats['lam_ans_mean']:.3f}" f"±{_v6stats['lam_ans_std']:.3f}" f"(sat_lo={_v6stats['lam_ans_sat_low']:.2f}," f"hi={_v6stats['lam_ans_sat_high']:.2f}," f"n={_v6stats['n_items']})") if int(_v6stats.get("n_bg_items", 0)) > 0: _msg += (f" lam_bg={_v6stats['lam_bg_mean']:.3f}" f"±{_v6stats['lam_bg_std']:.3f}" f"(n={_v6stats['n_bg_items']})") if "v6_lambda_head_bias" in _v6stats: _msg += (f" head_bias=" f"{_v6stats['v6_lambda_head_bias']:+.3f}") if _v4d_mode: _xattn = raw_model.d_stage.cross_attn _attn_t = getattr(_xattn, "_last_attn", None) if _attn_t is not None: with torch.no_grad(): _ent = -(_attn_t.clamp(min=1e-12) * _attn_t.clamp(min=1e-12).log() ).sum(-1).mean().item() _msg += (f" attn_ent={_ent:.2f} " f"sal_std={_xattn._last_sal_logits_std:.3f} " f"base_std={_xattn._last_base_logits_std:.3f}") if _ent < 1.5: print(_msg) raise RuntimeError( f"V4d attention saturation: mean attn entropy " f"{_ent:.2f} < 1.5 nats at step {step + 1}. " f"Attention collapsed onto a single slot. " f"Lower v4d_logit_cap (4.0 → 2.0) or zero-init W_sal.") print(_msg) if _v4d_mode and _log_this_step: raw_model.d_stage.store_v4d_attn = False _v4c_log_idx += 1 del _lm_grad_snapshot except Exception as _e: raise RuntimeError( f"V4c aux forward failed at step {step + 1} " f"(tag={tag!r}, variant={getattr(raw_model, 'v4c_variant', '?')}): " f"{_e}") from _e torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) lr = get_lr(step, cfg) for i, pg in enumerate(optimizer.param_groups): if pg.get("bank_head_lr_constant", False): continue multiplier = 1.0 if (not is_cedl or i == 0) else 1.5 pg['lr'] = lr * multiplier if xm is not None: xm.optimizer_step(optimizer) else: optimizer.step() step += 1 if is_cedl: if step < cedl_cfg.feedback_warmup_start: raw_model.feedback_alpha.fill_(0.0) elif step < cedl_cfg.feedback_warmup_end: frac = (step - cedl_cfg.feedback_warmup_start) / ( cedl_cfg.feedback_warmup_end - cedl_cfg.feedback_warmup_start) raw_model.feedback_alpha.fill_(frac) else: raw_model.feedback_alpha.fill_(1.0) if getattr(raw_model, 'use_salience', False): if step < 1000: raw_model.sl_alpha.fill_(0.0) elif step < 2000: raw_model.sl_alpha.fill_((step - 1000) / 1000.0) else: raw_model.sl_alpha.fill_(1.0) if getattr(raw_model, "salience_mode", "v0") in ( "v4c", "v4c_gate", "v4d"): if step < 500: raw_model.v4c_bce_alpha.fill_(0.0) elif step < 1500: raw_model.v4c_bce_alpha.fill_((step - 500) / 1000.0) else: raw_model.v4c_bce_alpha.fill_(1.0) if step < 1500: raw_model.v4c_nce_alpha.fill_(0.0) elif step < 3000: raw_model.v4c_nce_alpha.fill_((step - 1500) / 1500.0) else: raw_model.v4c_nce_alpha.fill_(1.0) anneal_step = int(cfg.max_steps * cedl_cfg.sparsity_anneal_frac) if step >= anneal_step and cfg.max_steps > anneal_step: frac = (step - anneal_step) / (cfg.max_steps - anneal_step) sparsity = cedl_cfg.e_sparsity + frac * ( cedl_cfg.sparsity_final - cedl_cfg.e_sparsity) raw_model.e_stage.set_sparsity(sparsity) if step % 100 == 0: elapsed = time.time() - t0 tokens_per_sec = (step * cfg.batch_size * cfg.grad_accum * cfg.max_seq) / elapsed print(f" [{step:>6d}/{cfg.max_steps}] loss={accum_loss/100:.4f} " f"lr={lr:.2e} tok/s={tokens_per_sec:.0f}") accum_loss = 0.0 if step % cfg.eval_interval == 0: val_ppl = evaluate(model, val_loader, device, cfg.eval_steps, tpu=cfg.tpu) print(f" >>> Val PPL: {val_ppl:.2f}") _is_best = val_ppl < best_val_ppl if _is_best: best_val_ppl = val_ppl best_step = step ckpt_path = os.path.join(cfg.save_dir, f"{tag}_best.pt") if xm is not None: xm.save(model.state_dict(), ckpt_path) else: torch.save(model.state_dict(), ckpt_path) print(f" >>> Saved best checkpoint: {ckpt_path}") if is_cedl and getattr(raw_model, "salience_mode", "v0") in ( "v4c", "v4c_gate", "v4d"): import json as _json cfg_path = os.path.join(cfg.save_dir, f"{tag}_config.json") with open(cfg_path, "w") as _f: _json.dump({ "tag": tag, "salience_mode": raw_model.salience_mode, "v4c_variant": getattr( raw_model, "v4c_variant", "base"), "v4c_margin_weight": float(getattr( raw_model._v4c_cfg, "v4c_margin_weight", 0.0)), "v4c_margin_target": float(getattr( raw_model._v4c_cfg, "v4c_margin_target", 1.0)), "v4c_warm_start_sigma": float(getattr( raw_model._v4c_cfg, "v4c_warm_start_sigma", 0.05)), "v4c_norm_cap": float(getattr( raw_model._v4c_cfg, "v4c_norm_cap", 0.3)), "v4d_w_sal_sigma": float(getattr( raw_model._v4c_cfg, "v4d_w_sal_sigma", 0.02)), "v4d_logit_cap": float(getattr( raw_model._v4c_cfg, "v4d_logit_cap", 4.0)), "v4d_noinject": bool(getattr( raw_model._v4c_cfg, "v4d_noinject", False)), "v4d_causal_weight": float(getattr( raw_model._v4c_cfg, "v4d_causal_weight", 0.0)), "v4d_causal_gap": float(getattr( raw_model._v4c_cfg, "v4d_causal_gap", 0.25)), "v4d_causal_z_weight": float(getattr( raw_model._v4c_cfg, "v4d_causal_z_weight", 0.0)), "v4d_causal_z_gap": float(getattr( raw_model._v4c_cfg, "v4d_causal_z_gap", 0.15)), "v4d_w_sal_rank": int(getattr( raw_model._v4c_cfg, "v4d_w_sal_rank", 0)), "v4d_swap_consistency_weight": float(getattr( raw_model._v4c_cfg, "v4d_swap_consistency_weight", 0.0)), "v4d_role_sep_weight": float(getattr( raw_model._v4c_cfg, "v4d_role_sep_weight", 0.0)), "v6_mixture": bool(getattr( raw_model, "v6_mixture", False)), "v6_lambda_init": float(getattr( raw_model, "v6_lambda_init", -4.0)), "v6_lambda_a_init": float(getattr( raw_model, "v6_lambda_a_init", 1.0)), "v6_lambda_head": bool(getattr( raw_model, "use_v6_lambda_head", False)), "v6_lambda_head_hidden": int(getattr( raw_model._v4c_cfg, "v6_lambda_head_hidden", 160)), "v6_lambda_head_bias_init": float(getattr( raw_model._v4c_cfg, "v6_lambda_head_bias_init", -7.0)), "v6_bg_weight": float(getattr( raw_model._v4c_cfg, "v6_bg_weight", 1.0)), "v6_bg_target": float(getattr( raw_model._v4c_cfg, "v6_bg_target", 0.01)), "v6_bce_objective": bool(getattr( raw_model._v4c_cfg, "v6_bce_objective", False)), "v6_sel_weight": float(getattr( raw_model._v4c_cfg, "v6_sel_weight", 1.0)), "v6_lambda_head_w_init_std": float(getattr( raw_model._v4c_cfg, "v6_lambda_head_w_init_std", 1e-3)), "v6_wt_sparsity_weight": float(getattr( raw_model._v4c_cfg, "v6_wt_sparsity_weight", 0.0)), "v6_wt_sparsity_target": float(getattr( raw_model._v4c_cfg, "v6_wt_sparsity_target", 0.0)), "v6_mem_head_bank": bool(getattr( raw_model._v4c_cfg, "v6_mem_head_bank", False)), "v6_bank_ce_weight": float(getattr( raw_model._v4c_cfg, "v6_bank_ce_weight", 0.0)), "v6_bank_pair_ce_weight": float(getattr( raw_model._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)), "v6_bank_query_source": str(getattr( raw_model._v4c_cfg, "v6_bank_query_source", "h_d")), "v6_bank_readout_mode": str(getattr( raw_model._v4c_cfg, "v6_bank_readout_mode", "bank")), "v6_source_adapter": bool(getattr( raw_model._v4c_cfg, "v6_source_adapter", False)), "v6_context_adapter": bool(getattr( raw_model._v4c_cfg, "v6_context_adapter", False)), "v6_bank_head_lr": float(getattr( cfg, "v6_bank_head_lr", 0.0)), "v6_specialist_from_b0": bool(getattr( cfg, "v6_specialist_from_b0", False)), "v6_specialist_freeze": str(getattr( cfg, "v6_specialist_freeze", "none")), "v6_specialist_noinject": bool(getattr( cfg, "v6_specialist_noinject", False)), "v6_specialist_base_checkpoint": getattr( raw_model, "_v6_specialist_base_checkpoint", getattr(cfg, "resume_checkpoint", None)), "v6_specialist_trainable_prefixes": getattr( raw_model, "_v6_specialist_trainable_prefixes", None), "v6_specialist_loaded_missing": getattr( raw_model, "_v6_specialist_loaded_missing", None), "v6_lambda_head_bias_live": ( float(raw_model.v6_lambda_head[2].bias.item()) if getattr(raw_model, "use_v6_lambda_head", False) else None), "v6_lambda_trajectory": getattr( raw_model, "_last_eval_lam_avg", None), "v6_aux_weight": float(getattr( raw_model._v4c_cfg, "v6_aux_weight", 0.0)), "v6_margin_target": float(getattr( raw_model._v4c_cfg, "v6_margin_target", 1.0)), "v6_mix_weight": float(getattr( raw_model._v4c_cfg, "v6_mix_weight", 0.5)), "v6_gate_weight": float(getattr( raw_model._v4c_cfg, "v6_gate_weight", 0.01)), "v6_lambda_floor": float(getattr( raw_model._v4c_cfg, "v6_lambda_floor", 0.05)), "v6_aux_trajectory": getattr( getattr(raw_model, "salience", None), "_last_v6_aux_stats", None), "v5_training": { "v5_grad_isolation": bool(getattr(cfg, "v5_grad_isolation", False)), "v5_trunk_aux_frac": float(getattr(cfg, "v5_trunk_aux_frac", 0.0)), "resume_checkpoint": getattr(cfg, "resume_checkpoint", None), "resume_start_step": int(getattr(cfg, "resume_start_step", 0) or 0), "schedule_total_steps": int(cfg.max_steps), "aux_scale_schedule": ( f"5-10k:{getattr(cfg,'v5_aux_scale_early',0.50)}, " f"10-20k:{getattr(cfg,'v5_aux_scale_mid',0.25)}, " f"20k+:{getattr(cfg,'v5_aux_scale_late',0.10)}"), "v4c_every_schedule": ( f"5-10k:{getattr(cfg,'v5_cadence_early',8)}, " f"10-20k:{getattr(cfg,'v5_cadence_mid',16)}, " f"20k+:{getattr(cfg,'v5_cadence_late',32)}"), }, "step": step, "val_ppl": float(val_ppl), }, _f, indent=2) if step % cfg.eval_interval == 0: try: import json as _sj _sc_path = os.path.join(cfg.save_dir, f"{tag}_trainscore.jsonl") with open(_sc_path, "a") as _scf: _row = { "step": step, "global_step": step, "model_tag": tag, "ppl": float(val_ppl), "is_best": bool(_is_best), "best_ppl": float(best_val_ppl), "val_loss": float(math.log(max(val_ppl, 1e-9))), "checkpoint_path": os.path.join(cfg.save_dir, f"{tag}_best.pt"), "run_started_at": _run_started_at, "run_start_step": _run_start_step, } _lam_avg = getattr(raw_model, "_last_eval_lam_avg", None) if _lam_avg is not None: _row["v6_lambda"] = _lam_avg _sal = getattr(raw_model, "salience", None) _v6_aux = getattr(_sal, "_last_v6_aux_stats", None) if _sal else None if _v6_aux is not None: _row["v6_aux"] = _v6_aux _scf.write(_sj.dumps(_row) + "\n") except Exception as _e: print(f" [scorecard] WARN: could not append train-score ({_e})") if getattr(cfg, "save_trainstate", False) and xm is None: _save_trainstate(step) if step % cfg.save_interval == 0: ckpt_path = os.path.join(cfg.save_dir, f"{tag}_step{step}.pt") if xm is not None: xm.save(model.state_dict(), ckpt_path) else: torch.save(model.state_dict(), ckpt_path) if getattr(cfg, "save_trainstate", False) and xm is None: _save_trainstate(step) if _v3m2_diagnostic_active and step > 0 and step % 250 == 0: _run_vvec_pathway_diagnostic( raw_model, _v3m2_eval_batch, step, device) final_path = os.path.join(cfg.save_dir, f"{tag}_final.pt") if xm is not None: xm.save(model.state_dict(), final_path) else: torch.save(model.state_dict(), final_path) print(f" >>> Saved final checkpoint: {final_path}") if getattr(cfg, "save_trainstate", False) and xm is None: _save_trainstate(step) if is_cedl and getattr(raw_model, "salience_mode", "v0") in ( "v4c", "v4c_gate", "v4d"): import json as _json cfg_path = os.path.join(cfg.save_dir, f"{tag}_config.json") if not os.path.exists(cfg_path): with open(cfg_path, "w") as _f: _json.dump({ "tag": tag, "salience_mode": raw_model.salience_mode, "v4c_variant": getattr(raw_model, "v4c_variant", "base"), "v4c_margin_weight": float(getattr( raw_model._v4c_cfg, "v4c_margin_weight", 0.0)), "v4c_margin_target": float(getattr( raw_model._v4c_cfg, "v4c_margin_target", 1.0)), "v4c_warm_start_sigma": float(getattr( raw_model._v4c_cfg, "v4c_warm_start_sigma", 0.05)), "v4c_norm_cap": float(getattr( raw_model._v4c_cfg, "v4c_norm_cap", 0.3)), "v4d_w_sal_sigma": float(getattr( raw_model._v4c_cfg, "v4d_w_sal_sigma", 0.02)), "v4d_logit_cap": float(getattr( raw_model._v4c_cfg, "v4d_logit_cap", 4.0)), "v4d_noinject": bool(getattr( raw_model._v4c_cfg, "v4d_noinject", False)), "v4d_causal_weight": float(getattr( raw_model._v4c_cfg, "v4d_causal_weight", 0.0)), "v4d_causal_gap": float(getattr( raw_model._v4c_cfg, "v4d_causal_gap", 0.25)), "v4d_causal_z_weight": float(getattr( raw_model._v4c_cfg, "v4d_causal_z_weight", 0.0)), "v4d_causal_z_gap": float(getattr( raw_model._v4c_cfg, "v4d_causal_z_gap", 0.15)), "v4d_w_sal_rank": int(getattr( raw_model._v4c_cfg, "v4d_w_sal_rank", 0)), "v4d_swap_consistency_weight": float(getattr( raw_model._v4c_cfg, "v4d_swap_consistency_weight", 0.0)), "v4d_role_sep_weight": float(getattr( raw_model._v4c_cfg, "v4d_role_sep_weight", 0.0)), "v6_mixture": bool(getattr( raw_model, "v6_mixture", False)), "v6_lambda_init": float(getattr( raw_model, "v6_lambda_init", -4.0)), "v6_lambda_a_init": float(getattr( raw_model, "v6_lambda_a_init", 1.0)), "v6_lambda_head": bool(getattr( raw_model, "use_v6_lambda_head", False)), "v6_lambda_head_hidden": int(getattr( raw_model._v4c_cfg, "v6_lambda_head_hidden", 160)), "v6_lambda_head_bias_init": float(getattr( raw_model._v4c_cfg, "v6_lambda_head_bias_init", -7.0)), "v6_bg_weight": float(getattr( raw_model._v4c_cfg, "v6_bg_weight", 1.0)), "v6_bg_target": float(getattr( raw_model._v4c_cfg, "v6_bg_target", 0.01)), "v6_bce_objective": bool(getattr( raw_model._v4c_cfg, "v6_bce_objective", False)), "v6_sel_weight": float(getattr( raw_model._v4c_cfg, "v6_sel_weight", 1.0)), "v6_lambda_head_w_init_std": float(getattr( raw_model._v4c_cfg, "v6_lambda_head_w_init_std", 1e-3)), "v6_wt_sparsity_weight": float(getattr( raw_model._v4c_cfg, "v6_wt_sparsity_weight", 0.0)), "v6_wt_sparsity_target": float(getattr( raw_model._v4c_cfg, "v6_wt_sparsity_target", 0.0)), "v6_mem_head_bank": bool(getattr( raw_model._v4c_cfg, "v6_mem_head_bank", False)), "v6_bank_ce_weight": float(getattr( raw_model._v4c_cfg, "v6_bank_ce_weight", 0.0)), "v6_bank_pair_ce_weight": float(getattr( raw_model._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)), "v6_bank_query_source": str(getattr( raw_model._v4c_cfg, "v6_bank_query_source", "h_d")), "v6_bank_readout_mode": str(getattr( raw_model._v4c_cfg, "v6_bank_readout_mode", "bank")), "v6_source_adapter": bool(getattr( raw_model._v4c_cfg, "v6_source_adapter", False)), "v6_context_adapter": bool(getattr( raw_model._v4c_cfg, "v6_context_adapter", False)), "v6_bank_head_lr": float(getattr( cfg, "v6_bank_head_lr", 0.0)), "v6_specialist_from_b0": bool(getattr( cfg, "v6_specialist_from_b0", False)), "v6_specialist_freeze": str(getattr( cfg, "v6_specialist_freeze", "none")), "v6_specialist_noinject": bool(getattr( cfg, "v6_specialist_noinject", False)), "v6_specialist_base_checkpoint": getattr( raw_model, "_v6_specialist_base_checkpoint", getattr(cfg, "resume_checkpoint", None)), "v6_specialist_trainable_prefixes": getattr( raw_model, "_v6_specialist_trainable_prefixes", None), "v6_specialist_loaded_missing": getattr( raw_model, "_v6_specialist_loaded_missing", None), "v6_lambda_head_bias_live": ( float(raw_model.v6_lambda_head[2].bias.item()) if getattr(raw_model, "use_v6_lambda_head", False) else None), "v6_lambda_trajectory": getattr( raw_model, "_last_eval_lam_avg", None), "v6_aux_weight": float(getattr( raw_model._v4c_cfg, "v6_aux_weight", 0.0)), "v6_margin_target": float(getattr( raw_model._v4c_cfg, "v6_margin_target", 1.0)), "v6_mix_weight": float(getattr( raw_model._v4c_cfg, "v6_mix_weight", 0.5)), "v6_gate_weight": float(getattr( raw_model._v4c_cfg, "v6_gate_weight", 0.01)), "v6_lambda_floor": float(getattr( raw_model._v4c_cfg, "v6_lambda_floor", 0.05)), "v6_aux_trajectory": getattr( getattr(raw_model, "salience", None), "_last_v6_aux_stats", None), "step": step, "val_ppl": float(best_val_ppl) if best_val_ppl != float('inf') else None, }, _f, indent=2) return best_val_ppl @torch.no_grad() def _run_vvec_pathway_diagnostic(raw_model, eval_batch, step, device): """V3m2 mid-training diagnostic. Runs the FIXED eval batch through the model twice (once normal, once with v_vec clamped to zero via a hook), measures the logit Effective-Contribution-Ratio (EVR), and prints salience_bias_up.norm. Switches to eval mode for the entire diagnostic so EMA / density / sl_alpha-gated branches don't fire.""" was_training = raw_model.training raw_model.eval() try: sb_up_norm = raw_model.d_stage.salience_bias_up.weight.norm().item() out_normal = raw_model(eval_batch) logits_normal = out_normal[0] if isinstance(out_normal, tuple) else out_normal def _zero_v_vec_hook(_module, _inputs, output): v_vec, v_scalar = output return torch.zeros_like(v_vec), v_scalar handle = raw_model.salience.register_forward_hook(_zero_v_vec_hook) try: out_clamped = raw_model(eval_batch) logits_clamped = (out_clamped[0] if isinstance(out_clamped, tuple) else out_clamped) finally: handle.remove() delta = (logits_normal - logits_clamped).norm().item() denom = logits_normal.norm().item() + 1e-9 evr = delta / denom print(f" [v_vec_pathway step={step:>5d}] " f"sb_up.norm={sb_up_norm:.4f} EVR={evr:.4f}") if step >= 3000 and evr < 0.05: print(f" [v_vec_pathway WARN] EVR<0.05 at step {step}: " f"pathway may be too weak; calibrate against V0 baseline " f"(load V0 ckpt + run diagnostic) before declaring failure.") finally: raw_model.train(was_training) def run_structured_benchmark(model, tokenizer, device, max_seq=1024, tpu=False): """Zero-shot structured benchmark: 12 tasks x 4 difficulty levels.""" import random as _rng _rng.seed(42) xm_mod = None if tpu: try: import torch_xla.core.xla_model as _xm xm_mod = _xm except ImportError: pass N_INSTANCES = 200 NAMES = ["Alice", "Bob", "Carol", "Dave", "Eve", "Frank", "Grace", "Hank", "Iris", "Jack", "Kate", "Leo", "Mia", "Nick", "Olivia", "Paul", "Quinn", "Rosa", "Sam", "Tina"] COLORS = ["red", "blue", "green", "pink", "gray", "brown", "white", "black", "gold", "silver", "orange", "purple", "yellow"] OBJECTS = ["car", "hat", "bag", "pen", "cup", "box", "ring", "lamp", "book", "ball", "shoe", "coat", "desk", "bell", "fork", "drum", "fish", "kite", "coin", "vase"] TRAITS = ["tall", "short", "fast", "slow", "kind", "bold", "calm", "warm", "cold", "rich", "poor", "wise", "young", "old", "brave"] ANIMALS = ["dog", "cat", "bird", "fish", "frog", "bear", "deer", "wolf", "duck", "fox", "pig", "cow", "ant", "bee", "rat"] FRUITS = ["apple", "banana", "grape", "lemon", "mango", "peach", "plum", "cherry", "melon", "berry"] def gen_A1(level): n = [2, 4, 6, 8][level] names = _rng.sample(NAMES, n) colors = [_rng.choice(COLORS) for _ in range(n)] objs = [_rng.choice(OBJECTS) for _ in range(n)] facts = ". ".join(f"{names[i]} has a {colors[i]} {objs[i]}" for i in range(n)) qi = _rng.randint(0, n - 1) return f"{facts}. Who has a {colors[qi]} {objs[qi]}? Answer:", f" {names[qi]}" def gen_A2(level): n = [2, 4, 6, 8][level] names = _rng.sample(NAMES, n) traits = [_rng.choice(TRAITS) for _ in range(n)] facts = ". ".join(f"{names[i]} is {traits[i]}" for i in range(n)) qi = _rng.randint(0, n - 1) is_true = _rng.random() < 0.5 ct = traits[qi] if is_true else _rng.choice([t for t in TRAITS if t != traits[qi]]) return f"{facts}. Claim: {names[qi]} is {ct}. True or false? Answer:", " true" if is_true else " false" def gen_A3(level): n_hops = [1, 2, 3, 3][level] n_dist = [0, 0, 0, 3][level] all_n = _rng.sample(NAMES, n_hops + 1 + n_dist * 2) chain = all_n[:n_hops + 1] rels = [f"{chain[i]} is parent of {chain[i+1]}" for i in range(n_hops)] for i in range(n_dist): rels.append(f"{all_n[n_hops+1+2*i]} is friend of {all_n[n_hops+2+2*i]}") _rng.shuffle(rels) label = {1: "child", 2: "grandchild", 3: "great grandchild"} return f"{'. '.join(rels)}. Who is {chain[0]}'s {label[n_hops]}? Answer:", f" {chain[-1]}" def gen_A4(level): n_pairs, noise = 4, [0, 2, 4, 8][level] keys = _rng.sample(FRUITS, n_pairs) vals = _rng.sample(COLORS, n_pairs) parts = [] for i in range(n_pairs): parts.append(f"key: {keys[i]} value: {vals[i]}.") if noise > 0: parts.append(" ".join(_rng.choice(ANIMALS) for _ in range(noise)) + ".") qi = _rng.randint(0, n_pairs - 1) return f"{' '.join(parts)} What is the value of {keys[qi]}? Answer:", f" {vals[qi]}" def gen_B1(level): if level == 0: a, b = _rng.randint(1, 9), _rng.randint(1, 9); r = a + b return f"What is {a} + {b}? Answer:", f" {r}" elif level == 1: a, b = _rng.randint(10, 49), _rng.randint(10, 49); r = a + b return f"What is {a} + {b}? Answer:", f" {r}" elif level == 2: a, b = _rng.randint(1, 30), _rng.randint(1, 30) op = _rng.choice(["+", "-"]) if op == "-": a, b = max(a, b), min(a, b) r = (a + b) if op == "+" else (a - b) return f"What is {a} {op} {b}? Answer:", f" {r}" else: a, b, c = _rng.randint(1, 20), _rng.randint(1, 20), _rng.randint(1, 15) s1 = a + b; op2 = _rng.choice(["+", "-"]) if op2 == "-" and s1 < c: c = _rng.randint(1, s1) r = (s1 + c) if op2 == "+" else (s1 - c) return f"What is {a} + {b} {op2} {c}? Answer:", f" {r}" def gen_B2(level): if level == 0: a, b = _rng.randint(1, 30), _rng.randint(1, 30) return f"If x = {a} + {b}, what is x? Answer:", f" {a+b}" elif level == 1: a, b, c = _rng.randint(1, 15), _rng.randint(1, 15), _rng.randint(1, 15) return f"If x = {a} + {b}, and y = x + {c}, what is y? Answer:", f" {a+b+c}" elif level == 2: a, b, c = _rng.randint(2, 9), _rng.randint(2, 9), _rng.randint(1, 15) return f"If x = {a} * {b} + {c}, what is x? Answer:", f" {a*b+c}" else: a, b, c, d = _rng.randint(1, 12), _rng.randint(1, 12), _rng.randint(1, 12), _rng.randint(1, 12) return f"If x = {a} + {b}, and y = {c} + {d}, what is x + y? Answer:", f" {a+b+c+d}" def gen_B3(level): if level == 0: a, b = _rng.randint(1, 99), _rng.randint(1, 99) while a == b: b = _rng.randint(1, 99) return f"Which is larger, {a} or {b}? Answer:", f" {max(a, b)}" elif level == 1: vals = _rng.sample(range(1, 99), 3) return f"Which is largest, {vals[0]}, {vals[1]}, or {vals[2]}? Answer:", f" {max(vals)}" elif level == 2: a, b, c, d = _rng.randint(1, 40), _rng.randint(1, 40), _rng.randint(1, 40), _rng.randint(1, 40) while a + b == c + d: d = _rng.randint(1, 40) return f"Which is larger, {a} + {b} or {c} + {d}? Answer:", f" {max(a+b, c+d)}" else: b = _rng.choice([3, 4, 5, 6, 7, 8, 9]); a = _rng.randint(10, 99) return f"What is {a} mod {b}? Answer:", f" {a % b}" def gen_B4(level): if level <= 2: target = _rng.choice(ANIMALS[:5]) others = [x for x in ANIMALS[:8] if x != target] seq_len = [_rng.randint(4, 8), _rng.randint(6, 10), _rng.randint(8, 15)][level] seq = [_rng.choice([target] + others) for _ in range(seq_len)] r = sum(1 for s in seq if s == target) return f"Count the number of times '{target}' appears: {' '.join(seq)}. Answer:", f" {r}" else: t1, t2 = _rng.sample(ANIMALS[:6], 2) others = [x for x in ANIMALS[:10] if x not in (t1, t2)] seq = [_rng.choice([t1, t2] + others) for _ in range(_rng.randint(6, 12))] r = sum(1 for s in seq if s in (t1, t2)) return f"Count the total times '{t1}' or '{t2}' appear: {' '.join(seq)}. Answer:", f" {r}" def gen_A5(level): """Memory update: present facts, then update some, query latest value.""" n_facts, n_updates = [(2, 1), (4, 2), (6, 3), (8, 4)][level] names = _rng.sample(NAMES, n_facts) attrs = _rng.sample(OBJECTS, n_facts) facts = ". ".join(f"{names[i]} owns a {attrs[i]}" for i in range(n_facts)) update_idx = set(_rng.sample(range(n_facts), n_updates)) new_attrs = {} for ui in update_idx: remaining = [o for o in OBJECTS if o != attrs[ui]] new_attrs[ui] = _rng.choice(remaining) updates = ". ".join(f"{names[ui]} now owns a {new_attrs[ui]}" for ui in sorted(update_idx)) query_updated = _rng.random() < 0.5 updated_list = list(update_idx) unchanged_list = [i for i in range(n_facts) if i not in update_idx] if query_updated and updated_list: qi = _rng.choice(updated_list) elif unchanged_list: qi = _rng.choice(unchanged_list) else: qi = _rng.choice(updated_list) answer = new_attrs[qi] if qi in update_idx else attrs[qi] return f"{facts}. Update: {updates}. What does {names[qi]} own now? Answer:", f" {answer}" def gen_C1(level): """Periodic pattern detection: predict next element in a repeating sequence. Probes C-stage grid-cell periodic retention.""" period, noise_per, length = [ (2, 0, 8), (3, 0, 15), (4, 2, 20), (6, 3, 30)][level] pattern_items = _rng.sample(OBJECTS[:10], period) sequence = [] for i in range(length): sequence.append(pattern_items[i % period]) if noise_per > 0 and _rng.random() < 0.3: for _ in range(min(noise_per, 2)): sequence.append(_rng.choice(ANIMALS[:6])) next_item = pattern_items[len([s for s in sequence if s in pattern_items]) % period] prompt = " ".join(sequence) return f"Repeating pattern: {prompt}. Next item:", f" {next_item}" def gen_C2(level): """Near-miss disambiguation: distinguish entities with overlapping attributes. Probes E-stage AHSD/CSR pattern separation.""" n_entities, n_shared = [(2, 1), (2, 2), (3, 2), (4, 3)][level] names = _rng.sample(NAMES, n_entities) shared_colors = _rng.sample(COLORS, n_shared) shared_objs = _rng.sample(OBJECTS[:10], n_shared) unique_objs = _rng.sample(OBJECTS[10:], n_entities) facts = [] for i, name in enumerate(names): for j in range(n_shared): facts.append(f"{name} has a {shared_colors[j]} {shared_objs[j]}") facts.append(f"{name} has a {unique_objs[i]}") _rng.shuffle(facts) qi = _rng.randint(0, n_entities - 1) return (f"{'. '.join(facts)}. Who has a {unique_objs[qi]}? Answer:", f" {names[qi]}") def gen_C3(level): """Partial cue completion: complete a degraded fact from stored patterns. Probes D-stage attractor pattern completion.""" n_facts, n_missing = [(2, 1), (3, 1), (4, 2), (5, 2)][level] names = _rng.sample(NAMES, n_facts) colors = [_rng.choice(COLORS) for _ in range(n_facts)] objs = _rng.sample(OBJECTS, n_facts) traits = _rng.sample(TRAITS, n_facts) facts = [f"{names[i]} is {traits[i]} and has a {colors[i]} {objs[i]}" for i in range(n_facts)] context = ". ".join(facts) qi = _rng.randint(0, n_facts - 1) cue_parts = [] query_attr = None attrs = [("trait", traits[qi]), ("color", colors[qi]), ("object", objs[qi])] mask_indices = _rng.sample(range(len(attrs)), min(n_missing, len(attrs))) query_idx = _rng.choice(mask_indices) query_attr_name, query_attr_val = attrs[query_idx] cue = f"{names[qi]}" for j, (aname, aval) in enumerate(attrs): if j in mask_indices: cue += f" ??? {aname}" else: cue += f" {aval}" if query_attr_name == "color": return (f"{context}. Partial cue: {cue}. What is {names[qi]}'s color? Answer:", f" {query_attr_val}") elif query_attr_name == "object": return (f"{context}. Partial cue: {cue}. What does {names[qi]} have? Answer:", f" {query_attr_val}") else: return (f"{context}. Partial cue: {cue}. What trait does {names[qi]} have? Answer:", f" {query_attr_val}") TASKS = { "A1_fact_retrieval": gen_A1, "A2_consistency": gen_A2, "A3_multihop": gen_A3, "A4_noisy_retrieval": gen_A4, "A5_memory_update": gen_A5, "B1_arithmetic": gen_B1, "B2_algebra": gen_B2, "B3_comparison": gen_B3, "B4_counting": gen_B4, "C1_periodic_pattern": gen_C1, "C2_near_miss": gen_C2, "C3_partial_completion": gen_C3, } CANDIDATES = { "A1_fact_retrieval": NAMES, "A2_consistency": ["true", "false"], "A3_multihop": NAMES, "A4_noisy_retrieval": COLORS, "A5_memory_update": OBJECTS, "B1_arithmetic": [str(i) for i in range(-50, 200)], "B2_algebra": [str(i) for i in range(-50, 200)], "B3_comparison": [str(i) for i in range(0, 200)], "B4_counting": [str(i) for i in range(0, 30)], "C1_periodic_pattern": OBJECTS[:10], "C2_near_miss": NAMES, "C3_partial_completion": COLORS + OBJECTS + TRAITS, } cand_ids_cache = {} for task_name, cands in CANDIDATES.items(): cand_ids_cache[task_name] = [] for c in cands: ids = tokenizer.encode(f" {c}", add_special_tokens=False) cand_ids_cache[task_name].append((c, ids)) model.eval() if hasattr(model, 'reset_memory'): model.reset_memory() results = {} for task_name, gen_fn in TASKS.items(): results[task_name] = {} cand_list = cand_ids_cache[task_name] for level in range(4): correct = 0 for _ in range(N_INSTANCES): if hasattr(model, 'reset_memory'): model.reset_memory() prompt_str, answer_str = gen_fn(level) answer_str_stripped = answer_str.strip() prompt_ids = tokenizer.encode(prompt_str) if len(prompt_ids) >= max_seq: prompt_ids = prompt_ids[-(max_seq - 1):] inp = torch.tensor([prompt_ids], dtype=torch.long, device=device) with torch.no_grad(): logits = model(inp) if isinstance(logits, tuple): logits = logits[0] log_probs = F.log_softmax(logits[0, -1], dim=-1) del logits best_score = float('-inf') best_cand = None for cand_str, cand_tok_ids in cand_list: if not cand_tok_ids: continue score = log_probs[cand_tok_ids[0]].item() if score > best_score: best_score = score best_cand = cand_str if best_cand == answer_str_stripped: correct += 1 if xm_mod is not None: xm_mod.mark_step() results[task_name][level] = correct / N_INSTANCES print(f"\n{'='*72}") print("Structured Benchmark Results (accuracy per task per level)") print(f"{'='*72}") header = f"{'Task':<22s}" + "".join(f" {'L'+str(i):>7s}" for i in range(4)) + f" {'Avg':>7s}" print(header) print("-" * len(header)) a_accs, b_accs, c_accs = [], [], [] for tn in TASKS: row = f"{tn:<22s}" avgs = [] for lv in range(4): acc = results[tn][lv] row += f" {acc:>6.1%}"; avgs.append(acc) row += f" {sum(avgs)/4:>6.1%}" print(row) if tn.startswith("A"): a_accs.extend(avgs) elif tn.startswith("B"): b_accs.extend(avgs) else: c_accs.extend(avgs) print("-" * len(header)) print(f"{'Content (A1-A5)':<22s}{' ':>31s} {sum(a_accs)/len(a_accs):>6.1%}") print(f"{'Math (B1-B4)':<22s}{' ':>31s} {sum(b_accs)/len(b_accs):>6.1%}") print(f"{'CEDL-probe (C1-C3)':<22s}{' ':>31s} {sum(c_accs)/len(c_accs):>6.1%}") all_accs = a_accs + b_accs + c_accs print(f"{'Overall':<22s}{' ':>31s} {sum(all_accs)/len(all_accs):>6.1%}") print(f"{'='*72}\n") return results def run_downstream_eval(model: nn.Module, tag: str, device: torch.device): """Zero-shot evaluation on LAMBADA, HellaSwag, ARC-Easy.""" try: from lm_eval import evaluator from lm_eval.api.model import LM from transformers import GPT2TokenizerFast except ImportError: print(" lm-eval not installed. Skipping downstream eval.") print(" Install: pip install lm-eval transformers") return {} tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") class WrappedModel(LM): def __init__(self, model, tokenizer, device): super().__init__() self._model = model self._tokenizer = tokenizer self._device = device @property def eot_token_id(self): return self._tokenizer.eos_token_id @property def max_length(self): return 1024 @property def max_gen_toks(self): return 256 @property def batch_size(self): return 8 @property def device(self): return self._device def tok_encode(self, string): return self._tokenizer.encode(string, add_special_tokens=False) def tok_decode(self, tokens): return self._tokenizer.decode(tokens) def loglikelihood(self, requests): results = [] for req in requests: if hasattr(req, 'args'): ctx, cont = req.args else: ctx, cont = req ctx_ids = self.tok_encode(ctx) cont_ids = self.tok_encode(cont) all_ids = ctx_ids + cont_ids if len(all_ids) > 1024: excess = len(all_ids) - 1024 all_ids = all_ids[excess:] ctx_ids = ctx_ids[excess:] input_ids = torch.tensor([all_ids], device=self._device) with torch.no_grad(): logits = self._model(input_ids) if isinstance(logits, tuple): logits = logits[0] start = len(ctx_ids) log_probs = F.log_softmax(logits[0], dim=-1) cont_log_prob = 0.0 is_greedy = True for j in range(start, len(all_ids)): tok = all_ids[j] lp = log_probs[j - 1, tok].item() cont_log_prob += lp if logits[0, j - 1].argmax().item() != tok: is_greedy = False results.append((cont_log_prob, is_greedy)) return results def loglikelihood_rolling(self, requests): results = [] for req in requests: if hasattr(req, 'args'): (text,) = req.args else: text = req[0] if isinstance(req, tuple) else req all_ids = self.tok_encode(text) max_len = 1024 stride = 512 total_lp = 0.0 scored = set() for start in range(0, max(1, len(all_ids) - 1), stride): end = min(start + max_len, len(all_ids)) chunk = all_ids[start:end] if len(chunk) < 2: continue input_t = torch.tensor([chunk], device=self._device) with torch.no_grad(): logits = self._model(input_t) if isinstance(logits, tuple): logits = logits[0] log_probs = F.log_softmax(logits[0], dim=-1) score_start = 1 if start == 0 else (end - start - stride) for j in range(max(1, score_start), len(chunk)): abs_pos = start + j if abs_pos not in scored: total_lp += log_probs[j - 1, chunk[j]].item() scored.add(abs_pos) if end == len(all_ids): break results.append((total_lp,)) return results def generate_until(self, requests): results = [] for req in requests: if hasattr(req, 'args'): ctx, kwargs = req.args else: ctx, kwargs = req input_ids = self.tok_encode(ctx)[-900:] input_t = torch.tensor([input_ids], device=self._device) max_gen = kwargs.get("max_gen_toks", 64) stop = kwargs.get("until", []) for _ in range(max_gen): with torch.no_grad(): logits = self._model(input_t[:, -1024:]) if isinstance(logits, tuple): logits = logits[0] next_tok = logits[0, -1].argmax().item() input_t = torch.cat([input_t, torch.tensor([[next_tok]], device=self._device)], dim=1) gen_text = self._tokenizer.decode( input_t[0, len(input_ids):].tolist()) if any(s in gen_text for s in stop): break results.append(gen_text) return results wrapped = WrappedModel(model, tokenizer, device) tasks = ["lambada_openai", "hellaswag", "arc_easy"] try: results = evaluator.simple_evaluate( model=wrapped, tasks=tasks, batch_size=8) print(f"\n Downstream Results ({tag}):") for task_name, task_res in results.get("results", {}).items(): acc = (task_res.get("acc,none", task_res.get("acc_norm,none", task_res.get("acc", task_res.get("acc_norm", task_res.get("perplexity,none", task_res.get("perplexity", "N/A"))))))) print(f" {task_name}: {acc}") return results except Exception as e: print(f" Downstream eval failed: {e}") return {} def _downstream_metric( downstream: Dict[str, Any], task_names: Tuple[str, ...], metric_names: Tuple[str, ...], ) -> Optional[float]: """Extract a scalar metric across lm-eval task/key naming variants.""" results = downstream.get("results", {}) if isinstance(downstream, dict) else {} for task_name in task_names: task_results = results.get(task_name, {}) if not isinstance(task_results, dict): continue for metric_name in metric_names: value = task_results.get(metric_name) if isinstance(value, (int, float)): return float(value) return None def _fmt_pct(value: Optional[float]) -> str: return "N/A" if value is None else f"{value:.1%}" def _structured_group_avg( structured: Dict[str, Dict[int, float]], prefix: str, ) -> Optional[float]: vals = [ v for task_name, levels in structured.items() if task_name.startswith(prefix) and isinstance(levels, dict) for v in levels.values() if isinstance(v, (int, float)) ] return sum(vals) / len(vals) if vals else None def print_final_results(results_table: Dict[str, Dict[str, Any]]) -> None: """Print only metrics that were actually run; skipped suites show as N/A.""" downstream_cols = [ ("LAMBADA", ("lambada_openai", "lambada"), ("acc,none", "acc")), ("HellaSwag", ("hellaswag",), ("acc_norm,none", "acc_norm", "acc,none", "acc")), ("ARC-Easy", ("arc_easy",), ("acc_norm,none", "acc_norm", "acc,none", "acc")), ] has_downstream = any( bool((res.get("downstream") or {}).get("results")) for res in results_table.values() ) has_structured = any(bool(res.get("structured")) for res in results_table.values()) print(f"\n{'='*70}") print("FINAL RESULTS") table_width = 90 if has_downstream else 46 print(f"{'='*table_width}") header = f"{'Model':20s} {'Params':>10s} {'Test PPL':>10s}" if has_downstream: for label, _, _ in downstream_cols: header += f" {label:>10s}" print(header) print("-" * len(header)) for tag, res in results_table.items(): row = f"{tag:20s} {res['params']/1e6:>8.1f}M {res['test_ppl']:>10.2f}" if has_downstream: downstream = res.get("downstream", {}) for _, task_names, metric_names in downstream_cols: metric = _downstream_metric(downstream, task_names, metric_names) row += f" {_fmt_pct(metric):>10s}" print(row) if not has_downstream: print("\nDownstream eval: N/A (lm-eval not installed or downstream eval failed).") if not has_structured: print("Structured benchmark: N/A (skipped for 100M path; use PPL/downstream metrics).") return print(f"\n{'='*70}") print("STRUCTURED RESULTS") print(f"{'Model':20s} {'Struct A':>9s} {'Struct B':>9s} {'Struct C':>9s} {'Overall':>9s}") print("-" * 62) for tag, res in results_table.items(): sr = res.get("structured", {}) a_avg = _structured_group_avg(sr, "A") b_avg = _structured_group_avg(sr, "B") c_avg = _structured_group_avg(sr, "C") all_vals = [ v for levels in sr.values() if isinstance(levels, dict) for v in levels.values() if isinstance(v, (int, float)) ] overall = sum(all_vals) / len(all_vals) if all_vals else None print( f"{tag:20s} {_fmt_pct(a_avg):>9s} {_fmt_pct(b_avg):>9s} " f"{_fmt_pct(c_avg):>9s} {_fmt_pct(overall):>9s}" ) def _run_v4c_pair_preview(args): """V4c synthetic-pair preview + audit. No model build, no training. Generates 500 pairs via data_v4c_pairs.generate(), runs the family-aware audit (hard-fail on span-validity violations), reports the duplicate- collision rate (soft-fail at ≥5%), and prints the first 10 decoded items. Exits 0 on PASS, prints failures and exits nonzero on hard-fail.""" try: import data_v4c_pairs as v4c_module from transformers import GPT2TokenizerFast except Exception as e: print(f"FATAL: V4c preview imports failed: {e}") raise SystemExit(1) print("=" * 70) print(f"V4c SYNTHETIC-PAIR PREVIEW — seed={args.seed}") print("=" * 70) tok = GPT2TokenizerFast.from_pretrained("gpt2") items = v4c_module.generate(tok, n=500, seed=args.seed) report = v4c_module.audit(items, tok, verbose=True) print() print(f" Total items: {report['total']}") print(f" By family: {report['by_family']}") print(f" Hard-fails: {len(report['hard_fails'])}") print(f" Anchorable: {report['n_anchorable']}") print(f" Dup collisions: {report['duplicate_collisions']}") print() if report["hard_fails"]: print(" STATUS: HARD-FAIL — fix data_v4c_pairs.py before training") raise SystemExit(1) print(f" STATUS: span-validity PASS") n_update_items = report["total"] - report["by_family"].get( "neutral_control", 0) dup_rate = report["duplicate_collisions"] / max(1, n_update_items) print(f"\n --- Duplicate-collision audit ---") print(f" Overall: {report['duplicate_collisions']} / {n_update_items} " f"update-family items = {dup_rate:.2%}") fams_over = [] for fam, dup in report.get("duplicate_collisions_by_family", {}).items(): fam_total = report["by_family"].get(fam, 0) fam_rate = dup / max(1, fam_total) flag = "FAIL" if fam_rate >= 0.10 else "warn" if fam_rate >= 0.05 else "ok" print(f" {fam:<22s} {dup:>3d}/{fam_total:>3d} = {fam_rate:>6.2%} [{flag}]") if fam_rate >= 0.10: fams_over.append((fam, fam_rate)) if report["top_repeated_pairs"]: print(f" Top repeated (family, query, current) triples:") for entry in report["top_repeated_pairs"][:8]: print(f" {entry['family']:<22s} q={entry['query']!r:<14s} " f"c={entry['current']!r:<14s} ×{entry['count']}") if fams_over: print(f"\n STATUS: HARD-FAIL — {len(fams_over)} families ≥10% dup rate:") for fam, rate in fams_over: print(f" {fam}: {rate:.2%}") print(f" Expand the corresponding pools in data_v4c_pairs.py " f"(NAMES/OBJECTS/CAPITALS_OLD/GATES/PARAPHRASE_*_CITIES) " f"before training. Reviewer-blocker: at this rate, InfoNCE " f"in-batch FN masking is doing major work every batch.") raise SystemExit(1) if dup_rate >= 0.05: print(f" WARN: overall rate {dup_rate:.2%} ≥ 5%; under per-family " f"10% hard-fail bar, but worth monitoring.") print() print(" --- First 10 decoded pairs ---") for i, it in enumerate(items[:10]): print(f"\n Item {i} family={it.family}") text = tok.decode(it.ids).replace("\n", " ").strip() print(f" text: {text!r}") decoded = v4c_module.decode_spans(it, tok) for k, v in decoded.items(): if v: print(f" {k:>21s}: {v!r}") def _run_v3m2_preview(args): """Build the V3m2 model, load shared backbone weights from an existing checkpoint (strict=False, with audit), run ONE forward in eval mode on a 16-seq WikiText sample, and print: (a) decoded anchor_negative_ids / label_negative_ids audit (b) belief-revision-lexicon eligibility check (must all be True) (c) per-sentence anchor + generated-label inspection for first 5 seqs. Aborts BEFORE training if (b) fails or if labels are still firing mostly on periods / function words.""" if args.model not in ("CEDL-V3m2", "CEDL-V3m2-NLL"): print(f"ERROR: --preview-only requires --model CEDL-V3m2 or " f"CEDL-V3m2-NLL, got {args.model!r}") return if not args.preview_shared_checkpoint: print("ERROR: --preview-only requires --preview-shared-checkpoint " "(path to V0/V3/V3m checkpoint to load shared backbone " "weights from). Random-init preview is meaningless.") return print("=" * 70) print(f"V3m2 TARGET-PREVIEW — model={args.model}") print(f" shared-backbone checkpoint: {args.preview_shared_checkpoint}") print("=" * 70) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") torch.manual_seed(args.seed) np.random.seed(args.seed) model = build_model(args.model, vocab=50257, max_seq=1024).to(device) model.eval() state = torch.load(args.preview_shared_checkpoint, map_location="cpu", weights_only=True) if any(k.startswith("_orig_mod.") for k in state): state = {k.replace("_orig_mod.", ""): v for k, v in state.items()} res = model.load_state_dict(state, strict=False) print(f"\n Loaded shared weights with strict=False:") print(f" missing keys (V3m2-new): {len(res.missing_keys)}") for k in res.missing_keys[:12]: print(f" - {k}") if len(res.missing_keys) > 12: print(f" ... ({len(res.missing_keys) - 12} more)") print(f" unexpected keys (V3/V3m-only): {len(res.unexpected_keys)}") for k in res.unexpected_keys[:12]: print(f" - {k}") if len(res.unexpected_keys) > 12: print(f" ... ({len(res.unexpected_keys) - 12} more)") raw = model if not hasattr(raw, "anchor_negative_ids"): print("\nFATAL: model has no anchor_negative_ids buffer; lexicon " "construction failed at build time. Cannot preview.") return try: from transformers import GPT2TokenizerFast tok = GPT2TokenizerFast.from_pretrained("gpt2") except Exception as e: print(f"\nWARN: transformers unavailable ({e}); decoded mask audit " f"skipped.") tok = None print(f"\n --- (a) Decoded mask audit ---") print(f" anchor_negative_ids (N={raw.anchor_negative_ids.numel()}):") if tok is not None and raw.anchor_negative_ids.numel() > 0: decoded = [repr(tok.decode([int(i)])) for i in raw.anchor_negative_ids[:8].tolist()] print(f" sample: {decoded}") print(f" label_negative_ids (N={raw.label_negative_ids.numel()}):") if tok is not None and raw.label_negative_ids.numel() > 0: decoded = [repr(tok.decode([int(i)])) for i in raw.label_negative_ids[:30].tolist()] print(f" first 30: {decoded}") print(f"\n --- (b) Belief-revision eligibility check ---") abort = False if tok is not None: belief_words = ["not", "no", "never", "but", "however", "although", "now", "currently", "previously", "instead", "all", "every", "some", "any"] label_neg_set = set(raw.label_negative_ids.tolist()) for w in belief_words: for form in (w, " " + w): ids = tok.encode(form, add_special_tokens=False) if not ids: continue tid = ids[0] eligible = tid not in label_neg_set marker = "OK " if eligible else "FAIL" print(f" [{marker}] {form!r:>14s} -> token_id={tid:>5d} " f"label_eligible={eligible}") if not eligible: abort = True if abort: print("\n ABORT: a belief-revision marker is in label_negative_ids. " "Fix V3M2_LABEL_NEGATIVE_WORDS in CEDL.py and re-run.") return print(f"\n --- (c) Per-sentence anchor + label inspection ---") try: from datasets import load_dataset ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="validation") texts = [t for t in ds["text"] if len(t.strip()) > 200][:128] if tok is None: print(" skipping (tokenizer unavailable)") return seqs = [] for t in texts: ids = tok.encode(t, add_special_tokens=False)[:512] if len(ids) >= 64: seqs.append(ids) if len(seqs) >= 64: break if not seqs: print(" no eligible sequences found; aborting preview") return T = min(len(s) for s in seqs) ids_tensor = torch.tensor([s[:T] for s in seqs], dtype=torch.long, device=device) except Exception as e: print(f" WARN: could not load WikiText ({e}); aborting preview") return calib_n = min(8, max(1, ids_tensor.size(0) // 4)) calib_batches = [ids_tensor[i * 4:(i + 1) * 4] for i in range(calib_n) if (i + 1) * 4 <= ids_tensor.size(0)] needs_calib = (raw.curve_C_ema_initialized.item() == 0.0) if needs_calib: print(f" EMA calibration: manifold buffers are UNINITIALIZED " f"(shared checkpoint was likely V0/V3, not V3m/V3m2). " f"Running {len(calib_batches)} warmup batches to populate " f"curve_C/curve_E/cosmis EMA before preview...") with torch.no_grad(): for cb in calib_batches: hC = raw.c_stage(cb, feedback=None) hE, _ = raw.e_stage(hC, feedback=None) _build_v3m2_targets( ids=cb, h_C_p0=hC, h_E_p0=hE, logits_mem_p0=None, model=raw, salience_mode=raw.salience_mode, update_ema=True, with_nll=False) cinit = raw.curve_C_ema_initialized.item() cvar = raw.curve_C_ema_var.item() evar = raw.curve_E_ema_var.item() cosvar = raw.cosmis_ema_var.item() print(f" after calib: curve_C_init={cinit:.4f} " f"curve_C_var={cvar:.4f} curve_E_var={evar:.4f} " f"cosmis_var={cosvar:.4f}") if cinit < 0.05 or cvar < 1e-6 or evar < 1e-6 or cosvar < 1e-6: print(" ABORT: EMA buffers did not move after warmup. " "Anchor threshold cannot be calibrated. Use a V3m/V3m2 " "checkpoint instead, or extend calibration batches.") return with torch.no_grad(): h_C = raw.c_stage(ids_tensor, feedback=None) h_E, _ = raw.e_stage(h_C, feedback=None) target, w, dbg = _build_v3m2_targets( ids=ids_tensor, h_C_p0=h_C, h_E_p0=h_E, logits_mem_p0=None, model=raw, salience_mode=raw.salience_mode, update_ema=False, with_nll=False, ) anchors = dbg["anchors"] strength = dbg["strength"] n_show = min(5, ids_tensor.size(0)) label_neg_set = set(raw.label_negative_ids.tolist()) sent_end_ids = set() if tok is not None: for s in (".", "!", "?", " .", " !", " ?"): try: ids_s = tok.encode(s, add_special_tokens=False) if ids_s: sent_end_ids.add(ids_s[0]) except Exception: pass n_pos = 0 n_pos_in_label_neg = 0 n_pos_sent_initial = 0 pos_mask = target > 0.05 for b in range(ids_tensor.size(0)): for p in pos_mask[b].nonzero(as_tuple=True)[0].tolist(): n_pos += 1 tid = int(ids_tensor[b, p].item()) if tid in label_neg_set: n_pos_in_label_neg += 1 prev_id = int(ids_tensor[b, p - 1].item()) if p > 0 else -1 if prev_id in sent_end_ids: n_pos_sent_initial += 1 for b in range(n_show): anchor_pos = anchors[b].nonzero(as_tuple=True)[0].tolist() print(f"\n seq {b} T={ids_tensor.size(1)} " f"#anchors={len(anchor_pos)}") for p in anchor_pos[:5]: tok_id = int(ids_tensor[b, p].item()) decoded = tok.decode([tok_id]) s = float(strength[b, p].item()) print(f" anchor @ pos={p:>3d} token={decoded!r:>12s} " f"strength={s:.3f}") pos_targets = (target[b] > 0.05).nonzero(as_tuple=True)[0].tolist() labels_show = pos_targets[:10] if labels_show: print(f" positive-label tokens:") for p in labels_show: tok_id = int(ids_tensor[b, p].item()) decoded = tok.decode([tok_id]) v = float(target[b, p].item()) print(f" pos={p:>3d} token={decoded!r:>12s} " f"target={v:.3f}") print(f"\n --- preview hard checks (N={n_pos} positive labels total) ---") if n_pos == 0: print(" WARN: no positive labels produced — anchor threshold may " "be too strict on uncalibrated EMA, or input is genuinely flat.") else: frac_neg = n_pos_in_label_neg / n_pos frac_si = n_pos_sent_initial / n_pos flag_neg = "PASS" if frac_neg <= 0.05 else "FAIL" flag_si = "PASS" if frac_si <= 0.50 else "FAIL" print(f" positive labels IN label_negative_ids: {frac_neg:>5.3f} " f"(criterion 3 target <= 0.05) [{flag_neg}]") print(f" positive labels at sentence-initial: {frac_si:>5.3f} " f"(preview soft target <= 0.50) [{flag_si}]") if flag_neg == "FAIL": print(" ABORT: label mask is letting punctuation/stopwords " "through. Fix V3M2_LABEL_NEGATIVE_WORDS in CEDL.py.") print(f"\n --- (d) Curated discourse-prompts preview ---") curated = [ ("Alice wanted to leave early, however she stayed until the end.", "explicit"), ("The forecast called for rain, but the day turned out sunny.", "explicit"), ("Although Tom had studied hard, the exam was very difficult.", "explicit"), ("John loves coffee. He never drinks it.", "implicit"), ("Mary said she would attend. She stayed home.", "implicit"), ("The cake was delicious. Nobody ate it.", "implicit"), ("Alice used to live in Paris. She now lives in Berlin.", "belief"), ("The capital was Bonn until 1990. Today it is Berlin.", "belief"), ("Previously the gate was A12. Currently it is C7.", "belief"), ("Alice owns a hat. Bob owns a pen. Carol owns a cup. " "But Alice now owns a car. What does Alice own now? Answer:", "stale_current"), ("Bob owns a ring. Dave owns a lamp. Eve owns a book. " "But Bob now owns a kite. What does Bob own now? Answer:", "stale_current"), ("Carol owns a bag. Frank owns a fork. Grace owns a desk. " "But Carol now owns a fish. What does Carol own now? Answer:", "stale_current"), ] MARKER_FORMS = ["however", "but", "although", "yet", "instead", "nevertheless", "though", "never", "no", "not", "nobody", "nothing", "now", "currently", "previously", "today", "originally"] marker_ids = set() for w in MARKER_FORMS: for f in (w, " " + w, w.capitalize(), " " + w.capitalize()): try: ids = tok.encode(f, add_special_tokens=False) if ids: marker_ids.add(ids[0]) except Exception: pass n_on_marker = 0 n_total_anchors = 0 for text, cat in curated: ids = tok.encode(text, add_special_tokens=False) ids_t = torch.tensor([ids], dtype=torch.long, device=device) with torch.no_grad(): hC = raw.c_stage(ids_t, feedback=None) hE, _ = raw.e_stage(hC, feedback=None) tgt, _w, dbg = _build_v3m2_targets( ids=ids_t, h_C_p0=hC, h_E_p0=hE, logits_mem_p0=None, model=raw, salience_mode=raw.salience_mode, update_ema=False, with_nll=False) apos = dbg["anchors"][0].nonzero(as_tuple=True)[0].tolist() lpos = (tgt[0] > 0.05).nonzero(as_tuple=True)[0].tolist() marker_hit = any(int(ids_t[0, p].item()) in marker_ids for p in apos) n_total_anchors += len(apos) n_on_marker += sum(1 for p in apos if int(ids_t[0, p].item()) in marker_ids) anc_str = ", ".join( f"{p}={repr(tok.decode([int(ids_t[0, p].item())]))}" for p in apos) or "(none)" lbl_str = ", ".join( f"{p}={repr(tok.decode([int(ids_t[0, p].item())]))}" f":{tgt[0, p].item():.2f}" for p in lpos[:8]) or "(none)" flag = " [MARKER]" if marker_hit else " [no-marker]" print(f" [{cat:>13s}] {text[:64]}...") print(f" anchors: {anc_str}{flag}") print(f" labels: {lbl_str}") frac_marker = (n_on_marker / max(1, n_total_anchors)) flag_m = "PASS" if frac_marker >= 0.40 else "FAIL" print(f"\n anchor-on-discourse-marker fraction: {frac_marker:>5.3f} " f"(curated target >= 0.40) [{flag_m}]") print(f" (anchors here should land on discourse markers — " f"however/but/now/currently/never/etc. — and labels should " f"land on the contrasting clause that follows.)") print(f"\n Preview complete. Review printed anchors/labels above for " f"qualitative confirmation.") def main(): parser = argparse.ArgumentParser(description="CEDL 100M Benchmark") parser.add_argument("--model", type=str, default="CEDL", help="Model tag, comma-separated list, or 'all'. " "Examples: CEDL | CEDL-B0 | CEDL,CEDL-B0 | all") parser.add_argument("--dataset", type=str, default="wikitext103") parser.add_argument("--tpu", action="store_true") parser.add_argument("--eval-only", action="store_true") parser.add_argument("--checkpoint", type=str, default=None) parser.add_argument("--verify-params", action="store_true", help="Print param counts and exit") parser.add_argument("--max-steps", type=int, default=30_000) parser.add_argument("--eval-interval", type=int, default=None, help="Eval cadence in steps (Config default 1000). " "Lower for short Stage-0 runs so the λ trajectory " "/ PPL trend has more datapoints.") parser.add_argument("--save-interval", type=int, default=None, help="Periodic checkpoint cadence in steps (Config " "default 5000).") parser.add_argument("--total-steps", type=int, default=None, help="V5-LM: GLOBAL target step (alias for --max-steps; " "schedules key off the global step). If set, " "overrides --max-steps.") parser.add_argument("--resume-checkpoint", type=str, default=None, help="V5-LM: resume from a checkpoint. Bare state_dict " "(e.g. CEDL-V4e_best.pt) → fresh AdamW; train-state " "dict {model,optimizer,global_step} → restore all.") parser.add_argument("--start-step", type=int, default=0, help="V5-LM: global step to resume the loop counter at " "(e.g. 5000). Schedules continue from here.") parser.add_argument("--run-until-step", type=int, default=None, help="V5-LM: STOP the loop at this global step WITHOUT " "compressing the schedule (which uses --total-steps " "as its denominator). Stage 0: --total-steps 30000 " "--run-until-step 6000.") parser.add_argument("--save-trainstate", action="store_true", help="V5-LM/B0-long: also write resume-only " "{model,optimizer,global_step} checkpoints for " "multi-session continuation.") parser.add_argument("--save-dir", type=str, default=None, help="Checkpoint/scorecard output dir. For Colab " "multi-session V5, point at Drive (e.g. " "/content/drive/MyDrive/cedl_ckpt) so train-state " "survives session resets. Use a SEPARATE dir per " "run to avoid mixing scorecard logs across runs.") parser.add_argument("--v5-grad-isolation", action="store_true", help="V5-LM: route the V4c aux gradient to the salience " "pathway only (trunk trains on LM-CE). Auto-enabled " "for CEDL-V5-LM* tags.") parser.add_argument("--v5-trunk-aux-frac", type=float, default=0.0, help="V5-LM: fraction of aux gradient still allowed into " "the trunk (0.0 = full isolation; Stage-0 may use " "0.05).") parser.add_argument("--v5-aux-scale-early", type=float, default=None, help="V5-LM: aux scale for global steps 5K-10K (default 0.50)") parser.add_argument("--v5-aux-scale-mid", type=float, default=None, help="V5-LM: aux scale for global steps 10K-20K (default 0.25)") parser.add_argument("--v5-aux-scale-late", type=float, default=None, help="V5-LM: aux scale floor for steps 20K+ (default 0.10)") parser.add_argument("--v5-cadence-early", type=int, default=None, help="V5-LM: v4c_every for steps 5K-10K (default 8)") parser.add_argument("--v5-cadence-mid", type=int, default=None, help="V5-LM: v4c_every for steps 10K-20K (default 16)") parser.add_argument("--v5-cadence-late", type=int, default=None, help="V5-LM: v4c_every for steps 20K+ (default 32)") parser.add_argument("--v5-family-reweight", action="store_true", help="V5-LM: oversample but_update (0.30→0.50) and " "drop neutral (0.20→0.10) in V4c batches — " "Path B for B1 (but_update is the weakest B2 " "family). No test-format contamination.") parser.add_argument("--v6-mixture", action="store_true", help="V6: documentary flag. The mixture is wired by " "the CEDL-V6 tag itself (build_model only " "registers bank_q_proj + v6_lambda_* for that " "tag), so this flag is auto-implied for CEDL-V6 " "and a NO-OP on every other tag — passing it " "with e.g. CEDL-V5-LM logs a warning and is " "ignored. Use --v6-lambda-init to override the " "initial λ logit.") parser.add_argument("--v6-lambda-init", type=float, default=-4.0, help="V6: BIAS LOGIT for the λ sigmoid gate " "(λ = sigmoid(v6_lambda_a · v_scalar + v6_lambda_b); " "v6_lambda_b inits to this value). Actual λ̄ also " "depends on v_scalar's distribution from the V5 " "trainstate — observed range at init: -4.0 → " "λ̄≈0.001, -2.0 → λ̄≈0.06 (empirical, NOT just " "the bias sigmoid). Try -2.0 if λ stays " "collapsed in Stage 0.") parser.add_argument("--v6-lambda-a-init", type=float, default=1.0, help="V6.1b: slope (a) of the λ affine over v_scalar. " "Default +1.0 matches V6.1 v1/v2/v3. V6.1b " "sign-correction sets -1.0 (V4c-ans has LOWER " "v_scalar than WikiText per d=-0.85 separability, " "so negative slope is needed for V4c-ans to get " "HIGHER λ). Only meaningful for the CEDL-V6 tag; " "ignored elsewhere with a WARN.") parser.add_argument("--v6-aux-weight", type=float, default=0.0, help="V6.1: global scale on the V4c-routed bank-aux " "loss (L_bank + v6_mix_weight·L_mix + " "v6_gate_weight·L_gate). 0.0 = V6 baseline " "(mixture-only, falsified). 0.05 = V6.1 main run.") parser.add_argument("--v6-margin-target", type=float, default=1.0, help="V6.1: hinge target (nats) for L_bank and L_mix.") parser.add_argument("--v6-mix-weight", type=float, default=0.5, help="V6.1: weight on the mixture-output margin L_mix. " "Mostly 0 in practice (trunk margins ≈ 30 nats > " "target=1), but trains v6_lambda_a/b when bank " "beats trunk on a few examples.") parser.add_argument("--v6-gate-weight", type=float, default=0.01, help="V6.1: weight on the LOGIT-SPACE λ floor " "regularizer L_gate (anti-collapse insurance). " "Sigmoid-space floor has vanishing gradient when " "λ collapses; logit-space gives constant gradient.") parser.add_argument("--v6-lambda-floor", type=float, default=0.05, help="V6.1: λ floor at answer positions (default 0.05; " "logit ≈ -2.94). L_gate fires when λ < this.") parser.add_argument("--v6-lambda-head", action="store_true", help="Stage 2a: replace the v_scalar-coupled affine " "gate with a separate λ head MLP(h_D). Gradient " "is aux-only (torch.no_grad() in main forward + " "V5 allowlist for the head's params), so LM-CE " "cannot drive λ→0 the way it did in V6/V6.1/V6.1b.") parser.add_argument("--v6-lambda-head-hidden", type=int, default=160, help="Stage 2a: hidden dim of the λ head MLP (d/4=160 " "default at d=640). ~102K total head params.") parser.add_argument("--v6-lambda-head-bias-init", type=float, default=-7.0, help="Stage 2a: final-layer bias init (sigmoid(-7)≈9e-4 " "so λ starts near zero everywhere; the head LEARNS " "h_D→λ selectivity from L_gate + L_bg supervision).") parser.add_argument("--v6-bg-weight", type=float, default=1.0, help="Stage 2a: weight on L_bg, the background-row " "sparsity hinge. Asymmetric counterpart to " "L_gate — pushes λ DOWN at sampled non-answer V4c " "rows so smoke tests measure selectivity rather " "than bias lift.") parser.add_argument("--v6-bg-target", type=float, default=0.01, help="Stage 2a: λ target at background V4c rows; " "L_bg = relu(λ_bg - target).mean().") parser.add_argument("--v6-bce-objective", action="store_true", help="Stage 2a-bce: replace hinge L_gate + L_bg with " "class-balanced BCE (target=1 at ans, target=0 " "at bg, 0.5/0.5 mean). Strong gradient at the " "floor exactly proportional to how wrong each " "row is — fixes the stalemate where the hinge " "objective could not lift λ_ans above floor.") parser.add_argument("--v6-sel-weight", type=float, default=1.0, help="Stage 2a-bce: weight on L_sel (subsumes " "gate_weight + bg_weight under BCE). Default 1.0 " "with v6_aux_weight=0.05 gives effective scale " "0.05 — 20× the failed gate=0.05 hinge run.") parser.add_argument("--v6-lambda-head-w-init-std", type=float, default=1e-3, help="Stage 2a-bce: final-layer W init std. Default " "1e-3 keeps prior Stage 2a runs bit-reproducible. " "Try 0.05 to let the head's output vary across " "positions from step 0 (BCE shapes existing " "differentiation instead of building it). Tail " "at bias=-3, std=0.05 → λ_max ≈ 0.25 at 3σ " "outliers, still well below 0.7 hard kill.") parser.add_argument("--v6-wt-sparsity-weight", type=float, default=0.0, help="Stage 2a-wts: WikiText-side sparsity loss weight. " "BCE alone supervises V4c-ans vs V4c-bg; the head " "learns ANY separating axis, which can over-fire " "on WikiText (empirically: w_init scan showed " "this is geometric, not optimization-driven). " "This term adds mean-λ-on-WikiText to aux_loss " "in the main forward, with gradient flowing to " "the head via a parallel non-no_grad path. " "Mixture lam stays detached → LM-CE still cannot " "drive the head. Try 0.05 (parity with v6_aux).") parser.add_argument("--v6-wt-sparsity-target", type=float, default=0.0, help="Stage 2a-wts: hinge target. 0 = mean-λ form " "(empirically crushes lam_ans because Adam " "normalizes consistent-sign gradient). >0 = " "hinge form `relu(λ-target).mean()`: only " "suppresses positions ABOVE target, leaving a " "free zone in [0,target] where BCE can lift " "lam_ans without WT pressure dragging it back. " "Try 0.05 (parity with v6_lambda_floor).") parser.add_argument("--v6-mem-head-bank", action="store_true", help="Stage 2b: enable untied bank vocab head. The " "existing bank readout uses l_stage.mem_head, " "whose .weight is WEIGHT-TIED to tok_emb. B2 " "probe confirmed the resulting bank distribution " "is near-uniform (entropy 9.15 / 10.83). This " "flag adds self.mem_head_bank = nn.Linear(d, V) " "as a SEPARATE trainable param initialized from " "tok_emb (copy, not alias) + l_stage.mem_head." "bias. Bank branch (V4c hook + main mixture) " "uses it when set. Aux-only via V5 allowlist.") parser.add_argument("--v6-bank-ce-weight", type=float, default=0.0, help="Stage 2b: weight on direct cross-entropy of " "logits_bank_a at V4c-answer rows vs cur_tok. " "The existing L_bank margin hinge is satisfiable " "through near-uniform logits with cur_lp slightly " "above stale_lp (B2 confirmed). Adding CE anchors " "the bank distribution to put high mass on the " "correct token. Try 1.0 (CE is on natural log " "scale; values <0.5 will be dominated by other " "losses).") parser.add_argument("--v6-bank-pair-ce-weight", type=float, default=0.0, help="Stage 2f: weight on pairwise CE over " "[bank_logit(cur), bank_logit(stale)] at V4c " "answer rows. This directly trains current-vs-" "stale role binding; use when full-vocab bank CE " "learns semantic candidates but leaves stale tied " "or above current. Try 2.0.") parser.add_argument("--v6-bank-query-source", type=str, default="h_d", choices=("h_d", "h_e", "q_attractor", "q_mem"), help="Stage 2c: source vector passed through " "bank_q_proj for the external V6 bank. h_d is " "the original V6 path; h_e uses the dense " "E-stage state; q_attractor/q_mem use DStage " "internal memory states.") parser.add_argument("--v6-bank-readout-mode", type=str, default="bank", choices=("bank", "direct"), help="Stage 2d: V6 external readout mode. bank keeps " "source→bank_q_proj→mem_vals→mem_head_bank. " "direct bypasses the 256-slot bank bottleneck " "and uses source→mem_head_bank.") parser.add_argument("--v6-source-adapter", action="store_true", help="Stage 2e: add a small nonlinear residual source " "adapter before the direct V6 readout head. Use " "with B0 specialist mode when direct-source probes " "rank current/stale candidates but cannot separate " "current from stale.") parser.add_argument("--v6-context-adapter", action="store_true", help="Stage 2g: add a causal prefix self-attention " "adapter before the direct V6 readout head. Use " "when answer-row source vectors rank plausible " "candidates but fail current-vs-stale role binding.") parser.add_argument("--v6-bank-head-lr", type=float, default=0.0, help="Stage 2b-lr: dedicated CONSTANT lr for " "mem_head_bank.* (separate AdamW group, no " "cosine schedule). The 50K-output head needs " "~10⁴+ CE examples to calibrate but V4c aux " "fires sparsely; without a dedicated higher " "lr the bank head doesn't move (Stage 2b " "smoke at main lr=1e-5 confirmed L_bank_ce " "stuck at log(V)=10.83 over 500 steps). " "Default 0.0 = use the main schedule lr. " "Try 3e-4 (~30× the tail lr).") parser.add_argument("--v6-specialist-from-b0", action="store_true", help="B0-preserving V6 specialist mode: build CEDL-V6, " "partially load matching weights from " "--resume-checkpoint, freshly initialize only the " "V6/salience specialist keys, and use a fresh " "optimizer. Requires CEDL-V6, direct q_mem readout, " "--v6-lambda-head, --v6-mem-head-bank, and " "--v6-specialist-noinject.") parser.add_argument("--v6-specialist-freeze", type=str, default="none", choices=("none", "trunk"), help="Freeze policy for --v6-specialist-from-b0. " "'trunk' freezes the B0-shared LM path and trains " "only v6_lambda_head.* and mem_head_bank.*.") parser.add_argument("--v6-specialist-noinject", action="store_true", help="For CEDL-V6 specialist mode, build the V4d shell " "with v4d_noinject=True so random/fresh salience " "injection cannot perturb the B0 trunk.") parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--grad-accum", type=int, default=4) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--preview-only", action="store_true", help="V3m2 only: build the model, load shared " "backbone weights via --preview-shared-checkpoint " "(strict=False), run ONE forward in eval mode on " "a 16-seq WikiText sample, print the decoded " "mask audit + first-5-sequence anchor / label " "inspection, then exit. No training, no save.") parser.add_argument("--preview-shared-checkpoint", type=str, default=None, help="Path to an existing checkpoint (V0/V3/V3m) to " "load shared backbone weights from for " "--preview-only mode. SL-specific keys (scalar_" "head, EMA buffers, lexicon buffers) will be " "missing — strict=False is used, and missing/" "unexpected keys are printed for audit.") parser.add_argument("--preview-v4c-pairs", action="store_true", help="V4c only: generate 500 synthetic contrastive " "pairs, run the family-aware audit (hard-fail on " "span-validity violations), and print the first " "10 pairs with decoded spans. No model build, " "no training. Use BEFORE training to verify the " "pair generator is well-formed.") args = parser.parse_args() if args.verify_params: verify_all_params() return if args.preview_only: return _run_v3m2_preview(args) if args.preview_v4c_pairs: return _run_v4c_pair_preview(args) cfg = Config( dataset=args.dataset, tpu=args.tpu, max_steps=(args.total_steps if args.total_steps is not None else args.max_steps), batch_size=args.batch_size, grad_accum=args.grad_accum, seed=args.seed, resume_checkpoint=args.resume_checkpoint, resume_start_step=args.start_step, save_trainstate=args.save_trainstate, run_until_step=args.run_until_step, **({"save_dir": args.save_dir} if args.save_dir else {}), **({"eval_interval": args.eval_interval} if args.eval_interval is not None else {}), **({"save_interval": args.save_interval} if args.save_interval is not None else {}), **({"v5_aux_scale_early": args.v5_aux_scale_early} if args.v5_aux_scale_early is not None else {}), **({"v5_aux_scale_mid": args.v5_aux_scale_mid} if args.v5_aux_scale_mid is not None else {}), **({"v5_aux_scale_late": args.v5_aux_scale_late} if args.v5_aux_scale_late is not None else {}), **({"v5_cadence_early": args.v5_cadence_early} if args.v5_cadence_early is not None else {}), **({"v5_cadence_mid": args.v5_cadence_mid} if args.v5_cadence_mid is not None else {}), **({"v5_cadence_late": args.v5_cadence_late} if args.v5_cadence_late is not None else {}), **({"v5_family_reweight": True} if args.v5_family_reweight else {}), **({"v6_specialist_from_b0": True} if args.v6_specialist_from_b0 else {}), **({"v6_specialist_freeze": args.v6_specialist_freeze} if args.v6_specialist_freeze != "none" else {}), **({"v6_specialist_noinject": True} if args.v6_specialist_noinject else {}), ) torch.manual_seed(cfg.seed) np.random.seed(cfg.seed) torch.set_float32_matmul_precision('medium') torch.backends.cudnn.benchmark = True if cfg.tpu: import torch_xla.core.xla_model as xm device = xm.xla_device() print(f"Using TPU: {device}") elif torch.cuda.is_available(): device = torch.device("cuda") print(f"Using GPU: {torch.cuda.get_device_name()}") else: device = torch.device("cpu") print("Using CPU (will be slow!)") train_loader, val_loader, test_loader, tokenizer = load_data(cfg) _train_ds = train_loader.dataset _val_ds = val_loader.dataset _test_ds = test_loader.dataset if args.model == "all": models_to_run = ALL_MODELS elif "," in args.model: models_to_run = [t.strip() for t in args.model.split(",") if t.strip()] else: models_to_run = [args.model] print(f"Models to run ({len(models_to_run)}): {models_to_run}") results_table = {} for tag in models_to_run: torch.manual_seed(cfg.seed) np.random.seed(cfg.seed) if device.type == 'cuda': torch.cuda.manual_seed_all(cfg.seed) if tag in ("CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m", "CEDL-V3m2", "CEDL-V3m2-NLL", "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold", "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M", "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject", "CEDL-V4d-strong", "CEDL-V4d-strong-noinject", "CEDL-V4d-causal", "CEDL-V4d-causal-strong", "CEDL-V4d-causalz", "CEDL-V4d-combo", "CEDL-V4e", "CEDL-V4e-noinject", "CEDL-V5-LM", "CEDL-V6", "CEDL-V3lex", "CEDL-B0") or tag.startswith("CEDL-V4e-k"): cfg.batch_size = min(args.batch_size, 4) else: cfg.batch_size = args.batch_size cfg.v5_grad_isolation = bool(args.v5_grad_isolation or tag.startswith("CEDL-V5-LM") or tag == "CEDL-V6") cfg.v5_trunk_aux_frac = float(args.v5_trunk_aux_frac) if args.v6_mixture and tag != "CEDL-V6": print(f" [V6] WARN: --v6-mixture is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_mixture = (tag == "CEDL-V6") cfg.v6_lambda_init = float(args.v6_lambda_init) if args.v6_lambda_a_init != 1.0 and tag != "CEDL-V6": print(f" [V6.1b] WARN: --v6-lambda-a-init={args.v6_lambda_a_init} " f"is a no-op on tag {tag!r} (only CEDL-V6 wires the " f"mixture). Ignoring.") cfg.v6_lambda_a_init = float(args.v6_lambda_a_init) if tag == "CEDL-V6" else 1.0 if args.v6_aux_weight > 0 and tag != "CEDL-V6": print(f" [V6.1] WARN: --v6-aux-weight={args.v6_aux_weight} is a " f"no-op on tag {tag!r} (only CEDL-V6 wires the mixture). " f"Ignoring.") cfg.v6_aux_weight = float(args.v6_aux_weight) if tag == "CEDL-V6" else 0.0 cfg.v6_margin_target = float(args.v6_margin_target) cfg.v6_mix_weight = float(args.v6_mix_weight) cfg.v6_gate_weight = float(args.v6_gate_weight) cfg.v6_lambda_floor = float(args.v6_lambda_floor) if args.v6_lambda_head and tag != "CEDL-V6": print(f" [Stage2a] WARN: --v6-lambda-head is a no-op on tag " f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_lambda_head = bool(args.v6_lambda_head) and tag == "CEDL-V6" cfg.v6_lambda_head_hidden = int(args.v6_lambda_head_hidden) cfg.v6_lambda_head_bias_init = float(args.v6_lambda_head_bias_init) cfg.v6_bg_weight = float(args.v6_bg_weight) if tag == "CEDL-V6" else 1.0 cfg.v6_bg_target = float(args.v6_bg_target) if args.v6_bce_objective and tag != "CEDL-V6": print(f" [Stage2a-bce] WARN: --v6-bce-objective is a no-op on " f"tag {tag!r} (only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bce_objective = bool(args.v6_bce_objective) and tag == "CEDL-V6" cfg.v6_sel_weight = float(args.v6_sel_weight) if tag == "CEDL-V6" else 1.0 cfg.v6_lambda_head_w_init_std = float(args.v6_lambda_head_w_init_std) if args.v6_wt_sparsity_weight > 0 and tag != "CEDL-V6": print(f" [Stage2a-wts] WARN: --v6-wt-sparsity-weight=" f"{args.v6_wt_sparsity_weight} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_wt_sparsity_weight = ( float(args.v6_wt_sparsity_weight) if tag == "CEDL-V6" else 0.0) cfg.v6_wt_sparsity_target = float(args.v6_wt_sparsity_target) if args.v6_mem_head_bank and tag != "CEDL-V6": print(f" [Stage2b] WARN: --v6-mem-head-bank is a no-op on tag " f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_mem_head_bank = bool(args.v6_mem_head_bank) and tag == "CEDL-V6" if args.v6_bank_ce_weight > 0 and tag != "CEDL-V6": print(f" [Stage2b] WARN: --v6-bank-ce-weight=" f"{args.v6_bank_ce_weight} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bank_ce_weight = ( float(args.v6_bank_ce_weight) if tag == "CEDL-V6" else 0.0) if args.v6_bank_pair_ce_weight > 0 and tag != "CEDL-V6": print(f" [Stage2f] WARN: --v6-bank-pair-ce-weight=" f"{args.v6_bank_pair_ce_weight} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bank_pair_ce_weight = ( float(args.v6_bank_pair_ce_weight) if tag == "CEDL-V6" else 0.0) if args.v6_bank_query_source != "h_d" and tag != "CEDL-V6": print(f" [Stage2c] WARN: --v6-bank-query-source=" f"{args.v6_bank_query_source!r} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bank_query_source = ( str(args.v6_bank_query_source) if tag == "CEDL-V6" else "h_d") if args.v6_bank_readout_mode != "bank" and tag != "CEDL-V6": print(f" [Stage2d] WARN: --v6-bank-readout-mode=" f"{args.v6_bank_readout_mode!r} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bank_readout_mode = ( str(args.v6_bank_readout_mode) if tag == "CEDL-V6" else "bank") if args.v6_source_adapter and tag != "CEDL-V6": print(f" [Stage2e] WARN: --v6-source-adapter is a no-op on tag " f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_source_adapter = bool(args.v6_source_adapter) and tag == "CEDL-V6" if args.v6_context_adapter and tag != "CEDL-V6": print(f" [Stage2g] WARN: --v6-context-adapter is a no-op on tag " f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_context_adapter = ( bool(args.v6_context_adapter) and tag == "CEDL-V6") if args.v6_bank_head_lr > 0 and tag != "CEDL-V6": print(f" [Stage2b-lr] WARN: --v6-bank-head-lr=" f"{args.v6_bank_head_lr} is a no-op on tag {tag!r} " f"(only CEDL-V6 wires the mixture). Ignoring.") cfg.v6_bank_head_lr = ( float(args.v6_bank_head_lr) if tag == "CEDL-V6" else 0.0) if args.v6_specialist_from_b0 and tag != "CEDL-V6": raise RuntimeError( "--v6-specialist-from-b0 is only valid with --model CEDL-V6") cfg.v6_specialist_from_b0 = ( bool(args.v6_specialist_from_b0) and tag == "CEDL-V6") cfg.v6_specialist_freeze = ( str(args.v6_specialist_freeze) if cfg.v6_specialist_from_b0 else "none") cfg.v6_specialist_noinject = ( bool(args.v6_specialist_noinject) and tag == "CEDL-V6") if cfg.v6_specialist_from_b0: if not cfg.resume_checkpoint: raise RuntimeError( "--v6-specialist-from-b0 requires --resume-checkpoint " "pointing at the B0 source checkpoint.") if cfg.v6_specialist_freeze != "trunk": raise RuntimeError( "--v6-specialist-from-b0 requires " "--v6-specialist-freeze trunk.") if not cfg.v6_specialist_noinject: raise RuntimeError( "--v6-specialist-from-b0 requires " "--v6-specialist-noinject.") if not cfg.v6_lambda_head: raise RuntimeError( "--v6-specialist-from-b0 requires --v6-lambda-head.") if not cfg.v6_mem_head_bank: raise RuntimeError( "--v6-specialist-from-b0 requires --v6-mem-head-bank.") if cfg.v6_bank_query_source != "q_mem": raise RuntimeError( "--v6-specialist-from-b0 requires " "--v6-bank-query-source q_mem.") if cfg.v6_bank_readout_mode != "direct": raise RuntimeError( "--v6-specialist-from-b0 requires " "--v6-bank-readout-mode direct.") print(" [V6 specialist] B0 bootstrap mode ON " "freeze=trunk source=q_mem readout=direct noinject=True") train_loader, val_loader, test_loader = make_loaders( _train_ds, _val_ds, _test_ds, cfg.batch_size, cfg.tpu, shuffle_seed=cfg.seed) print(f"\n{'='*60}") print(f" Model: {tag} (batch_size={cfg.batch_size})") print(f"{'='*60}") model = build_model( tag, cfg.vocab_size, cfg.max_seq, v6_mixture=bool(getattr(cfg, "v6_mixture", False)), v6_lambda_init=float(getattr(cfg, "v6_lambda_init", -4.0)), v6_lambda_a_init=float(getattr(cfg, "v6_lambda_a_init", 1.0)), v6_aux_weight=float(getattr(cfg, "v6_aux_weight", 0.0)), v6_margin_target=float(getattr(cfg, "v6_margin_target", 1.0)), v6_mix_weight=float(getattr(cfg, "v6_mix_weight", 0.5)), v6_gate_weight=float(getattr(cfg, "v6_gate_weight", 0.01)), v6_lambda_floor=float(getattr(cfg, "v6_lambda_floor", 0.05)), v6_lambda_head=bool(getattr(cfg, "v6_lambda_head", False)), v6_lambda_head_hidden=int(getattr(cfg, "v6_lambda_head_hidden", 160)), v6_lambda_head_bias_init=float(getattr(cfg, "v6_lambda_head_bias_init", -7.0)), v6_bg_weight=float(getattr(cfg, "v6_bg_weight", 1.0)), v6_bg_target=float(getattr(cfg, "v6_bg_target", 0.01)), v6_bce_objective=bool(getattr(cfg, "v6_bce_objective", False)), v6_sel_weight=float(getattr(cfg, "v6_sel_weight", 1.0)), v6_lambda_head_w_init_std=float(getattr(cfg, "v6_lambda_head_w_init_std", 1e-3)), v6_wt_sparsity_weight=float(getattr(cfg, "v6_wt_sparsity_weight", 0.0)), v6_wt_sparsity_target=float(getattr(cfg, "v6_wt_sparsity_target", 0.0)), v6_mem_head_bank=bool(getattr(cfg, "v6_mem_head_bank", False)), v6_bank_ce_weight=float(getattr(cfg, "v6_bank_ce_weight", 0.0)), v6_bank_pair_ce_weight=float(getattr( cfg, "v6_bank_pair_ce_weight", 0.0)), v6_bank_query_source=str(getattr(cfg, "v6_bank_query_source", "h_d")), v6_bank_readout_mode=str(getattr(cfg, "v6_bank_readout_mode", "bank")), v6_source_adapter=bool(getattr(cfg, "v6_source_adapter", False)), v6_context_adapter=bool(getattr(cfg, "v6_context_adapter", False)), v6_specialist_noinject=bool(getattr(cfg, "v6_specialist_noinject", False)), ).to(device) n_params = count_params(model) print(f" Parameters: {n_params:,} ({n_params/1e6:.1f}M)") if device.type == 'cuda' and not args.eval_only and tag not in ( "Transformer-XL", "CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m", "CEDL-V3m2", "CEDL-V3m2-NLL", "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold", "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M", "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject", "CEDL-V4d-strong", "CEDL-V4d-strong-noinject", "CEDL-V4d-causal", "CEDL-V4d-causal-strong", "CEDL-V4d-causalz", "CEDL-V4d-combo", "CEDL-V4e", "CEDL-V4e-noinject", "CEDL-V5-LM", "CEDL-V6", "CEDL-V3lex", "CEDL-B0", "Mamba") \ or tag.startswith("CEDL-V4e-k"): try: model = torch.compile(model) print(" torch.compile: enabled") except Exception as e: print(f" torch.compile: skipped ({e})") if args.eval_only: ckpt = args.checkpoint or os.path.join(cfg.save_dir, f"{tag}_best.pt") if os.path.exists(ckpt): state = torch.load(ckpt, map_location='cpu', weights_only=True) if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict): wrapper_keys = sorted(k for k in state.keys() if k != "model") wrapper_msg = f" (wrapper keys={wrapper_keys})" if wrapper_keys else "" print(f" [eval] unwrapping checkpoint['model'] -> model weights{wrapper_msg}") state = state["model"] if any(k.startswith('_orig_mod.') for k in state): state = {k.replace('_orig_mod.', ''): v for k, v in state.items()} model.load_state_dict(state) print(f" Loaded checkpoint: {ckpt}") else: print(f" WARNING: no checkpoint found at {ckpt}, skipping {tag}") del model if torch.cuda.is_available(): torch.cuda.empty_cache() continue else: best_val_ppl = train(model, tag, cfg, device, train_loader, val_loader) best_path = os.path.join(cfg.save_dir, f"{tag}_best.pt") if os.path.exists(best_path): model.load_state_dict(torch.load(best_path, map_location='cpu', weights_only=True)) raw_m = model._orig_mod if hasattr(model, '_orig_mod') else model if isinstance(raw_m, CEDLTwoLoop100M): raw_m.feedback_alpha.fill_(1.0) test_ppl = evaluate(model, test_loader, device, max_steps=500, tpu=cfg.tpu) print(f"\n Test Perplexity: {test_ppl:.2f}") structured_reason = "skipped for 100M path; use PPL/downstream metrics" print(f" Structured benchmark: {structured_reason}.") struct_results = {} downstream = run_downstream_eval(model, tag, device) results_table[tag] = { "params": n_params, "test_ppl": test_ppl, "structured": struct_results, "structured_reason": structured_reason, "downstream": downstream, } del model if torch.cuda.is_available(): torch.cuda.empty_cache() print_final_results(results_table) if __name__ == "__main__": main()