FIX: Use vectorized span masking in notebook (no Python loop over batch)"

MuseMorphic_Training.ipynb

@@ -20,26 +20,13 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "# 🎵 MuseMorphic: Lightweight MIDI Generator\n",
+        "# 🎵 MuseMorphic: Lightweight MIDI Generator (v0.2)\n",
         "\n",
         "**A novel consumer-grade architecture for controllable, infinite-length MIDI generation.**\n",
         "\n",
-        "Key features:\n",
-        "- **~33M parameters** — trains on free Colab T4, inference <1GB VRAM\n",
-        "- **O(n) complexity** — Mamba SSM backbone, no quadratic attention\n",
-        "- **Two-stage hierarchical** — PhraseVAE (compress) + LatentMamba (generate)\n",
-        "- **Music-native** — FME embeddings with harmonic awareness\n",
-        "- **Controllable** — tempo, key, density, style conditioning\n",
-        "- **Infinite generation** — fixed-size recurrent state, no memory growth\n",
-        "- **Training-stable by design** — σReparam + ZClip + Pre-LN + BF16\n",
+        "v0.2 Performance fixes: weight_norm (not spectral_norm), chunked SSM scan, vectorized masking.\n",
         "\n",
-        "📄 [Architecture Paper/README](https://huggingface.co/asdf98/MuseMorphic)\n",
-        "\n",
-        "---\n",
-        "\n",
-        "## Setup\n",
-        "\n",
-        "Run this cell first to install all dependencies."
+        "📄 [Architecture Details](https://huggingface.co/asdf98/MuseMorphic)"
       ]
     },
     {
@@ -52,15 +39,14 @@
         "# 1. Install Dependencies\n",
         "# ============================================================\n",
         "!pip install -q torch torchvision torchaudio\n",
-        "!pip install -q einops datasets pretty_midi midiutil\n",
-        "!pip install -q huggingface_hub\n",
+        "!pip install -q einops datasets pretty_midi midiutil huggingface_hub\n",
         "\n",
-        "# Clone MuseMorphic repo\n",
-        "!git clone https://huggingface.co/asdf98/MuseMorphic /content/MuseMorphic 2>/dev/null || (cd /content/MuseMorphic && git pull)\n",
+        "# Clone MuseMorphic repo (fresh pull to get v0.2 fixes)\n",
+        "!rm -rf /content/MuseMorphic\n",
+        "!git clone https://huggingface.co/asdf98/MuseMorphic /content/MuseMorphic\n",
         "\n",
         "import sys\n",
         "sys.path.insert(0, '/content/MuseMorphic/musemorphic')\n",
-        "\n",
         "print('✅ Dependencies installed!')"
       ]
     },
@@ -74,36 +60,20 @@
         "# 2. Check GPU & Hardware\n",
         "# ============================================================\n",
         "import torch\n",
-        "import os\n",
-        "\n",
-        "print(f'PyTorch version: {torch.__version__}')\n",
-        "print(f'CUDA available: {torch.cuda.is_available()}')\n",
+        "print(f'PyTorch: {torch.__version__}')\n",
+        "print(f'CUDA: {torch.cuda.is_available()}')\n",
         "if torch.cuda.is_available():\n",
         "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
         "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')\n",
-        "    print(f'BF16 support: {torch.cuda.is_bf16_supported()}')\n",
-        "else:\n",
-        "    print('⚠️ No GPU detected. Training will be slow but functional on CPU.')\n",
+        "    print(f'BF16: {torch.cuda.is_bf16_supported()}')\n",
         "\n",
-        "# Auto-detect best dtype\n",
         "if torch.cuda.is_available() and torch.cuda.is_bf16_supported():\n",
-        "    DTYPE = 'bf16'\n",
-        "    print('\\n✅ Using BFloat16 (optimal: no loss scaling needed)')\n",
+        "    DTYPE = 'bf16'; amp_dtype = torch.bfloat16\n",
         "elif torch.cuda.is_available():\n",
-        "    DTYPE = 'fp16'\n",
-        "    print('\\n⚠️ Using Float16 (T4 GPU — BF16 not supported, using FP16 with gradient scaling)')\n",
+        "    DTYPE = 'fp16'; amp_dtype = torch.float16\n",
         "else:\n",
-        "    DTYPE = 'fp32'\n",
-        "    print('\\n📋 Using Float32 (CPU mode)')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Configuration\n",
-        "\n",
-        "Adjust these settings based on your GPU and desired output quality."
+        "    DTYPE = 'fp32'; amp_dtype = torch.float32\n",
+        "print(f'\\nUsing: {DTYPE}')"
       ]
     },
     {
@@ -115,35 +85,17 @@
         "# ============================================================\n",
         "# 3. Configuration\n",
         "# ============================================================\n",
-        "\n",
-        "# ---- Model Size Presets ----\n",
-        "# 'tiny'  : ~8M params  — Fast experiments, lower quality\n",
-        "# 'small' : ~33M params — Default, good quality (recommended for Colab T4)\n",
-        "# 'medium': ~65M params — Better quality, needs more VRAM\n",
-        "\n",
         "MODEL_SIZE = 'small'  # @param ['tiny', 'small', 'medium']\n",
-        "\n",
-        "# ---- Dataset ----\n",
-        "# 'auto'     : Auto-select best available\n",
-        "# 'maestro'  : Classical piano (MAESTRO)\n",
-        "# 'synthetic': Generated data (for testing)\n",
-        "DATASET = 'auto'  # @param ['auto', 'maestro', 'synthetic']\n",
-        "MAX_PIECES = 500  # @param {type: 'integer'}\n",
-        "\n",
-        "# ---- Training ----\n",
-        "VAE_EPOCHS = 15      # @param {type: 'integer'}\n",
-        "MAMBA_EPOCHS = 30    # @param {type: 'integer'}\n",
-        "BATCH_SIZE = 32      # @param {type: 'integer'}\n",
-        "LEARNING_RATE = 3e-4 # @param {type: 'number'}\n",
-        "\n",
-        "# ---- Output ----\n",
-        "OUTPUT_DIR = '/content/checkpoints'  # @param {type: 'string'}\n",
-        "PUSH_TO_HUB = False   # @param {type: 'boolean'}\n",
-        "HUB_MODEL_ID = ''     # @param {type: 'string'}\n",
-        "\n",
-        "print(f'Model size: {MODEL_SIZE}')\n",
-        "print(f'Dataset: {DATASET}')\n",
-        "print(f'Training: VAE {VAE_EPOCHS}ep + Mamba {MAMBA_EPOCHS}ep, batch={BATCH_SIZE}')"
+        "DATASET = 'auto'      # @param ['auto', 'maestro', 'synthetic']\n",
+        "MAX_PIECES = 500      # @param {type: 'integer'}\n",
+        "VAE_EPOCHS = 15       # @param {type: 'integer'}\n",
+        "MAMBA_EPOCHS = 30     # @param {type: 'integer'}\n",
+        "BATCH_SIZE = 32       # @param {type: 'integer'}\n",
+        "LEARNING_RATE = 3e-4  # @param {type: 'number'}\n",
+        "OUTPUT_DIR = '/content/checkpoints'\n",
+        "PUSH_TO_HUB = False\n",
+        "HUB_MODEL_ID = ''\n",
+        "print(f'Config: {MODEL_SIZE} model, {DATASET} data, VAE {VAE_EPOCHS}ep + Mamba {MAMBA_EPOCHS}ep')"
       ]
     },
     {
@@ -153,63 +105,25 @@
       "outputs": [],
       "source": [
         "# ============================================================\n",
-        "# 4. Build Model Configuration\n",
+        "# 4. Build Model\n",
         "# ============================================================\n",
         "from model import MuseMorphicConfig, MuseMorphic, model_summary\n",
         "\n",
-        "# Model size presets\n",
         "SIZE_CONFIGS = {\n",
-        "    'tiny': MuseMorphicConfig(\n",
-        "        d_model=128,\n",
-        "        vae_encoder_layers=2,\n",
-        "        vae_decoder_layers=2,\n",
-        "        vae_n_heads=4,\n",
-        "        vae_d_ff=256,\n",
-        "        latent_dim=32,\n",
-        "        mamba_d_model=128,\n",
-        "        mamba_n_layers=4,\n",
-        "        mamba_d_state=8,\n",
-        "        mamba_expand=2,\n",
-        "    ),\n",
-        "    'small': MuseMorphicConfig(\n",
-        "        d_model=256,\n",
-        "        vae_encoder_layers=3,\n",
-        "        vae_decoder_layers=3,\n",
-        "        vae_n_heads=4,\n",
-        "        vae_d_ff=512,\n",
-        "        latent_dim=64,\n",
-        "        mamba_d_model=256,\n",
-        "        mamba_n_layers=8,\n",
-        "        mamba_d_state=16,\n",
-        "        mamba_expand=2,\n",
-        "    ),\n",
-        "    'medium': MuseMorphicConfig(\n",
-        "        d_model=384,\n",
-        "        vae_encoder_layers=4,\n",
-        "        vae_decoder_layers=4,\n",
-        "        vae_n_heads=6,\n",
-        "        vae_d_ff=768,\n",
-        "        latent_dim=96,\n",
-        "        mamba_d_model=384,\n",
-        "        mamba_n_layers=12,\n",
-        "        mamba_d_state=16,\n",
-        "        mamba_expand=2,\n",
-        "    ),\n",
+        "    'tiny': MuseMorphicConfig(d_model=128, vae_encoder_layers=2, vae_decoder_layers=2,\n",
+        "        vae_n_heads=4, vae_d_ff=256, latent_dim=32, mamba_d_model=128,\n",
+        "        mamba_n_layers=4, mamba_d_state=8, mamba_expand=2),\n",
+        "    'small': MuseMorphicConfig(d_model=256, vae_encoder_layers=3, vae_decoder_layers=3,\n",
+        "        vae_n_heads=4, vae_d_ff=512, latent_dim=64, mamba_d_model=256,\n",
+        "        mamba_n_layers=8, mamba_d_state=16, mamba_expand=2),\n",
+        "    'medium': MuseMorphicConfig(d_model=384, vae_encoder_layers=4, vae_decoder_layers=4,\n",
+        "        vae_n_heads=6, vae_d_ff=768, latent_dim=96, mamba_d_model=384,\n",
+        "        mamba_n_layers=12, mamba_d_state=16, mamba_expand=2),\n",
         "}\n",
-        "\n",
         "config = SIZE_CONFIGS[MODEL_SIZE]\n",
         "model = model_summary(config)"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Data Preparation\n",
-        "\n",
-        "Automatically downloads and preprocesses MIDI data."
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -221,60 +135,27 @@
         "# ============================================================\n",
         "import logging\n",
         "logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')\n",
-        "\n",
-        "from data_pipeline import prepare_training_data, auto_select_dataset, load_dataset_notes\n",
-        "from data_pipeline import preprocess_dataset, _generate_synthetic_dataset\n",
+        "from data_pipeline import auto_select_dataset, load_dataset_notes, preprocess_dataset, _generate_synthetic_dataset\n",
         "from tokenizer import REMIPlusTokenizer\n",
         "\n",
-        "# Select dataset\n",
-        "if DATASET == 'auto':\n",
-        "    dataset_name = auto_select_dataset()\n",
-        "elif DATASET == 'maestro':\n",
-        "    dataset_name = 'maestro_v1_sustain'\n",
-        "elif DATASET == 'synthetic':\n",
-        "    dataset_name = None\n",
-        "else:\n",
-        "    dataset_name = DATASET\n",
-        "\n",
-        "# Load and preprocess\n",
         "tokenizer = REMIPlusTokenizer()\n",
+        "if DATASET == 'auto': dataset_name = auto_select_dataset()\n",
+        "elif DATASET == 'maestro': dataset_name = 'maestro_v1_sustain'\n",
+        "else: dataset_name = None\n",
         "\n",
-        "if dataset_name is not None:\n",
+        "if dataset_name:\n",
         "    try:\n",
         "        pieces = load_dataset_notes(dataset_name, max_pieces=MAX_PIECES)\n",
         "        print(f'✅ Loaded {len(pieces)} pieces from {dataset_name}')\n",
         "    except Exception as e:\n",
-        "        print(f'⚠️ Failed to load {dataset_name}: {e}')\n",
-        "        print('Falling back to synthetic data...')\n",
+        "        print(f'⚠️ {e}\\nFalling back to synthetic data...')\n",
         "        pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
         "else:\n",
         "    pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
         "    print(f'✅ Generated {len(pieces)} synthetic pieces')\n",
         "\n",
-        "# Preprocess\n",
         "phrases, controls = preprocess_dataset(pieces, tokenizer, max_phrase_len=config.vae_max_seq_len)\n",
-        "print(f'\\n📊 Dataset Summary:')\n",
-        "print(f'  Total phrases: {len(phrases)}')\n",
-        "print(f'  Avg phrase length: {sum(len(p) for p in phrases)/len(phrases):.1f} tokens')\n",
-        "print(f'  Vocab size: {tokenizer.vocab_size}')\n",
-        "print(f'  Sample phrase (first 10 tokens): {phrases[0][:10]}')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Training\n",
-        "\n",
-        "Two-stage training with curriculum:\n",
-        "\n",
-        "**Stage 1 — PhraseVAE** (compress phrases to latent vectors):\n",
-        "- 1a. Span-infilling pretraining (learn REMI grammar)\n",
-        "- 1b. Autoencoder (pure reconstruction, KL=0)\n",
-        "- 1c. VAE fine-tuning (KL weight = 0.01)\n",
-        "\n",
-        "**Stage 2 — LatentMamba** (generate latent sequences):\n",
-        "- Predict next phrase latent from history, O(n) complexity"
+        "print(f'📊 {len(phrases)} phrases, avg {sum(len(p) for p in phrases)/len(phrases):.1f} tokens')"
       ]
     },
     {
@@ -286,160 +167,95 @@
         "# ============================================================\n",
         "# 6. Training — Stage 1: PhraseVAE\n",
         "# ============================================================\n",
-        "import time\n",
-        "import random\n",
-        "import math\n",
-        "import numpy as np\n",
-        "import torch\n",
+        "import time, random, math, numpy as np\n",
         "import torch.nn.functional as F\n",
         "from torch.utils.data import DataLoader, Dataset\n",
-        "from model import PhraseVAE, ZClip\n",
+        "from model import PhraseVAE, ZClip, apply_span_mask_vectorized\n",
         "\n",
-        "# Set seed\n",
+        "# Seed\n",
         "SEED = 42\n",
-        "random.seed(SEED)\n",
-        "np.random.seed(SEED)\n",
-        "torch.manual_seed(SEED)\n",
-        "if torch.cuda.is_available():\n",
-        "    torch.cuda.manual_seed_all(SEED)\n",
+        "random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\n",
+        "if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)\n",
         "\n",
-        "# Device & dtype\n",
         "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
-        "if DTYPE == 'bf16' and torch.cuda.is_available() and torch.cuda.is_bf16_supported():\n",
-        "    amp_dtype = torch.bfloat16\n",
-        "elif DTYPE == 'fp16' and torch.cuda.is_available():\n",
-        "    amp_dtype = torch.float16\n",
-        "else:\n",
-        "    amp_dtype = torch.float32\n",
+        "use_scaler = (amp_dtype == torch.float16)\n",
+        "scaler = torch.amp.GradScaler() if use_scaler else None\n",
         "\n",
-        "# Dataset\n",
         "class PhraseDS(Dataset):\n",
         "    def __init__(self, phrases, max_len, pad_id=0):\n",
-        "        self.phrases = phrases\n",
-        "        self.max_len = max_len\n",
-        "        self.pad_id = pad_id\n",
-        "    def __len__(self):\n",
-        "        return len(self.phrases)\n",
+        "        self.phrases = phrases; self.max_len = max_len; self.pad_id = pad_id\n",
+        "    def __len__(self): return len(self.phrases)\n",
         "    def __getitem__(self, idx):\n",
         "        ids = self.phrases[idx][:self.max_len]\n",
-        "        padded = ids + [self.pad_id] * (self.max_len - len(ids))\n",
-        "        return torch.tensor(padded, dtype=torch.long)\n",
+        "        return torch.tensor(ids + [self.pad_id] * (self.max_len - len(ids)), dtype=torch.long)\n",
         "\n",
         "train_ds = PhraseDS(phrases, config.vae_max_seq_len)\n",
-        "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, \n",
+        "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,\n",
         "                          num_workers=2, pin_memory=True, drop_last=True)\n",
         "\n",
-        "# Create VAE\n",
         "vae = PhraseVAE(config).to(device)\n",
-        "vae_params = sum(p.numel() for p in vae.parameters())\n",
-        "print(f'PhraseVAE parameters: {vae_params:,} ({vae_params/1e6:.2f}M)')\n",
+        "print(f'PhraseVAE: {sum(p.numel() for p in vae.parameters()):,} params')\n",
         "\n",
-        "# Optimizer\n",
         "optimizer = torch.optim.AdamW(vae.parameters(), lr=LEARNING_RATE, weight_decay=0.01)\n",
         "scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=500, T_mult=2, eta_min=1e-6)\n",
         "zclip = ZClip(z_thresh=2.5)\n",
         "\n",
-        "# FP16 needs GradScaler, BF16 does not\n",
-        "use_scaler = (amp_dtype == torch.float16)\n",
-        "scaler = torch.amp.GradScaler() if use_scaler else None\n",
+        "# ---- Training Loop (uses vectorized span masking — NO Python loop over batch) ----\n",
+        "print('\\n' + '='*60 + '\\nStarting PhraseVAE Training\\n' + '='*60)\n",
         "\n",
-        "# Span masking helper\n",
-        "def apply_span_mask(token_ids, mask_prob=0.15, mask_id=3, span_len=3):\n",
-        "    masked = token_ids.clone()\n",
-        "    B, L = masked.shape\n",
-        "    for b in range(B):\n",
-        "        n_masks = max(1, int(L * mask_prob / span_len))\n",
-        "        for _ in range(n_masks):\n",
-        "            start = random.randint(1, max(1, L - span_len - 1))\n",
-        "            end = min(start + span_len, L)\n",
-        "            masked[b, start:end] = mask_id\n",
-        "    return masked\n",
-        "\n",
-        "# ---- Training Loop ----\n",
-        "print('\\n' + '='*60)\n",
-        "print('Starting PhraseVAE Training')\n",
-        "print('='*60)\n",
-        "\n",
-        "# Compute total epochs\n",
         "pretrain_epochs = max(1, VAE_EPOCHS // 5)\n",
         "ae_epochs = max(1, VAE_EPOCHS * 3 // 5)\n",
-        "vae_epochs = max(1, VAE_EPOCHS - pretrain_epochs - ae_epochs)\n",
-        "\n",
-        "stages = [\n",
-        "    ('1a-Pretrain', pretrain_epochs, 0.0, True),\n",
-        "    ('1b-AE', ae_epochs, 0.0, False),\n",
-        "    ('1c-VAE', vae_epochs, 0.01, False),\n",
-        "]\n",
+        "vae_epochs_count = max(1, VAE_EPOCHS - pretrain_epochs - ae_epochs)\n",
+        "stages = [('1a-Pretrain', pretrain_epochs, 0.0, True),\n",
+        "          ('1b-AE', ae_epochs, 0.0, False),\n",
+        "          ('1c-VAE', vae_epochs_count, 0.01, False)]\n",
         "\n",
         "global_step = 0\n",
         "history = {'loss': [], 'recon': [], 'kl': []}\n",
         "\n",
         "for stage_name, n_epochs, kl_weight, use_masking in stages:\n",
         "    print(f'\\n--- Stage {stage_name} ({n_epochs} epochs, KL={kl_weight}) ---')\n",
-        "    \n",
-        "    # Lower LR for VAE fine-tuning stage\n",
         "    if stage_name == '1c-VAE':\n",
-        "        for pg in optimizer.param_groups:\n",
-        "            pg['lr'] = LEARNING_RATE * 0.1\n",
-        "    \n",
+        "        for pg in optimizer.param_groups: pg['lr'] = LEARNING_RATE * 0.1\n",
+        "\n",
         "    for epoch in range(n_epochs):\n",
-        "        vae.train()\n",
-        "        epoch_loss = 0\n",
-        "        n_batches = 0\n",
-        "        t0 = time.time()\n",
-        "        \n",
+        "        vae.train(); epoch_loss = 0; n_batches = 0; t0 = time.time()\n",
         "        for batch in train_loader:\n",
         "            token_ids = batch.to(device)\n",
-        "            \n",
-        "            # Apply masking for pretraining\n",
-        "            if use_masking:\n",
-        "                input_ids = apply_span_mask(token_ids)\n",
-        "            else:\n",
-        "                input_ids = token_ids\n",
-        "            \n",
-        "            optimizer.zero_grad()\n",
-        "            \n",
+        "            # VECTORIZED masking — runs entirely on GPU, no Python loop\n",
+        "            input_ids = apply_span_mask_vectorized(token_ids, mask_id=config.mask_token_id) if use_masking else token_ids\n",
+        "\n",
+        "            optimizer.zero_grad(set_to_none=True)\n",
         "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
         "                outputs = vae(input_ids, target_tokens=token_ids, kl_weight=kl_weight)\n",
         "                loss = outputs['loss']\n",
-        "            \n",
-        "            # NaN check\n",
+        "\n",
         "            if torch.isnan(loss) or torch.isinf(loss):\n",
-        "                print(f'⚠️ NaN/Inf at step {global_step}! Skipping...')\n",
-        "                optimizer.zero_grad()\n",
-        "                continue\n",
-        "            \n",
+        "                print(f'⚠️ NaN at step {global_step}! Skipping...')\n",
+        "                optimizer.zero_grad(set_to_none=True); continue\n",
+        "\n",
         "            if use_scaler:\n",
         "                scaler.scale(loss).backward()\n",
-        "                scaler.unscale_(optimizer)\n",
-        "                zclip(vae)\n",
-        "                scaler.step(optimizer)\n",
-        "                scaler.update()\n",
+        "                scaler.unscale_(optimizer); zclip(vae)\n",
+        "                scaler.step(optimizer); scaler.update()\n",
         "            else:\n",
-        "                loss.backward()\n",
-        "                zclip(vae)\n",
-        "                optimizer.step()\n",
-        "            \n",
+        "                loss.backward(); zclip(vae); optimizer.step()\n",
         "            scheduler.step()\n",
-        "            \n",
-        "            epoch_loss += loss.item()\n",
-        "            n_batches += 1\n",
-        "            global_step += 1\n",
-        "            \n",
-        "            # Log\n",
+        "\n",
+        "            epoch_loss += loss.item(); n_batches += 1; global_step += 1\n",
         "            history['loss'].append(loss.item())\n",
         "            history['recon'].append(outputs['recon_loss'].item())\n",
         "            history['kl'].append(outputs['kl_loss'].item())\n",
-        "        \n",
+        "\n",
         "        elapsed = time.time() - t0\n",
         "        avg_loss = epoch_loss / max(n_batches, 1)\n",
-        "        lr = optimizer.param_groups[0]['lr']\n",
         "        print(f'  Epoch {epoch+1}/{n_epochs} | Loss: {avg_loss:.4f} | '\n",
         "              f'Recon: {outputs[\"recon_loss\"].item():.4f} | '\n",
         "              f'KL: {outputs[\"kl_loss\"].item():.4f} | '\n",
-        "              f'LR: {lr:.2e} | Time: {elapsed:.1f}s')\n",
+        "              f'LR: {optimizer.param_groups[0][\"lr\"]:.2e} | {elapsed:.1f}s | '\n",
+        "              f'{n_batches/elapsed:.1f} batch/s')\n",
         "\n",
-        "print(f'\\n✅ PhraseVAE training complete! ({global_step} total steps)')"
+        "print(f'\\n✅ PhraseVAE training complete! ({global_step} steps)')"
       ]
     },
     {
@@ -453,122 +269,77 @@
         "# ============================================================\n",
         "from model import LatentMamba\n",
         "\n",
-        "# Freeze VAE encoder\n",
         "vae.eval()\n",
-        "for p in vae.parameters():\n",
-        "    p.requires_grad = False\n",
+        "for p in vae.parameters(): p.requires_grad = False\n",
         "\n",
-        "# Encode all phrases to latent space\n",
         "print('Encoding phrases to latent space...')\n",
         "all_latents = []\n",
-        "encode_loader = DataLoader(train_ds, batch_size=64, shuffle=False, num_workers=2)\n",
-        "\n",
+        "encode_loader = DataLoader(train_ds, batch_size=128, shuffle=False, num_workers=2)\n",
         "with torch.no_grad():\n",
         "    for batch in encode_loader:\n",
-        "        token_ids = batch.to(device)\n",
         "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
-        "            z, _, _ = vae.encode(token_ids)\n",
+        "            z, _, _ = vae.encode(batch.to(device))\n",
         "        all_latents.append(z.cpu())\n",
-        "\n",
         "all_z = torch.cat(all_latents, dim=0)\n",
-        "print(f'Encoded {all_z.shape[0]} phrases to latent dim {all_z.shape[1]}')\n",
+        "print(f'Encoded {all_z.shape[0]} phrases → latent dim {all_z.shape[1]}')\n",
         "\n",
-        "# Create latent sequences (group phrases into chunks)\n",
         "SEQ_LEN = min(64, len(all_z) // 4)\n",
-        "latent_seqs = []\n",
-        "for i in range(0, len(all_z) - SEQ_LEN, SEQ_LEN // 2):  # 50% overlap\n",
-        "    latent_seqs.append(all_z[i:i+SEQ_LEN])\n",
-        "\n",
-        "print(f'Created {len(latent_seqs)} latent sequences of length {SEQ_LEN}')\n",
+        "latent_seqs = [all_z[i:i+SEQ_LEN] for i in range(0, len(all_z) - SEQ_LEN, SEQ_LEN // 2)]\n",
+        "print(f'{len(latent_seqs)} latent sequences of length {SEQ_LEN}')\n",
         "\n",
         "class LatentDS(Dataset):\n",
-        "    def __init__(self, seqs):\n",
-        "        self.seqs = seqs\n",
-        "    def __len__(self):\n",
-        "        return len(self.seqs)\n",
-        "    def __getitem__(self, idx):\n",
-        "        return self.seqs[idx]\n",
+        "    def __init__(self, seqs): self.seqs = seqs\n",
+        "    def __len__(self): return len(self.seqs)\n",
+        "    def __getitem__(self, idx): return self.seqs[idx]\n",
         "\n",
-        "latent_ds = LatentDS(latent_seqs)\n",
-        "latent_loader = DataLoader(latent_ds, batch_size=min(BATCH_SIZE, len(latent_seqs)),\n",
-        "                           shuffle=True, drop_last=True)\n",
+        "latent_loader = DataLoader(LatentDS(latent_seqs),\n",
+        "    batch_size=min(BATCH_SIZE, len(latent_seqs)), shuffle=True, drop_last=True)\n",
         "\n",
-        "# Create LatentMamba\n",
         "mamba = LatentMamba(config).to(device)\n",
-        "mamba_params = sum(p.numel() for p in mamba.parameters())\n",
-        "print(f'LatentMamba parameters: {mamba_params:,} ({mamba_params/1e6:.2f}M)')\n",
+        "print(f'LatentMamba: {sum(p.numel() for p in mamba.parameters()):,} params')\n",
         "\n",
-        "# Optimizer\n",
-        "mamba_optimizer = torch.optim.AdamW(mamba.parameters(), lr=LEARNING_RATE * 0.5, weight_decay=0.01)\n",
-        "mamba_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n",
-        "    mamba_optimizer, T_0=300, T_mult=2, eta_min=1e-6)\n",
+        "mamba_opt = torch.optim.AdamW(mamba.parameters(), lr=LEARNING_RATE*0.5, weight_decay=0.01)\n",
+        "mamba_sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(mamba_opt, T_0=300, T_mult=2, eta_min=1e-6)\n",
         "mamba_zclip = ZClip(z_thresh=2.5)\n",
-        "\n",
         "mamba_scaler = torch.amp.GradScaler() if use_scaler else None\n",
         "\n",
-        "# Training loop\n",
-        "print('\\n' + '='*60)\n",
-        "print('Starting LatentMamba Training')\n",
-        "print('='*60)\n",
-        "\n",
+        "print('\\n' + '='*60 + '\\nStarting LatentMamba Training\\n' + '='*60)\n",
         "mamba_history = {'mse': [], 'cos': []}\n",
         "\n",
         "for epoch in range(MAMBA_EPOCHS):\n",
-        "    mamba.train()\n",
-        "    epoch_loss = 0\n",
-        "    n_batches = 0\n",
-        "    t0 = time.time()\n",
-        "    \n",
+        "    mamba.train(); epoch_loss = 0; n_batches = 0; t0 = time.time()\n",
         "    for batch in latent_loader:\n",
         "        z_seq = batch.to(device)\n",
-        "        z_input = z_seq[:, :-1]\n",
-        "        z_target = z_seq[:, 1:]\n",
-        "        \n",
-        "        mamba_optimizer.zero_grad()\n",
-        "        \n",
+        "        z_input, z_target = z_seq[:, :-1], z_seq[:, 1:]\n",
+        "        mamba_opt.zero_grad(set_to_none=True)\n",
         "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
         "            z_pred = mamba(z_input)\n",
         "            mse_loss = F.mse_loss(z_pred, z_target)\n",
         "            cos_loss = 1.0 - F.cosine_similarity(\n",
         "                z_pred.reshape(-1, z_pred.shape[-1]),\n",
-        "                z_target.reshape(-1, z_target.shape[-1]), dim=-1\n",
-        "            ).mean()\n",
+        "                z_target.reshape(-1, z_target.shape[-1]), dim=-1).mean()\n",
         "            loss = mse_loss + 0.1 * cos_loss\n",
-        "        \n",
+        "\n",
         "        if torch.isnan(loss) or torch.isinf(loss):\n",
-        "            print(f'⚠️ NaN/Inf at epoch {epoch}! Skipping...')\n",
-        "            mamba_optimizer.zero_grad()\n",
-        "            continue\n",
-        "        \n",
+        "            mamba_opt.zero_grad(set_to_none=True); continue\n",
+        "\n",
         "        if use_scaler:\n",
         "            mamba_scaler.scale(loss).backward()\n",
-        "            mamba_scaler.unscale_(mamba_optimizer)\n",
-        "            mamba_zclip(mamba)\n",
-        "            mamba_scaler.step(mamba_optimizer)\n",
-        "            mamba_scaler.update()\n",
+        "            mamba_scaler.unscale_(mamba_opt); mamba_zclip(mamba)\n",
+        "            mamba_scaler.step(mamba_opt); mamba_scaler.update()\n",
         "        else:\n",
-        "            loss.backward()\n",
-        "            mamba_zclip(mamba)\n",
-        "            mamba_optimizer.step()\n",
-        "        \n",
-        "        mamba_scheduler.step()\n",
-        "        \n",
-        "        epoch_loss += loss.item()\n",
-        "        n_batches += 1\n",
-        "        \n",
+        "            loss.backward(); mamba_zclip(mamba); mamba_opt.step()\n",
+        "        mamba_sched.step()\n",
+        "        epoch_loss += loss.item(); n_batches += 1\n",
         "        mamba_history['mse'].append(mse_loss.item())\n",
         "        mamba_history['cos'].append(cos_loss.item())\n",
-        "    \n",
-        "    elapsed = time.time() - t0\n",
-        "    avg_loss = epoch_loss / max(n_batches, 1)\n",
-        "    lr = mamba_optimizer.param_groups[0]['lr']\n",
-        "    \n",
+        "\n",
         "    if (epoch + 1) % 5 == 0 or epoch == 0:\n",
-        "        print(f'  Epoch {epoch+1}/{MAMBA_EPOCHS} | Loss: {avg_loss:.6f} | '\n",
-        "              f'MSE: {mse_loss.item():.6f} | Cos: {cos_loss.item():.4f} | '\n",
-        "              f'LR: {lr:.2e} | Time: {elapsed:.1f}s')\n",
+        "        elapsed = time.time() - t0\n",
+        "        print(f'  Epoch {epoch+1}/{MAMBA_EPOCHS} | Loss: {epoch_loss/max(n_batches,1):.6f} | '\n",
+        "              f'MSE: {mse_loss.item():.6f} | Cos: {cos_loss.item():.4f} | {elapsed:.1f}s')\n",
         "\n",
-        "print(f'\\n✅ LatentMamba training complete!')"
+        "print('\\n✅ LatentMamba training complete!')"
       ]
     },
     {
@@ -581,51 +352,18 @@
         "# 8. Plot Training Curves\n",
         "# ============================================================\n",
         "import matplotlib.pyplot as plt\n",
-        "\n",
         "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
         "\n",
-        "# VAE Loss\n",
-        "axes[0].plot(history['loss'], alpha=0.3, color='blue')\n",
-        "window = min(50, len(history['loss']) // 5) if len(history['loss']) > 10 else 1\n",
-        "if window > 1:\n",
-        "    smoothed = np.convolve(history['loss'], np.ones(window)/window, mode='valid')\n",
-        "    axes[0].plot(smoothed, color='blue', linewidth=2)\n",
-        "axes[0].set_title('PhraseVAE Loss')\n",
-        "axes[0].set_xlabel('Step')\n",
-        "axes[0].set_ylabel('Loss')\n",
-        "axes[0].grid(True, alpha=0.3)\n",
-        "\n",
-        "# VAE KL\n",
-        "axes[1].plot(history['kl'], alpha=0.5, color='red')\n",
-        "axes[1].set_title('KL Divergence')\n",
-        "axes[1].set_xlabel('Step')\n",
-        "axes[1].set_ylabel('KL')\n",
-        "axes[1].grid(True, alpha=0.3)\n",
-        "\n",
-        "# Mamba Loss\n",
-        "axes[2].plot(mamba_history['mse'], alpha=0.3, color='green')\n",
-        "window = min(50, len(mamba_history['mse']) // 5) if len(mamba_history['mse']) > 10 else 1\n",
-        "if window > 1:\n",
-        "    smoothed = np.convolve(mamba_history['mse'], np.ones(window)/window, mode='valid')\n",
-        "    axes[2].plot(smoothed, color='green', linewidth=2)\n",
-        "axes[2].set_title('LatentMamba MSE Loss')\n",
-        "axes[2].set_xlabel('Step')\n",
-        "axes[2].set_ylabel('MSE')\n",
-        "axes[2].grid(True, alpha=0.3)\n",
-        "\n",
-        "plt.tight_layout()\n",
-        "plt.savefig('/content/training_curves.png', dpi=150)\n",
-        "plt.show()\n",
-        "print('📊 Training curves saved to /content/training_curves.png')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Generation\n",
+        "for ax, data, title, color in [\n",
+        "    (axes[0], history['loss'], 'PhraseVAE Loss', 'blue'),\n",
+        "    (axes[1], history['kl'], 'KL Divergence', 'red'),\n",
+        "    (axes[2], mamba_history['mse'], 'LatentMamba MSE', 'green')]:\n",
+        "    ax.plot(data, alpha=0.3, color=color)\n",
+        "    w = min(50, max(1, len(data)//5))\n",
+        "    if w > 1: ax.plot(np.convolve(data, np.ones(w)/w, 'valid'), color=color, lw=2)\n",
+        "    ax.set_title(title); ax.set_xlabel('Step'); ax.grid(True, alpha=0.3)\n",
         "\n",
-        "Generate MIDI music using the trained model!"
+        "plt.tight_layout(); plt.savefig('/content/training_curves.png', dpi=150); plt.show()"
       ]
     },
     {
@@ -638,80 +376,42 @@
         "# 9. Generate Music!\n",
         "# ============================================================\n",
         "from model import MuseMorphic\n",
-        "from tokenizer import REMIPlusTokenizer, notes_to_midi_file\n",
+        "from tokenizer import notes_to_midi_file\n",
         "\n",
-        "# Assemble full model\n",
         "full_model = MuseMorphic(config).to(device)\n",
         "full_model.phrase_vae = vae\n",
         "full_model.latent_mamba = mamba\n",
         "full_model.eval()\n",
         "\n",
-        "# Unfreeze VAE for generation (was frozen for Stage 2)\n",
-        "# (No gradient computation needed for generation anyway)\n",
-        "\n",
-        "# Generation settings\n",
-        "N_PHRASES = 16    # @param {type: 'integer'}\n",
-        "TEMPERATURE = 0.7 # @param {type: 'number'}\n",
-        "\n",
+        "N_PHRASES = 16; TEMPERATURE = 0.7\n",
         "print(f'Generating {N_PHRASES} phrases at temperature {TEMPERATURE}...')\n",
         "\n",
         "with torch.no_grad():\n",
-        "    # Generate latent sequence\n",
-        "    z_generated = mamba.generate(\n",
-        "        n_phrases=N_PHRASES,\n",
-        "        temperature=TEMPERATURE,\n",
-        "        batch_size=1,\n",
-        "    )\n",
-        "    print(f'Generated latent shape: {z_generated.shape}')\n",
-        "    \n",
-        "    # Decode each phrase latent to tokens\n",
+        "    z_gen = mamba.generate(n_phrases=N_PHRASES, temperature=TEMPERATURE, batch_size=1)\n",
+        "    print(f'Latent shape: {z_gen.shape}')\n",
+        "\n",
         "    all_tokens = []\n",
-        "    for t in range(z_generated.shape[1]):\n",
-        "        z = z_generated[:, t]  # (1, latent_dim)\n",
-        "        \n",
-        "        # Autoregressive decode\n",
-        "        generated_ids = [config.bos_token_id]\n",
-        "        max_decode_len = 128\n",
-        "        \n",
-        "        for _ in range(max_decode_len):\n",
-        "            input_tensor = torch.tensor([generated_ids], dtype=torch.long, device=device)\n",
+        "    for t in range(z_gen.shape[1]):\n",
+        "        z = z_gen[:, t]\n",
+        "        gen_ids = [config.bos_token_id]\n",
+        "        for _ in range(128):\n",
+        "            inp = torch.tensor([gen_ids], dtype=torch.long, device=device)\n",
         "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
-        "                logits = vae.decode(z, input_tensor)\n",
-        "            \n",
-        "            next_logits = logits[0, -1] / max(TEMPERATURE, 0.1)\n",
-        "            probs = F.softmax(next_logits, dim=-1)\n",
-        "            next_token = torch.multinomial(probs, 1).item()\n",
-        "            generated_ids.append(next_token)\n",
-        "            \n",
-        "            if next_token == config.eos_token_id:\n",
-        "                break\n",
-        "        \n",
-        "        phrase_tokens = tokenizer.decode(generated_ids)\n",
-        "        all_tokens.extend(phrase_tokens)\n",
-        "\n",
-        "print(f'\\nGenerated {len(all_tokens)} REMI+ tokens')\n",
-        "print(f'Sample tokens: {all_tokens[:20]}')\n",
-        "\n",
-        "# Convert to MIDI notes\n",
+        "                logits = vae.decode(z, inp)\n",
+        "            probs = F.softmax(logits[0, -1] / max(TEMPERATURE, 0.1), dim=-1)\n",
+        "            tok = torch.multinomial(probs, 1).item()\n",
+        "            gen_ids.append(tok)\n",
+        "            if tok == config.eos_token_id: break\n",
+        "        all_tokens.extend(tokenizer.decode(gen_ids))\n",
+        "\n",
         "notes = tokenizer.tokens_to_midi_notes(all_tokens)\n",
-        "print(f'Extracted {len(notes)} notes')\n",
+        "print(f'Generated {len(notes)} notes from {len(all_tokens)} tokens')\n",
         "\n",
         "if notes:\n",
-        "    # Write MIDI file\n",
-        "    output_midi = '/content/generated_music.mid'\n",
-        "    success = notes_to_midi_file(notes, output_midi)\n",
-        "    if success:\n",
-        "        print(f'\\n🎵 MIDI file saved to: {output_midi}')\n",
-        "        print(f'  Notes: {len(notes)}')\n",
-        "        if notes:\n",
-        "            total_duration = max(n[\"start\"] + n[\"duration\"] for n in notes)\n",
-        "            print(f'  Duration: ~{total_duration/480:.1f} beats')\n",
-        "            pitches = [n[\"pitch\"] for n in notes]\n",
-        "            print(f'  Pitch range: {min(pitches)}-{max(pitches)}')\n",
-        "    else:\n",
-        "        print('⚠️ MIDI writing failed. Install midiutil: pip install midiutil')\n",
+        "    notes_to_midi_file(notes, '/content/generated_music.mid')\n",
+        "    print('🎵 MIDI saved to /content/generated_music.mid')\n",
         "else:\n",
-        "    print('⚠️ No notes generated. Try training for more epochs or adjusting temperature.')"
+        "    print('⚠️ No notes. Try more training epochs or different temperature.')"
       ]
     },
     {
@@ -725,36 +425,16 @@
         "# ============================================================\n",
         "import os\n",
         "from dataclasses import asdict\n",
-        "\n",
         "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
-        "\n",
-        "# Save model\n",
         "save_path = os.path.join(OUTPUT_DIR, 'musemorphic_model.pt')\n",
-        "torch.save({\n",
-        "    'vae_state_dict': vae.state_dict(),\n",
-        "    'mamba_state_dict': mamba.state_dict(),\n",
-        "    'config': asdict(config),\n",
-        "    'training_history': {\n",
-        "        'vae': history,\n",
-        "        'mamba': mamba_history,\n",
-        "    }\n",
-        "}, save_path)\n",
-        "print(f'✅ Model saved to {save_path}')\n",
-        "print(f'  File size: {os.path.getsize(save_path)/1e6:.1f} MB')\n",
-        "\n",
-        "# Save tokenizer\n",
+        "torch.save({'vae_state_dict': vae.state_dict(), 'mamba_state_dict': mamba.state_dict(),\n",
+        "            'config': asdict(config), 'training_history': {'vae': history, 'mamba': mamba_history}}, save_path)\n",
         "tokenizer.save(os.path.join(OUTPUT_DIR, 'tokenizer'))\n",
-        "print(f'✅ Tokenizer saved')\n",
+        "print(f'✅ Model saved to {save_path} ({os.path.getsize(save_path)/1e6:.1f} MB)')\n",
         "\n",
-        "# Optional: Push to HF Hub\n",
         "if PUSH_TO_HUB and HUB_MODEL_ID:\n",
         "    from huggingface_hub import HfApi\n",
-        "    api = HfApi()\n",
-        "    api.upload_folder(\n",
-        "        folder_path=OUTPUT_DIR,\n",
-        "        repo_id=HUB_MODEL_ID,\n",
-        "        repo_type='model',\n",
-        "    )\n",
+        "    HfApi().upload_folder(folder_path=OUTPUT_DIR, repo_id=HUB_MODEL_ID, repo_type='model')\n",
         "    print(f'✅ Pushed to https://huggingface.co/{HUB_MODEL_ID}')"
       ]
     },
@@ -765,52 +445,17 @@
       "outputs": [],
       "source": [
         "# ============================================================\n",
-        "# 11. Listen to Generated MIDI (in Colab)\n",
+        "# 11. Listen to Generated MIDI\n",
         "# ============================================================\n",
         "try:\n",
         "    from IPython.display import Audio, display\n",
         "    import pretty_midi\n",
-        "    \n",
-        "    # Synthesize MIDI to audio using FluidSynth\n",
         "    pm = pretty_midi.PrettyMIDI('/content/generated_music.mid')\n",
         "    audio = pm.fluidsynth(fs=22050)\n",
-        "    \n",
-        "    print('🎧 Listen to generated music:')\n",
-        "    display(Audio(audio, rate=22050))\n",
+        "    print('🎧 Listen:'); display(Audio(audio, rate=22050))\n",
         "except Exception as e:\n",
-        "    print(f'Audio playback not available: {e}')\n",
-        "    print('Download the MIDI file and play it in any MIDI player.')\n",
-        "    print('File: /content/generated_music.mid')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "---\n",
-        "\n",
-        "## Architecture Summary\n",
-        "\n",
-        "### Novel Contributions\n",
-        "\n",
-        "1. **First SSM-based latent music generator**: Mamba operating on compressed phrase latents\n",
-        "2. **FME with log-frequency encoding**: Physics-aware embeddings respecting harmonic series\n",
-        "3. **Multi-attribute control via latent conditioning**: Tempo, key, density, style\n",
-        "4. **Guaranteed training stability stack**: σReparam + ZClip + Pre-LN + BF16 + label smoothing\n",
-        "5. **Three-stage PhraseVAE curriculum**: Prevents posterior collapse\n",
-        "6. **Sub-1GB inference**: Phrase-level Mamba recurrence with fixed-size state\n",
-        "\n",
-        "### Key References\n",
-        "\n",
-        "- Gu & Dao (2023). **Mamba**: Linear-Time Sequence Modeling with Selective State Spaces\n",
-        "- **MIDI-RWKV** (2025). Personalizable Long-Context Symbolic Music Infilling\n",
-        "- **PhraseVAE** (2024). Phrase-level latent diffusion for music\n",
-        "- **FME** (2022). Domain-Knowledge-Inspired Music Embedding\n",
-        "- **σReparam** (2023). Stabilizing Transformer Training\n",
-        "- **ZClip** (2025). Adaptive Spike Mitigation for LLM Pre-Training\n",
-        "- **REMI** (2020). Pop Music Transformer\n",
-        "\n",
-        "📄 Full architecture document: [https://huggingface.co/asdf98/MuseMorphic](https://huggingface.co/asdf98/MuseMorphic)"
+        "    print(f'Audio playback unavailable: {e}')\n",
+        "    print('Download /content/generated_music.mid and play in any MIDI player.')"
       ]
     }
   ]