ebm-splats 2.0.0__py3-none-any.whl

ebm/__init__.py

@@ -0,0 +1,8 @@
+"""
+EBM-Splats: Energy-Based Model with Gaussian Splats on a Hypersphere.
+
+Archived research project (2026). See README.md for details.
+"""
+
+__version__ = "2.0.0"
+__status__ = "archived"

ebm/config.py

@@ -0,0 +1,109 @@
+"""
+EBM Configuration — V2 Specification Compliant.
+
+Fixes applied:
+  - Removed duplicate init_alpha field
+  - Added V2 spec parameters (lambda context weights, adaptive noise, etc.)
+  - Explicit sign convention documentation
+  - Proper default values from V2 Especificacion Tecnica
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+
+@dataclass
+class EBMConfig:
+    """
+    Configuration for the EBM continuous language model.
+    
+    Sign Convention (V2 Spec §2.1):
+        p(x) = exp(-E(x)) / Z
+        score(x) = ∇_x log p(x) = -∇_x E(x)
+        Low energy = high probability (near splat centers)
+    """
+    
+    # ── Environment ──
+    device: str = "cpu"
+    
+    # ── Latent Space ──
+    latent_dim: int = 640
+    n_splats_init: int = 10000
+    max_splats: int = 100000
+    knn_k: int = 64
+    
+    # ── Splat Initialization (Phase 1 improvements) ──
+    vocab_embedding_path: Optional[str] = None
+    init_from_vocab_embeddings: bool = False
+    init_alpha: float = 1.0
+    init_kappa: float = 10.0
+    
+    # ── Splat Regularization ──
+    splat_temperature: float = 0.1
+    splat_weight_decay: float = 0.0
+    splat_weight_decay_start: float = 1.0
+    min_kappa: float = 1.0
+    max_kappa: float = 50.0
+    
+    # ── Training ──
+    learning_rate: float = 1e-3
+    reg_weight: float = 0.01
+    grad_clip: float = 1.0
+    temperature: float = 0.1
+    
+    # ── Curriculum Learning ──
+    enable_curriculum_learning: bool = True
+    curriculum_epochs: int = 5
+    curriculum_target_splats: int = 50000
+    splat_convergence_threshold: float = 0.95
+    
+    # ── Noise Schedule (V2 §3.3 Adaptive Noise) ──
+    noise_levels: Tuple[float, ...] = (0.01, 0.05, 0.1, 0.2, 0.5)
+    sigma_base: float = 0.1
+    sigma_scale: float = 1.0
+    sigma_rho_ref: float = 1.0
+    
+    # ── Langevin Dynamics ──
+    langevin_steps: int = 200
+    langevin_dt: float = 0.001
+    langevin_gamma: float = 0.1
+    langevin_T: float = 1.0
+    
+    # ── SOC (Self-Organized Criticality) ──
+    soc_threshold: float = 0.8
+    soc_check_interval: int = 100
+    min_splat_distance: float = 0.1
+    
+    # ── Hierarchical Context (V2 §4) ──
+    context_local_window: int = 16       # 8-16 tokens
+    context_medium_window: int = 128     # 64-128 tokens
+    context_global_window: int = 512     # 512+ tokens
+    beta_local: float = 0.5              # Fast adaptation
+    beta_medium: float = 0.8             # Moderate
+    beta_global: float = 0.95            # Slow, stable
+    lambda_context_local: float = 1.0
+    lambda_context_medium: float = 0.5
+    lambda_context_global: float = 0.2
+    
+    # ── Decoder (V2 §5) ──
+    vocab_size: int = 50257  # GPT-2 vocab
+    moe_experts: int = 4
+    moe_active: int = 2
+    hidden_dim: int = 1024
+    
+    # ── Collapse Regularization (V2 §2.4) ──
+    lambda_reg: float = 0.01
+    theta_threshold: float = 0.9  # Min angular separation between splat centers
+    
+    # ── Energy Weights ──
+    lambda_geom: float = 0.01
+    lambda_comp: float = 0.05
+    
+    # ── EMA ──
+    ema_decay: float = 0.999
+    
+    # ── Error Recovery (V2 §7) ──
+    energy_stagnation_window: int = 5
+    energy_stagnation_epsilon: float = 1e-4
+    gradient_magnitude_threshold: float = 1e-5
+    max_recovery_attempts: int = 3

ebm/context.py

@@ -0,0 +1,117 @@
+"""
+Hierarchical Context System — V2 §4.
+
+Three-level context for long-range dependencies:
+  - Local:  last 8-16 tokens  (β≈0.5, fast adaptation)
+  - Medium: last 64-128 tokens (β≈0.8, moderate)
+  - Global: full sequence       (β≈0.95, slow, stable)
+
+Integration:
+  E_trans(x_t) = -Σ_{l∈{local,medium,global}} λ_l · (x_t · c_l)
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict
+
+from .config import EBMConfig
+from .geometry import normalize_sphere
+
+
+class HierarchicalContext(nn.Module):
+    """Three-level context system for EBM language model."""
+
+    def __init__(self, config: EBMConfig):
+        super().__init__()
+        self.config = config
+        self.latent_dim = config.latent_dim
+
+        # Persistent context vectors (on hypersphere)
+        self.register_buffer('c_local', torch.zeros(1, config.latent_dim))
+        self.register_buffer('c_medium', torch.zeros(1, config.latent_dim))
+        self.register_buffer('c_global', torch.zeros(1, config.latent_dim))
+
+        self.beta_local = config.beta_local
+        self.beta_medium = config.beta_medium
+        self.beta_global = config.beta_global
+
+        # Update counters
+        self.register_buffer('_step', torch.tensor(0, dtype=torch.long))
+        self.register_buffer('_local_count', torch.tensor(0, dtype=torch.long))
+        self.register_buffer('_medium_count', torch.tensor(0, dtype=torch.long))
+
+        # Buffers stored as lists to avoid device/shape issues
+        self._local_buffer = []
+        self._medium_buffer = []
+
+    def reset(self, batch_size: int, device: torch.device):
+        """Reset context for a new sequence."""
+        self.c_local = normalize_sphere(torch.randn(batch_size, self.latent_dim, device=device))
+        self.c_medium = normalize_sphere(torch.randn(batch_size, self.latent_dim, device=device))
+        self.c_global = normalize_sphere(torch.randn(batch_size, self.latent_dim, device=device))
+        self._step.zero_()
+        self._local_count.zero_()
+        self._medium_count.zero_()
+        self._local_buffer = []
+        self._medium_buffer = []
+
+    @torch.no_grad()
+    def update(self, x: torch.Tensor):
+        """Update context with new token embedding x: [B, D]."""
+        B = x.size(0)
+        device = x.device
+        step = self._step.item()
+
+        # Ensure batch/device consistency
+        if self.c_local.size(0) != B or self.c_local.device != device:
+            self.c_local = normalize_sphere(torch.randn(B, self.latent_dim, device=device))
+            self.c_medium = normalize_sphere(torch.randn(B, self.latent_dim, device=device))
+            self.c_global = normalize_sphere(torch.randn(B, self.latent_dim, device=device))
+
+        # ── Local: every step ──
+        self.c_local = normalize_sphere(
+            self.beta_local * self.c_local.to(device) + (1 - self.beta_local) * x
+        )
+
+        # Store batch-averaged embedding for medium context
+        self._local_buffer.append(x.mean(dim=0).detach().cpu())
+        if len(self._local_buffer) > self.config.context_local_window:
+            self._local_buffer.pop(0)
+        self._local_count.add_(1)
+
+        # ── Medium: every 4 steps ──
+        if step > 0 and step % 4 == 0 and len(self._local_buffer) > 0:
+            local_avg = torch.stack(self._local_buffer).mean(dim=0).to(device)
+            self.c_medium = normalize_sphere(
+                self.beta_medium * self.c_medium.to(device) + (1 - self.beta_medium) * local_avg.unsqueeze(0).expand(B, -1)
+            )
+            self._medium_buffer.append(self.c_medium.mean(dim=0).detach().cpu())
+            if len(self._medium_buffer) > self.config.context_medium_window:
+                self._medium_buffer.pop(0)
+            self._medium_count.add_(1)
+
+        # ── Global: every 16 steps ──
+        if step > 0 and step % 16 == 0 and len(self._medium_buffer) > 0:
+            medium_avg = torch.stack(self._medium_buffer).mean(dim=0).to(device)
+            self.c_global = normalize_sphere(
+                self.beta_global * self.c_global.to(device) + (1 - self.beta_global) * medium_avg.unsqueeze(0).expand(B, -1)
+            )
+
+        self._step.add_(1)
+
+    def get_context(self) -> Dict[str, torch.Tensor]:
+        """Return dict of context vectors."""
+        return {
+            'local': self.c_local,
+            'medium': self.c_medium,
+            'global': self.c_global,
+        }
+
+    def get_combined_context(self) -> torch.Tensor:
+        """Weighted average context."""
+        c = (
+            self.config.lambda_context_local * self.c_local +
+            self.config.lambda_context_medium * self.c_medium +
+            self.config.lambda_context_global * self.c_global
+        )
+        return normalize_sphere(c)

ebm/cuda/__init__.py

@@ -0,0 +1 @@
+"""CUDA-accelerated energy computation (requires PyTorch with CUDA)."""

ebm/cuda/energy.py

@@ -0,0 +1,102 @@
+"""
+CUDA-native Energy Function — V2 Compliant.
+
+Same as energy.py but uses CUDA AMP-safe autograd instead of manual gradients.
+The V2 manual gradient approach in energy.py is preferred for training stability.
+This module serves as CUDA fallback when V2 manual gradients aren't needed.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..config import EBMConfig
+from ..geometry import project_to_tangent
+
+
+class EnergyFunctionCUDA(nn.Module):
+    """CUDA-native energy with V2 sign conventions."""
+
+    def __init__(self, config: EBMConfig, splat_store):
+        super().__init__()
+        self.config = config
+        self.splats = splat_store
+        self.W_comp = nn.Linear(3, 1)
+
+    def compute_splat_energy(self, x: torch.Tensor) -> torch.Tensor:
+        original_shape = x.shape[:-1]
+        if x.dim() > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        neighbors_mu, neighbors_alpha, neighbors_kappa = self.splats.find_neighbors(x, self.config.knn_k)
+
+        dot_products = torch.bmm(neighbors_mu, x.unsqueeze(-1)).squeeze(-1)
+
+        # Importance weighting (V2 §2.2)
+        importance = neighbors_kappa.clamp(min=1e-4)
+        weights = importance / importance.sum(dim=-1, keepdim=True)
+
+        exponent = neighbors_alpha * (dot_products - 1.0) / self.config.temperature
+        weighted_exponent = exponent + torch.log(weights.clamp(min=1e-8))
+
+        energy = -torch.logsumexp(weighted_exponent, dim=-1)
+        return energy.view(original_shape)
+
+    def compute_geom_energy(self, x: torch.Tensor) -> torch.Tensor:
+        original_shape = x.shape[:-1]
+        if x.dim() > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        if x.size(0) < 2:
+            return torch.tensor(0.0, device=x.device)
+
+        batch_sims = torch.mm(x, x.T)
+        mask = ~torch.eye(x.size(0), dtype=torch.bool, device=x.device)
+        off_diag = batch_sims[mask]
+        spread_energy = -torch.log(1.0 - off_diag.clamp(max=1.0 - 1e-4) + 1e-4).mean()
+
+        return spread_energy
+
+    def compute_comp_energy(self, x: torch.Tensor) -> torch.Tensor:
+        original_shape = x.shape[:-1]
+        if x.dim() > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        neighbors_mu, _, _ = self.splats.find_neighbors(x, 2)
+        u = torch.bmm(neighbors_mu[:, 0:1, :], x.unsqueeze(-1)).squeeze(-1)
+        v = torch.bmm(neighbors_mu[:, 1:2, :], x.unsqueeze(-1)).squeeze(-1)
+        uv_concat = torch.cat([u, v, u * v], dim=-1)
+        comp_energy = torch.sigmoid(self.W_comp(uv_concat)).squeeze(-1)
+        return comp_energy.view(original_shape)
+
+    def compute_context_energy(self, x: torch.Tensor, context_vecs: dict) -> torch.Tensor:
+        energy = torch.tensor(0.0, device=x.device)
+        for level, vec in context_vecs.items():
+            lam = getattr(self.config, f'lambda_context_{level}', 0.0)
+            energy = energy - lam * (x * vec).sum(dim=-1)
+        return energy
+
+    def forward(self, x: torch.Tensor, context_vecs: dict = None) -> torch.Tensor:
+        e_splat = self.compute_splat_energy(x)
+        e_geom = self.compute_geom_energy(x) * self.config.lambda_geom
+        e_comp = self.compute_comp_energy(x) * self.config.lambda_comp
+        e_total = e_splat + e_geom + e_comp
+
+        if context_vecs is not None:
+            e_trans = self.compute_context_energy(x, context_vecs)
+            e_total = e_total + e_trans
+
+        return e_total
+
+    def compute_score(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Riemannian score via AMP-safe autograd.
+        For training stability, prefer energy.py manual gradient approach.
+        """
+        with torch.amp.autocast('cuda', enabled=False):
+            x_f32 = x.float().detach().requires_grad_(True)
+            energy = self.forward(x_f32.float())
+            grad_e = torch.autograd.grad(energy.sum(), x_f32, create_graph=False)[0]
+            grad_riemann = project_to_tangent(x_f32, grad_e)
+
+        return -grad_riemann

ebm/data.py

@@ -0,0 +1,67 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class TextDataset(Dataset):
+    def __init__(self, tokenized_data, seq_len: int = 16):
+        """
+        Dynamically slices a continuous stream of tokens into chunks of `seq_len`.
+        """
+        self.seq_len = seq_len
+        self.tokens = tokenized_data
+        self.total_chunks = len(self.tokens) // seq_len
+
+    def __len__(self):
+        return self.total_chunks
+
+    def __getitem__(self, idx):
+        start = idx * self.seq_len
+        end = start + self.seq_len
+        return torch.tensor(self.tokens[start:end], dtype=torch.long)
+
+def get_dataloader(tokenizer_name: str = "gpt2", dataset_name: str = "wikitext", 
+                   config_name: str = "wikitext-2-raw-v1", split: str = "train",
+                   batch_size: int = 8, seq_len: int = 16, max_samples: int = 5000):
+    """
+    Downloads, tokenizes, and maps text blocks to PyTorch DataLoader structs.
+    """
+    logger.info(f"Loading '{dataset_name}' dataset ({split})...")
+    
+    try:
+        raw_datasets = load_dataset(dataset_name, config_name)
+        data_split = raw_datasets[split]
+    except Exception as e:
+        logger.error(f"Failed to load dataset: {e}")
+        raise
+        
+    logger.info(f"Loading tokenizer '{tokenizer_name}'...")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    
+    # Optional pad token setting if required by specific batching, but we utilize 
+    # exact chunks here.
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        
+    logger.info("Tokenizing active text samples...")
+    # Concatenate texts dropping empty lines
+    text_corpus = " ".join([row["text"] for row in data_split.select(range(min(len(data_split), max_samples)))])
+    
+    # Process without truncation targeting a continuous token stream
+    tokenized = tokenizer.encode(text_corpus, add_special_tokens=False)
+    
+    logger.info(f"Generated {len(tokenized)} tokens.")
+    
+    # Construct sequence slicer dataset
+    dataset = TextDataset(tokenized, seq_len=seq_len)
+    
+    # Drop_last=True ensures all batches are uniformly sized (B, seq_len)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
+    
+    logger.info(f"Dataloader yielding {len(dataloader)} batches of dimension [{batch_size}, {seq_len}]")
+    
+    return dataloader, tokenizer

ebm/data_loader.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""
+Optimized dataset loader for EBM training. Uses local text files
+instead of downloading from HuggingFace every time.
+
+Supported datasets:
+  - TinyStories: D:/datasets/ebm/tinystories_train.txt (1.9GB, 2.1M stories)
+  - TinyStories val: D:/datasets/ebm/tinystories_val.txt (19MB, 22K stories)
+  - WikiText-103: cached via HuggingFace (fallback)
+"""
+
+import os
+import logging
+import warnings
+
+# Suppress GPT-2 tokenizer warnings about sequence length
+warnings.filterwarnings("ignore", message="Token indices sequence length")
+
+logger = logging.getLogger(__name__)
+
+# Lazy imports for heavy dependencies
+def _get_torch_utils():
+    """Lazy import torch utilities."""
+    import torch
+    from torch.utils.data import Dataset, DataLoader
+    return torch, Dataset, DataLoader
+
+def _get_transformers():
+    """Lazy import transformers."""
+    from transformers import AutoTokenizer
+    return AutoTokenizer
+
+# Default paths
+DATASETS_DIR = "D:/datasets/ebm"
+TINYSTORIES_TRAIN = os.path.join(DATASETS_DIR, "tinystories_train.txt")
+TINYSTORIES_VAL = os.path.join(DATASETS_DIR, "tinystories_val.txt")
+
+
+class TextFileDataset:
+    """Memory-mapped text file sliced into token chunks."""
+
+    def __init__(self, token_ids: list, seq_len: int):
+        self.seq_len = seq_len
+        self.tokens = token_ids
+        self.total_chunks = len(self.tokens) // seq_len
+        # Register with torch Dataset class
+        torch, Dataset, _ = _get_torch_utils()
+        self._Dataset = Dataset
+
+    def __len__(self):
+        return self.total_chunks
+
+    def __getitem__(self, idx):
+        torch, _, _ = _get_torch_utils()
+        start = idx * self.seq_len
+        end = start + self.seq_len
+        return torch.tensor(self.tokens[start:end], dtype=torch.long)
+
+
+def tokenize_file(filepath: str, tokenizer, max_chars: int = None, chunk_size: int = 50_000_000) -> list:
+    """Tokenize a text file into a flat list of token IDs, reading in chunks to avoid OOM."""
+    logger.info(f"Reading {filepath}...")
+    
+    all_token_ids = []
+    chars_read = 0
+    
+    with open(filepath, 'r', encoding='utf-8') as f:
+        while True:
+            if max_chars and chars_read >= max_chars:
+                break
+            read_size = min(chunk_size, max_chars - chars_read) if max_chars else chunk_size
+            text = f.read(read_size)
+            if not text:
+                break
+            chars_read += len(text)
+            
+            # Tokenize this chunk
+            chunk_ids = tokenizer.encode(text, add_special_tokens=False)
+            all_token_ids.extend(chunk_ids)
+            logger.info(f"  {chars_read:,} chars → {len(all_token_ids):,} tokens")
+    
+    logger.info(f"  Total: {chars_read:,} characters, {len(all_token_ids):,} tokens")
+    return all_token_ids
+
+
+def get_tinystories_dataloader(
+    seq_len: int = 64,
+    batch_size: int = 64,
+    split: str = "train",
+    max_chars: int = None,
+    tokenizer_name: str = "gpt2",
+    shuffle: bool = True,
+):
+    """Load TinyStories dataset from local D: drive."""
+
+    if split == "train":
+        filepath = TINYSTORIES_TRAIN
+    elif split == "val":
+        filepath = TINYSTORIES_VAL
+    else:
+        raise ValueError(f"Unknown split: {split}")
+
+    if not os.path.exists(filepath):
+        raise FileNotFoundError(f"Dataset not found: {filepath}. Run download script first.")
+
+    # Lazy load transformers and torch
+    AutoTokenizer = _get_transformers()
+    torch, _, DataLoader = _get_torch_utils()
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    token_ids = tokenize_file(filepath, tokenizer, max_chars)
+    dataset = TextFileDataset(token_ids, seq_len)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=True,
+                            num_workers=0, pin_memory=True)
+
+    logger.info(f"TinyStories {split}: {len(dataset):,} chunks of [{batch_size} x {seq_len}]")
+    return dataloader, tokenizer
+
+
+def get_wikitext_dataloader(
+    seq_len: int = 32,
+    batch_size: int = 64,
+    split: str = "train",
+    max_samples: int = 5000,
+    tokenizer_name: str = "gpt2",
+):
+    """Fallback: load WikiText-103 from HuggingFace cache."""
+    from datasets import load_dataset
+
+    AutoTokenizer = _get_transformers()
+    torch, _, DataLoader = _get_torch_utils()
+
+    raw = load_dataset("wikitext", "wikitext-103-raw-v1")
+    data = raw[split]
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    text_corpus = " ".join(
+        row["text"] for row in data.select(range(min(len(data), max_samples)))
+    )
+    token_ids = tokenizer.encode(text_corpus, add_special_tokens=False)
+
+    dataset = TextFileDataset(token_ids, seq_len)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True,
+                            num_workers=0, pin_memory=True)
+    return dataloader, tokenizer
+
+
+def get_dataloader(
+    dataset_name: str = "tinystories",
+    seq_len: int = 64,
+    batch_size: int = 64,
+    split: str = "train",
+    max_chars: int = None,
+    tokenizer_name: str = "gpt2",
+    **kwargs,
+):
+    """Unified dataloader factory."""
+    if dataset_name == "tinystories":
+        return get_tinystories_dataloader(
+            seq_len=seq_len, batch_size=batch_size, split=split,
+            max_chars=max_chars, tokenizer_name=tokenizer_name,
+        )
+    elif dataset_name == "wikitext":
+        return get_wikitext_dataloader(
+            seq_len=seq_len, batch_size=batch_size, split=split,
+            max_samples=kwargs.get("max_samples", 5000),
+            tokenizer_name=tokenizer_name,
+        )
+    else:
+        raise ValueError(f"Unknown dataset: {dataset_name}")
+
+
+if __name__ == "__main__":
+    # Quick test
+    logging.basicConfig(level=logging.INFO, format='%(message)s')
+    
+    print("=== TinyStories (100K chars test) ===")
+    dl, tok = get_dataloader("tinystories", seq_len=64, batch_size=8, max_chars=100_000)
+    batch = next(iter(dl))
+    print(f"Batch shape: {batch.shape}")
+    print(f"Sample tokens (first 10): {batch[0][:10].tolist()}")
+    print(f"Decoded: {tok.decode(batch[0][:10])}")
+    
+    print("\n=== WikiText fallback ===")
+    dl2, tok2 = get_dataloader("wikitext", seq_len=32, batch_size=8, max_samples=100)
+    batch2 = next(iter(dl2))
+    print(f"Batch shape: {batch2.shape}")

ebm/decoder.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+from .config import EBMConfig
+
+class MoELayer(nn.Module):
+    def __init__(self, config: EBMConfig, in_features: int):
+        super().__init__()
+        self.num_experts = config.moe_experts
+        self.k = config.moe_active
+        self.hidden_dim = config.hidden_dim
+        
+        # Router
+        self.router = nn.Linear(in_features, self.num_experts)
+        
+        # Experts (simplified as parallel linear layers for CPU efficiency)
+        self.experts_w1 = nn.Parameter(torch.randn(self.num_experts, in_features, self.hidden_dim) / in_features**0.5)
+        self.experts_b1 = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+        
+        self.experts_w2 = nn.Parameter(torch.randn(self.num_experts, self.hidden_dim, in_features) / self.hidden_dim**0.5)
+        self.experts_b2 = nn.Parameter(torch.zeros(self.num_experts, in_features))
+        
+        self.activation = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x is [B, D]
+        B, D = x.shape
+        router_logits = self.router(x)
+        routing_weights = torch.softmax(router_logits, dim=-1)
+
+        # Top-K routing
+        top_k_weights, top_k_indices = torch.topk(routing_weights, self.k, dim=-1)
+        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
+
+        # Vectorized MoE: compute all experts, then gather top-K
+        # x: [B, D] → expand to [B, 1, D] → [B, num_experts, D]
+        x_expanded = x.unsqueeze(1).expand(-1, self.num_experts, -1)  # [B, E, D]
+
+        # Expert FFN (batched over all experts)
+        # experts_w1: [E, D, H], x_expanded: [B, E, D] → h: [B, E, H]
+        h = torch.einsum('bed,edh->beh', x_expanded, self.experts_w1) + self.experts_b1.unsqueeze(0)
+        h = self.activation(h)
+        # experts_w2: [E, H, D], h: [B, E, H] → o: [B, E, D]
+        o = torch.einsum('beh,ehd->bed', h, self.experts_w2) + self.experts_b2.unsqueeze(0)
+
+        # Gather top-K expert outputs: [B, E, D] → [B, K, D]
+        top_k_idx = top_k_indices.unsqueeze(-1).expand(-1, -1, D)  # [B, K, D]
+        selected = torch.gather(o, 1, top_k_idx)  # [B, K, D]
+
+        # Weighted sum: [B, K, 1] * [B, K, D] → [B, D]
+        out = (top_k_weights.unsqueeze(-1) * selected).sum(dim=1)
+        return out
+
+class EBMDecoder(nn.Module):
+    def __init__(self, config: EBMConfig):
+        super().__init__()
+        
+        # Context concat -> [X; c_total] means dim is 2 * latent_dim
+        in_dim = config.latent_dim * 2 
+        
+        self.moe = MoELayer(config, in_dim)
+        self.output_layer = nn.Linear(in_dim, config.vocab_size)
+        
+    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        """
+        x: Current latent state [B, D]
+        context: Context hierarchical vector [B, D]
+        """
+        combined = torch.cat([x, context], dim=-1)
+        moe_out = self.moe(combined)
+        logits = self.output_layer(moe_out)
+        return logits