I found a critical issue during review — the `start_emb` is never used during training, so it won't receive gradients. I need to fix this by prepending it to the training input sequence. Let me fix this:


Edit

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema[:, 1:]  # [B, T-1, d] — EMA targets, shifted

        # Predictor input: online embeddings, with optional noise injection
        input_emb = z[:, :-1]     # [B, T-1, d]
        if self.training and self.input_noise > 0:
            input_emb = input_emb + torch.randn_like(input_emb) * self.input_noise

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if pred_mask is not None:
            valid = ~pred_mask
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)
        B = z.shape[0]
        T_down = z.shape[1]

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Prepend learnable start embedding to input sequence
        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
        # This trains the model to predict the first audio embedding from start_emb + text
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

        # Input noise: only on positions 1+ (real embeddings), not start_emb
        # Simulates the imperfect predictions seen during AR inference
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            # Predictor input mask: start_emb (pos 0) is always unmasked
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
            # Loss mask: follows original ds_mask for target positions
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)


The file /home/ubuntu/lewm-tts/model.py has been updated successfully.

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)
        B = z.shape[0]
        T_down = z.shape[1]

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Prepend learnable start embedding to input sequence
        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
        # This trains the model to predict the first audio embedding from start_emb + text
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

        # Input noise: only on positions 1+ (real embeddings), not start_emb
        # Simulates the imperfect predictions seen during AR inference
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            # Predictor input mask: start_emb (pos 0) is always unmasked
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
            # Loss mask: follows original ds_mask for target positions
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

# ─── Text Encoder ───────────────────────────────────────────────────────────

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

# ─── JEPA Predictor ─────────────────────────────────────────────────────────

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

    def forward(self, x):
        return x + self.pe[:, :x.shape[1]]

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Learnable start-of-sequence embedding for AR inference
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Prediction weight — primary JEPA loss
        self.pred_weight = config.get("pred_weight", 10.0)

        # Input noise injection to bridge train/inference gap
        self.input_noise = config.get("input_noise", 0.0)

        # Free-bits KL: allow KL up to this threshold without penalty
        self.free_bits = config.get("free_bits", 1.0)

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    @torch.no_grad()
    def update_ema(self):
        """Update EMA target encoder weights."""
        for p_online, p_ema in zip(self.audio_encoder.parameters(),
                                   self.ema_audio_encoder.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema[:, 1:]  # [B, T-1, d] — EMA targets, shifted

        # Predictor input: online embeddings, with optional noise injection
        input_emb = z[:, :-1]     # [B, T-1, d]
        if self.training and self.input_noise > 0:
            input_emb = input_emb + torch.randn_like(input_emb) * self.input_noise

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if pred_mask is not None:
            valid = ~pred_mask
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
        # Prevents both KL collapse (too low) and KL explosion (too high)
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid_kl = (~ds_mask).unsqueeze(-1).float()
            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = kl_free.mean()

        # ─── Loss 3: Mel reconstruction ───
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Multi-resolution spectral loss ───
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (self.pred_weight * prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.05,
            "pred_weight": 10.0,
            "ema_decay": 0.998,
            "free_bits": 1.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

-            target_emb = mu_ema[:, 1:]  # [B, T-1, d] — EMA targets, shifted

+            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

-        # Predictor input: online embeddings, with optional noise injection

+        # Prepend learnable start embedding to input sequence

+        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions

+        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions

+        # This trains the model to predict the first audio embedding from start_emb + text

+        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]

+        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

+        # Input noise: only on positions 1+ (real embeddings), not start_emb

+        # Simulates the imperfect predictions seen during AR inference

-            input_emb = input_emb + torch.randn_like(input_emb) * self.input_noise

+            noise = torch.zeros_like(input_emb)

+            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],

+                                        device=input_emb.device) * self.input_noise

+            input_emb = input_emb + noise

-        # Adjust mask for downsampled sequence

+        # Adjust masks for downsampled sequence

-            pred_mask = ds_mask[:, :-1]

+            # Predictor input mask: start_emb (pos 0) is always unmasked

+            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)

+            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]

+            # Loss mask: follows original ds_mask for target positions

-            valid = valid.unsqueeze(-1)  # [B, T-1, 1]

+            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]


Now verify this change works correctly:


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
from model import build_model
import torch

config = {
    'd_model': 256, 'nhead': 4, 'n_mels': 100,
    'text_vocab_size': 256,
    'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
    'dropout': 0.1, 'downsample_factor': 4,
    'kl_weight': 0.05, 'recon_weight': 1.0, 'spectral_weight': 0.5,
    'pred_weight': 10.0, 'ema_decay': 0.998, 'free_bits': 1.0, 'input_noise': 0.1,
}
model, config = build_model(config)
model = model.cuda()

# Test with padding
B = 4
mel = torch.randn(B, 100, 200).cuda()
text = torch.randint(0, 256, (B, 50)).cuda()
mel_mask = torch.zeros(B, 200, dtype=torch.bool).cuda()
mel_mask[2, 150:] = True  # some padding
mel_mask[3, 100:] = True
text_mask = torch.zeros(B, 50, dtype=torch.bool).cuda()
text_mask[2, 30:] = True

model.train()
losses = model(mel, text, mel_mask, text_mask)
print('Forward OK with padding')
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')

losses['total_loss'].backward()
print('Backward OK')

# Check start_emb gets gradients now
print(f'start_emb grad: {model.start_emb.grad is not None} (should be True)')
if model.start_emb.grad is not None:
    print(f'  grad norm: {model.start_emb.grad.norm().item():.6f}')

# Test EMA update
model.update_ema()

# Verify sequence shapes in predictor
z, mu, logvar = model.audio_encoder(mel, mel_mask)
T_down = z.shape[1]
print(f'z shape: {z.shape}')
print(f'T_down: {T_down}, input would be [B, {T_down}, d], target would be [B, {T_down}, d]')
"

Verify updated forward with start_emb gradients


LeWM TTS model: 24.00M parameters
Forward OK with padding
  total_loss: 7.9432
  prediction_loss: 0.4056
  kl_loss: 0.0102
  recon_loss: 0.7997
  spectral_loss: 6.1746
Backward OK
start_emb grad: True (should be True)
  grad norm: 0.033308
z shape: torch.Size([4, 50, 256])
T_down: 50, input would be [B, 50, d], target would be [B, 50, d]
Shell cwd was reset to /home/ubuntu/pipecat


`start_emb` now receives gradients. Let me do a final comprehensive sanity check from a researcher's perspective:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model import build_model

# ============================================================
# RESEARCH REVIEW: Verify training dynamics are correct
# ============================================================

config = {
    'd_model': 256, 'nhead': 4, 'n_mels': 100,
    'text_vocab_size': 256,
    'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
    'dropout': 0.1, 'downsample_factor': 4,
    'kl_weight': 0.05, 'recon_weight': 1.0, 'spectral_weight': 0.5,
    'pred_weight': 10.0, 'ema_decay': 0.998, 'free_bits': 1.0, 'input_noise': 0.15,
}
model, config = build_model(config)
model = model.cuda()

B, T_mel, T_text = 8, 400, 80
mel = torch.randn(B, 100, T_mel).cuda()
text = torch.randint(0, 256, (B, T_text)).cuda()
mel_mask = torch.zeros(B, T_mel, dtype=torch.bool).cuda()
text_mask = torch.zeros(B, T_text, dtype=torch.bool).cuda()

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

print('=== CHECK 1: EMA diverges from online encoder ===')
# After init, EMA == online. After a few steps, they should diverge
online_p0 = model.audio_encoder.conv_pre[0].weight.data.clone()
ema_p0 = model.ema_audio_encoder.conv_pre[0].weight.data.clone()
print(f'  Init diff: {(online_p0 - ema_p0).abs().max().item():.8f} (should be ~0)')

for step in range(10):
    model.train()
    losses = model(mel, text, mel_mask, text_mask)
    optimizer.zero_grad()
    losses['total_loss'].backward()
    optimizer.step()
    model.update_ema()

online_p0 = model.audio_encoder.conv_pre[0].weight.data.clone()
ema_p0 = model.ema_audio_encoder.conv_pre[0].weight.data.clone()
print(f'  After 10 steps diff: {(online_p0 - ema_p0).abs().max().item():.6f} (should be small but nonzero)')

print()
print('=== CHECK 2: Free-bits KL behavior ===')
# With free_bits=1.0, if all KL dimensions are <1.0, loss should be 0
mu_low = torch.zeros(2, 10, 256).cuda()
logvar_low = torch.zeros(2, 10, 256).cuda()  # KL=0 per dim
kl_per_dim = -0.5 * (1 + logvar_low - mu_low.pow(2) - logvar_low.exp())
kl_free = torch.clamp(kl_per_dim - 1.0, min=0.0)
print(f'  KL per dim when mu=0,logvar=0: {kl_per_dim[0,0,0].item():.4f} (should be 0)')
print(f'  Free-bits clipped: {kl_free.mean().item():.4f} (should be 0)')

# With large KL, should penalize
mu_high = torch.ones(2, 10, 256).cuda() * 3.0
kl_per_dim = -0.5 * (1 + logvar_low - mu_high.pow(2) - logvar_low.exp())
kl_free = torch.clamp(kl_per_dim - 1.0, min=0.0)
print(f'  KL per dim when mu=3: {kl_per_dim[0,0,0].item():.4f}')
print(f'  Free-bits clipped: {kl_free.mean().item():.4f} (should be >0)')

print()
print('=== CHECK 3: Gradient flow verification ===')
model.zero_grad()
model.train()
losses = model(mel, text, mel_mask, text_mask)
losses['total_loss'].backward()

components = {
    'text_encoder': model.text_encoder,
    'audio_encoder': model.audio_encoder,
    'predictor': model.predictor,
    'mel_decoder': model.mel_decoder,
    'start_emb': None,
}

for name, module in components.items():
    if name == 'start_emb':
        g = model.start_emb.grad
        print(f'  {name}: grad_norm={g.norm().item():.6f}')
    else:
        grads = [p.grad.norm().item() for p in module.parameters() if p.grad is not None]
        if grads:
            print(f'  {name}: {len(grads)} params with grads, max_norm={max(grads):.4f}')
        else:
            print(f'  {name}: NO GRADIENTS!')

# EMA should NOT have gradients
ema_grads = [p.grad for p in model.ema_audio_encoder.parameters() if p.grad is not None]
print(f'  ema_audio_encoder: {len(ema_grads)} params with grads (should be 0)')

print()
print('=== CHECK 4: Prediction loss should be higher with EMA targets ===')
# EMA targets are different from online targets, so prediction is harder
# This is good — it forces the predictor to learn robust features
model.eval()
with torch.no_grad():
    z, mu, _ = model.audio_encoder(mel)
    _, mu_ema, _ = model.ema_audio_encoder(mel)
    diff = (mu - mu_ema).abs().mean().item()
    print(f'  Online vs EMA embedding diff: {diff:.6f}')
    print(f'  (small after 10 steps, will grow as training progresses)')

print()
print('=== ALL CHECKS PASSED ===')
"

Comprehensive research verification of training dynamics


LeWM TTS model: 24.00M parameters
=== CHECK 1: EMA diverges from online encoder ===
  Init diff: 0.00000000 (should be ~0)
  After 10 steps diff: 0.000969 (should be small but nonzero)

=== CHECK 2: Free-bits KL behavior ===
  KL per dim when mu=0,logvar=0: -0.0000 (should be 0)
  Free-bits clipped: 0.0000 (should be 0)
  KL per dim when mu=3: 4.5000
  Free-bits clipped: 3.5000 (should be >0)

=== CHECK 3: Gradient flow verification ===
  text_encoder: 51 params with grads, max_norm=0.1941
  audio_encoder: 58 params with grads, max_norm=0.6508
  predictor: 114 params with grads, max_norm=3.2114
  mel_decoder: 22 params with grads, max_norm=6.6089
  start_emb: grad_norm=0.000861
  ema_audio_encoder: 0 params with grads (should be 0)

=== CHECK 4: Prediction loss should be higher with EMA targets ===
  Online vs EMA embedding diff: 0.436841
  (small after 10 steps, will grow as training progresses)

=== ALL CHECKS PASSED ===
Shell cwd was reset to /home/ubuntu/pipecat


All checks pass. Let me complete the research review task:


Updated task #4 status


Now let me save the key findings to memory for future reference:


Write

# Project Memory

## lewm-tts (JEPA TTS)
- Location: `/home/ubuntu/lewm-tts/`
- JEPA-based TTS for Hindi, 24M params, Vocos vocoder
- See [lewm-tts.md](lewm-tts.md) for details


# LeWM TTS - JEPA-based Hindi TTS

## Key Issues Found (v3/v3b/v3c)
1. Missing EMA target encoder (core JEPA component)
2. Prediction loss was <1% of total loss — predictor learned identity mapping
3. KL divergence exploded to 2.45 (kl_weight=0.01 too low)
4. v3c's downsample_factor=2 made consecutive embeddings too similar
5. Random start embedding at inference (never trained)
6. No input noise to bridge train/inference gap

## v4 Fixes Applied
- **EMA target encoder**: deepcopy of AudioEncoder, updated with momentum 0.998
- **Rebalanced losses**: pred_weight=10.0 (now ~36% of total), kl_weight=0.05
- **Free-bits KL**: threshold=1.0, prevents both collapse and explosion
- **Learnable start_emb**: prepended to training sequence, receives gradients
- **Input noise**: annealed 0→0.2 over first 30% of training
- **Downsample factor**: back to 4 (more temporal separation)

## Training Command (v4)
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v4 \
  --epochs 200 --batch_size 256 \
  --pred_weight 10.0 --kl_weight 0.05 \
  --recon_weight 1.0 --spectral_weight 0.5 \
  --ema_decay 0.998 --free_bits 1.0 --input_noise_max 0.2
```

## Architecture
TextEncoder(4L,256d) + AudioEncoder(CNN+4L,256d,stride4) + JEPAPredictor(6L decoder) + MelDecoder(512h,ResConv) → Vocos 24kHz

## Dataset
7009 segments, 17.88h single-speaker Hindi, 100-mel, processed_data_100mel/


File created successfully at: /home/ubuntu/.claude/projects/-home-ubuntu-pipecat/memory/MEMORY.md

# Project Memory

## lewm-tts (JEPA TTS)
- Location: `/home/ubuntu/lewm-tts/`
- JEPA-based TTS for Hindi, 24M params, Vocos vocoder
- See [lewm-tts.md](lewm-tts.md) for details


File created successfully at: /home/ubuntu/.claude/projects/-home-ubuntu-pipecat/memory/lewm-tts.md

# LeWM TTS - JEPA-based Hindi TTS

## Key Issues Found (v3/v3b/v3c)
1. Missing EMA target encoder (core JEPA component)
2. Prediction loss was <1% of total loss — predictor learned identity mapping
3. KL divergence exploded to 2.45 (kl_weight=0.01 too low)
4. v3c's downsample_factor=2 made consecutive embeddings too similar
5. Random start embedding at inference (never trained)
6. No input noise to bridge train/inference gap

## v4 Fixes Applied
- **EMA target encoder**: deepcopy of AudioEncoder, updated with momentum 0.998
- **Rebalanced losses**: pred_weight=10.0 (now ~36% of total), kl_weight=0.05
- **Free-bits KL**: threshold=1.0, prevents both collapse and explosion
- **Learnable start_emb**: prepended to training sequence, receives gradients
- **Input noise**: annealed 0→0.2 over first 30% of training
- **Downsample factor**: back to 4 (more temporal separation)

## Training Command (v4)
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v4 \
  --epochs 200 --batch_size 256 \
  --pred_weight 10.0 --kl_weight 0.05 \
  --recon_weight 1.0 --spectral_weight 0.5 \
  --ema_decay 0.998 --free_bits 1.0 --input_noise_max 0.2
```

## Architecture
TextEncoder(4L,256d) + AudioEncoder(CNN+4L,256d,stride4) + JEPAPredictor(6L decoder) + MelDecoder(512h,ResConv) → Vocos 24kHz

## Dataset
7009 segments, 17.88h single-speaker Hindi, 100-mel, processed_data_100mel/


---

## Research Review Summary

Here's the complete picture of what was wrong and what I've fixed, from a researcher's perspective:

### What Was Wrong (Root Causes)

| Issue | Impact | Evidence |
|-------|--------|----------|
| **No EMA target encoder** | Predictor targets are identical to inputs → trivially predictable | pred_loss=0.017, near-identity mapping |
| **pred_weight missing** (implicitly 1.0) | Prediction was 0.86% of total loss | 97.9% loss from recon+spectral |
| **kl_weight=0.01** | KL exploded to 2.45, latent space uncontrollable at inference | AR generation starts from wrong distribution |
| **No free-bits** | KL penalty fights reconstruction — either collapses or explodes | KL monotonically increased throughout training |
| **downsample_factor=2** (v3c) | Consecutive embeddings 21ms apart → near-identical → trivial prediction | Prediction loss dropped to 0.017 |
| **Random start_emb** | `randn * 5.0` at inference — completely untrained | First AR step produces garbage, cascades |
| **No input noise** | Training sees perfect embeddings; inference sees predicted ones | Exposure bias → AR drift |

### What I Fixed

1. **EMA Target Encoder** — `model.ema_audio_encoder` is a frozen deepcopy of the online encoder, updated with momentum 0.998 after each optimizer step. Prediction targets come from EMA (stable, prevents collapse). Verified: zero gradients, correct divergence rate.

2. **Loss Rebalancing** — `pred_weight=10.0` makes prediction ~36% of total loss (was <1%). The predictor now has meaningful training signal.

3. **Free-bits KL** — `free_bits=1.0` means per-dimension KL below 1.0 is free (no penalty). This prevents KL collapse while capping the explosion. Verified: KL=0 → loss=0, KL=4.5 → loss=3.5.

4. **Learnable start_emb** — Prepended to the training input sequence so position 0 predicts the first audio embedding from `start_emb + text`. Verified: receives gradients during backprop.

5. **Input noise injection** — Gaussian noise (annealed 0→0.2 over 30% of training) added to positions 1+ of the input sequence. Simulates the imperfect AR predictions at inference.

6. **Downsample factor = 4** — Consecutive embeddings are now ~42ms apart (vs 21ms in v3c), making prediction a genuinely non-trivial task.

### Expected Training Dynamics

- **Prediction loss** should stabilize around 0.05-0.15 (not 0.017) — the EMA targets and higher weight force the predictor to learn actual speech dynamics
- **KL** should settle around 0.5-2.0 with free-bits preventing explosion
- **AR synthesis** should produce coherent speech because the predictor learned meaningful transitions and the start embedding is trained

### Verified
- Forward/backward pass: works
- EMA divergence: correct (starts at 0, grows slowly)
- Free-bits math: correct
- Gradient flow: all components receive gradients, EMA does not
- Checkpoint save/load: roundtrip works including EMA state

The model is ready to train. Want me to kick off a v4 training run?


I wanted to check again because I'm very confident that it will work. If it works, I can go tell anyone that I trained this model, that it is very good for And good for you.


The user wants me to do another thorough check before training. They're very confident this approach will work and want to make sure everything is solid. Let me do a deep, researcher-level review of the entire codebase — checking for subtle bugs, mathematical correctness, potential training pitfalls, and anything that could prevent good results.

Let me re-read all the modified files carefully and think about potential issues.


Let me do a deep, line-by-line audit of everything. This needs to be bulletproof before we train.