2025 - Thesis - ASR Knowledge Distillation

Sonnet 64

And nothing ’gainst Time’s scythe can make defence,
From each of life moment, to brave him when he takes thee hence.

**Shakespeare, William.** “Sonnet 64,” lines 11–12, in The Sonnets * (1609).*

Training Loop
    ↓
[ Forward pass ]
    ↓
[ Compute loss ]
    ↓
[ Backward pass: compute gradients ]
    ↓
[ **Gradient Clipping** ]       ←— `clip_grad_norm_(model.params, max_norm)`
    ↓
[ **AdamW Update** ]            ←— `optimizer = AdamW(lr=…, weight_decay=…)`
    ↓
[ Zero Gradients ]          ←— `optimizer.zero_grad()`
    ↓
[ **Cosine LR Annealing** ]     ←— `scheduler = CosineAnnealingLR(optimizer, T_max, eta_min)`
    ↓
[ Next batch ]

the parameters that dominate memory footprint and inference latency
- fused LoRA into the backbone and applied quantization:

1. Embedding matrix
Quantized token‑to‑hidden lookup table (+ scale/zero‑point metadata)

2. Fused Transformer weights
Self‑Attention: Q, K, V, and output projection matrices (each with fused LoRA ΔW), plus biases; all quantized
Feed‑Forward: two linear layers (with fused LoRA updates), plus biases; all quantized

3. LayerNorm parameters
γ (scale) and β (shift) for each layer—typically kept in FP16/FP32 for stability

4. Output projection head
Final hidden‑to‑vocab weight matrix and bias, quantized (+ metadata)

5. Quantization metadata
Per‑tensor (or per‑channel) scale and zero‑point arrays that map integer ops back to real values

- Extract the spectral or time domain features, then train the U-Net, Conv-TasNet, Demucs and other networks to output multiple audio streams

- Give it a "label" to tell it how many channels to split (vocals/drums/bass/other), and it will split the signal

%load_ext tensorboard
%tensorboard --logdir /content/drive/MyDrive/distil_run_cell2.x/tb

Raw Audio
    ↓
[ Encoder (frozen) ]
    ↓
[ Decoder + LoRA adapters ]    ←— LoRA injects low-rank ΔW here, directly modifying student logits
    ↓
Student logits zˢᵗᵘᵈᵉⁿᵗ      ←— LoRA’s effect appears in these logits
    ↓
÷ T   (temperature scaling)     ←— Temperature T then scales these logits
    ↓
softmax(z/T)                   ←— T smooths the distribution, controlling gradient strength
    ↓
KL(pᵗᵉᵃᶜʰᵉʳ ‖ pˢᵗᵘᵈᵉⁿᵗ)
    ↓
Backpropagate to update LoRA adapter parameters

🍩 Transformer = Non-Sugar Donut Factory Assembly Line  
═══════════════════════════════════════════════════════
Raw Donuts → Community Check → Solo Decoration → Finished Donuts  
(Input)       (Attention)       (FFN)            (Output)
    ↓             ↓                ↓               ↓
┌─────────┐  ┌──────────────┐  ┌──────────────┐    ┌─────────┐
│ Plain   │→ │👥 Community  │→ │ 🎨 Solo       │ →  │Gourmet  │
│ Donuts  │  │  Analysis    │  │ Decoration   │    │Donuts   │
└─────────┘  └──────────────┘  └──────────────┘    └─────────┘
   ↓                ↓                 ↓                 ↓
   X₀          X₁ = Attention     X₂ = FFN            Output

 1. X₁ᵢ = Σⱼ αᵢⱼ × V_j               (Global Linear)  
 2. X₂ᵢ = W₂·ReLU(W₁·X₁ᵢ + b₁) + b₂  (Local Nonlinear)

Attention: Convex combination → Stays within input space
FFN: Nonlinear transformation → Can transcend input space

Activation Function Characteristics Comparison:
═════════════════════════════════════════════════
┌──────────┬────────────┬───────────────┬──────────────┬─────────────┐
│Function  │ Smoothness │ Computational │ Gradient     │ Performance │
│          │            │ Complexity    │ Properties   │             │
├──────────┼────────────┼───────────────┼──────────────┼─────────────┤
│ ReLU     │ Non-smooth │ Minimal       │ Sparse       │ Baseline    │
│ GELU     │ Smooth     │ Moderate      │ Dense        │ Better      │
│ SwiGLU   │ Smooth     │ High          │ Gated        │ Best        │
│ Mish     │ Very Smooth│ High          │ Adaptive     │ Very Good   │
│ Swish    │ Smooth     │ Moderate      │ Self-gated   │ Good        │
│ ELU      │ Smooth     │ Moderate      │ Negative-safe│ Good        │
└──────────┴────────────┴───────────────┴──────────────┴─────────────┘

Raw Material (Input) → Processing (Distill) → Packaging (Loss) → Finished Product (Student)
    ↓                       ↓                       ↓                    ↓
X₀ (Teacher Output)    X₁ (Distillation)    X₂ (Loss Computation)     Student

**B** = batch size
**C** = number of classes
**D** = hidden feature dimension
**n** = sequence length or number of nodes

Teacher (Whisper-large-v2)                     Student (distil-small.en + LoRA + Hidden Align)
─────────────────────────────                  ──────────────────────────────────────────────────

Audio Input                                     Audio Input  
1 × T samples                                   1 × T samples
     │                                               │
     ▼                                               ▼
Whisper Encoder                                 Whisper Encoder
1280-d hidden, T~1500 frames                    768-d hidden, T~499 frames
(32 layers, FROZEN)                            (12 layers, FROZEN)
     │                                               │
     │                                               │
     ├─────── Hidden States ──────────────────────── ├─── Projection Layer ───┐
     │        (B,1500,1280)                          │    (768→1280)          │
     │                                               │                        │
     ▼                                               ▼                        ▼
Whisper Decoder                                 Whisper Decoder              Aligned Hidden
(32 layers, FROZEN)                            (4 layers +/ LoRA)            (B,1500,1280)
     │                                               │                        │
     │                                               │                        │
     ▼                                               ▼                        │
Teacher Logits ────── Soft Targets ─────────▶ Student Logits                  │
(B,seq,vocab)         (KL Loss)               (B,seq,vocab)                   │
     │                 T=temperature              │                           │
     │                                            │                           │
     │                                            ▼                           │
     │                                      Hard Labels ◀── Ground Truth      │
     │                                      (CTC Loss)                        │
     │                                            │                           │
     │                                            │                           │
     │                                            ▼                           │
     │                                       Student Loss ◀─── MSE Loss ──────┘
     │                                            │         (Hidden Align)
     │                                            │
     ▼                                            ▼
No parameter updates                        (LoRA) + Projection parameters
(Inference only)                             ONLY these are trained

Teacher Encoder Output (t_h)
┌─────────────────────────────────────────┐
│  t_h: shape = (B, T≈1500, 1280)         │
│                                         │
│  ┌─────┐ ┌─────────────┐ ┌─────────────┐│
│  │ B   │ │  T_frames   │ │ Hidden_dim  ││
│  │     ├─│ ≈1500       ├─│ 1280        ││
│  └─────┘ └─────────────┘ └─────────────┘│
└─────────────────────────────────────────┘

Student Encoder Output (s_h)
┌──────────────────────────────────────────┐
│  s_h: shape = (B, T≈499, 768)            │
│                                          │
│  ┌─────┐ ┌─────────────┐ ┌─────────────┐ │
│  │ B   │ │  T_frames   │ │ Hidden_dim  │ │
│  │     ├─│  ≈499       ├─│ 768         │ │
│  └─────┘ └─────────────┘ └─────────────┘ │
└──────────────────────────────────────────┘

↓↓  Student Model: LoRA Injection Points for Encoder & Decoder  ↓↓

           ┌───────────────────────┐
           │   Whisper Encoder     │   ← 12 frozen transformer layers
           │   (layer i)           │
           │                       │
           │  ┌───────────────┐    │    ← If i ∈ {top-4 layers}, 
           │  │ Self-Attn     │◀───┼───┐      inject LoRA into:
           │  │  ┌───┬───┬───┬───┐ │   │      • self_attn.q_proj
           │  │  │ Q │ K │ V │ Out││   │      • self_attn.k_proj
           │  └──┴───┴───┴───┴───┘ │   │      • self_attn.v_proj
           │  ┌───────────────┐    │   │      • self_attn.out_proj
           │  │ Feed-Forward  │    │   │      • fc1, fc2
           │  │  ┌────┐  ┌─────┐│  │   └───────┘
           │  │  │fc1 │  │ fc2 ││  │
           │  └──┴────┴──┴─────┴┘  │
           └───────────────────────┘

                    ↓  (pass projected hidden to Decoder)

           ┌───────────────────────┐
           │  Whisper Decoder      │   ← 4 frozen transformer layers
           │  (layer j)            │
           │                       │
           │  ┌───────────────┐    │    ← inject LoRA into every
           │  │ Self-Attn     │◀───┼───┐      • self_attn.q_proj
           │  │  ┌───┬───┬───┬───┐ │   │      • self_attn.k_proj
           │  │  │ Q │ K │ V │ Out││   │      • self_attn.v_proj
           │  └──┴───┴───┴───┴───┘ │   │      • self_attn.out_proj
           │  ┌───────────────┐    │   │
           │  │ Cross-Attn    │◀───┼───┐      • encoder_attn.q_proj
           │  │  ┌───┬───┬───┬───┐ │   │      • encoder_attn.k_proj
           │  │  │ Q │ K │ V │ Out││   │      • encoder_attn.v_proj
           │  └──┴───┴───┴───┴───┘ │   │      • encoder_attn.out_proj
           │  ┌───────────────┐    │   └───────┘
           │  │ Feed-Forward  │◀───┼───┐      • fc1, fc2
           │  │  ┌────┐  ┌─────┐│  │   └───────┘
           │  │  │fc1 │  │ fc2 ││  │
           │  └──┴────┴──┴─────┴┘  │
           └───────────────────────┘

(All adapters B·A learnable; original weights frozen)

30 s audio
    ⇓ (extract 80-dim log-Mel every 10 ms)
3000 frames
    ⇓ (convolutional layer with stride=2)
1500 frames
    ⇓ (further down-sampling/merging inside the Transformer encoder ≈×3)
    ⇓ (Pooling or Conv1d: kernel_size=3, stride=3)
≈500 frames  (actually 499 frames)

1. Audio frame rate is typically high
sample_rate = 16000      # 16 kHz sampling rate
frame_rate = 100         # 100 frames per second
frame_duration = 10      # 10 ms per frame

2. 30 seconds of audio
total_frames = 30 * frame_rate  # 3000 frames

3. Adjacent frames are highly correlated
correlation_coefficient ≈ 0.9  # typical inter-frame correlation

ΔW = A · B -> only low-rank increments are made to W_q and W_v in the attention

decoder.layers.*.encoder_attn.q_proj
decoder.layers.*.encoder_attn.v_proj
decoder.layers.*.self_attn.q_proj
decoder.layers.*.self_attn.v_proj

decoder.layers.*.encoder_attn.q_proj, encoder_attn.k_proj, encoder_attn.v_proj
decoder.layers.*.self_attn.q_proj, self_attn.k_proj, self_attn.v_proj
decoder.layers.*.fc2

decoder.layers.*.encoder_attn.q_proj, encoder_attn.k_proj, encoder_attn.v_proj, encoder_attn.out_proj
decoder.layers.*.self_attn.q_proj, self_attn.k_proj, self_attn.v_proj, self_attn.out_proj
decoder.layers.*.fc1, fc2

import optuna
from optuna.pruners import MedianPruner
from optuna.samplers import TPESampler

def objective(trial):
    # Distillation loss weights
    alpha = trial.suggest_loguniform("alpha", 1e-3, 1e1)
    beta  = trial.suggest_loguniform("beta",  1e-3, 1e1)
    # Optimization hyperparameters
    lr        = trial.suggest_loguniform("lr",        1e-5, 1e-3)
    batch_size= trial.suggest_categorical("batch_size", [4, 8, 16, 32])
    dropout   = trial.suggest_float("dropout", 0.0, 0.5)
    
    # Train & evaluate with these settings (implement train_and_evaluate accordingly)
    wer = train_and_evaluate(
        alpha=alpha,
        beta=beta,
        learning_rate=lr,
        batch_size=batch_size,
        dropout=dropout,
        pruner=trial  # for early stopping
    )
    return wer

# Pruner to stop unpromising trials early
pruner  = MedianPruner(n_startup_trials=5, n_warmup_steps=100)
sampler = TPESampler()

study = optuna.create_study(
    direction="minimize",
    sampler=sampler,
    pruner=pruner
)
study.optimize(objective, n_trials=100)

print("Best hyperparameters:", study.best_params)

**Gradient Underflow**
feats = b["input_features"].half().to(device) <- FP32 to 16

self.scaler = **GradScaler()**
...
with autocast():
    loss = model(input)  # loss = float16
self.scaler.scale(loss).backward() 
self.scaler.step(optimizer)  
self.scaler.update()

- Local weights
w_ij = exp(−(d(x_i, x_j) − ρ_i) / σ_i)  
w_ji = exp(−(d(x_j, x_i) − ρ_j) / σ_j)

- Fuse into a single “strength” score
μ_ij = w_ij + w_ji − w_ij * w_ji

🍩 Vocabulary Mismatch = Multilingual "Plain" vs English "Decorated" Pipeline
═══════════════════════════════════════════════════════════════════════════════

Large Vocab → Filter English → Student English‑Only → Aligned Output
(Original)    (Trim Embedding)  (Slice/Filter)        (Consistent vocab)
    ↓                 ↓                      ↓                      ↓
┌─────────────┐   ┌─────────────────┐   ┌─────────────────┐   ┌─────────────┐
│Multi‑Lingual│→  │🍰 Trim Large    │→  │🎯 Keep English   │→  │ Distil Vocab│
│51866 tokens │   │Embedding ↓ to   │   │Tokens Only ↓    │   │51864 size   │
│(v3)         │   │51864            │   │keep 51864       │   │             │
└─────────────┘   └─────────────────┘   └─────────────────┘   └─────────────┘
    ↓                 ↓                      ↓                      ↓
  V₀ (51866)      V₁ (trimmed)         V₂ (student)         V_out (match)

Pipeline Steps:
1. **V₀**: Teacher `openai/whisper-large-v3` uses a 51 866‑token multilingual vocabulary  
2. **V₁**: Resize the teacher’s embedding matrix down to **51 864** slots (remove non‑English tokens)  
3. **V₂**: Student model (English‑only) loads that 51 864‑sized embedding  
4. **V_out**: Ensures output token IDs align exactly between teacher‑derived embeddings and student’s vocab  

Vocabulary Alignment Strategy:
┌──────────────┐    ┌──────────────┐    ┌──────────────┐
│   Teacher    │ →  │   Student    │ →  │   Output     │
│ Whisper-v3   │    │ English‑Only │    │  Aligned     │
│ 51866 tokens │    │  51864 toks  │    │  51864 toks  │
└──────────────┘    └──────────────┘    └──────────────┘

Key Technical Details:
- **Teacher**: `openai/whisper-large-v3` (51 866‑token vocab)  
- **Student**: English‑only distilled model (51 864‑token vocab)  
- **Solution**: Align token spaces by trimming the teacher’s embedding to the student’s English subset  
- **Result**: Consistent token IDs & embedding dimensions for seamless knowledge distillation

VQ-VAE

Motivation for Discrete Latent Audio Representations

- **Information Bottleneck & Compression**  
  Continuous waveforms contain abundant redundancies; modeling them directly (e.g., WaveNet) is computationally expensive.  
  VQ-VAE quantizes latent vectors into a fixed-size codebook (e.g., 512 or 1024 entries), achieving lossy compression while preserving salient features.  
  Subsequent priors (PixelCNN/Transformer) operate on shorter discrete sequences, yielding far greater efficiency.

- **Advantages of Discrete Sequence Modeling**  
  Autoregressive Transformers, PixelCNNs, and RNNs excel at modeling discrete tokens.  
  Mapping audio to discrete tokens enables these models to capture long-range dependencies, linguistic structures, and prosody, without struggling in high-dimensional continuous spaces.

- **Denoising & Semantic Extraction**  
  Codebook entries tend to represent prototypical acoustic units (e.g., phonemes).  
  Quantization suppresses minor noise and emphasizes semantic/prosodic features, yielding more robust representations for synthesis and recognition.

- **Hierarchical & Multi-Rate Structure**  
  VQ-VAE-2 employs multi-scale codebooks:  
  - **Lower levels:** fine details (timbre, noise)  
  - **Higher levels:** global structure (intonation, rhythm)  
  This hierarchical discrete design is hard to achieve with continuous latent models.

1. NVIDIA Megatron-LM (GPT / BERT Pretraining)

- **Compute**: All forward/backward kernels (e.g., matrix multiplies, attention, activations) execute in **FP16** on Tensor Cores.  
- **Master Weights & Optimizer State**: A **FP32** “master copy” of model weights is maintained; FP16 gradients are accumulated into FP32 weights, then cast back to FP16 for the next iteration.  
- **Numerical Stability**: Overflow-sensitive ops (e.g., Softmax) remain in **FP32** to avoid divergence, without meaningful speed penalty.  
- **Implementation**: Uses NVIDIA Apex/AMP or NeMo’s automatic mixed-precision modules to orchestrate autocasting, dynamic loss scaling, and master-weight management :contentReference[oaicite:0]{index=0}.  

2. Google PaLM / AlphaFold2 (Transformer & Scientific Models)

- **Core Operators** (e.g., GEMM, LayerNorm, Dropout) run in **bfloat16** (BF16)—a 16-bit format with the same exponent range as FP32 but reduced mantissa.  
- **Critical Sections** (e.g., logits normalization, loss computation) remain in **FP32** to prevent underflow/overflow.  
- **Framework Support**: JAX/Flax users enable BF16 globally via flags like `jax.experimental.enable_x64=False` and per-op control via `jax.lax.cond` :contentReference[oaicite:1]{index=1}.  


3. Stable Diffusion / GAN Inference

- **Inference Precision**: Entire U-Net, scheduler, and CLIP text encoder run in **FP16**, halving memory and nearly doubling throughput compared to FP32, with negligible impact on image quality.  
- **Selective FP32 Fallback**: For maximum stability or visual fidelity (e.g., in the VAE decoder), critical layers can be temporarily cast back to **FP32**.  

4. Frontiers: FP8-LM (8-bit Mixed-Precision)

- **FP8 Framework**: Base weights, gradients, and optimizer states use **8-bit FP8** (e.g., NVIDIA’s NF4 or Microsoft’s custom FP8); key stability ops (Softmax, LayerNorm) stay in FP16/BF16 or FP32.  
- **Efficiency Gains**: In GPT-175B training, FP8 mixed-precision achieved a **75% speedup** and **39% memory reduction** over BF16 baselines—without changing hyperparameters :contentReference[oaicite:2]{index=2}.  
- **Adoption**: Offers a drop-in replacement for existing FP16/BF16 pipelines, open-sourced via MS-AMP (aka.ms/MS.AMP).  

*References*  
- Micikevicius *et al.*, “Mixed Precision Training,” NVIDIA Developer Blog, 2018.  
- Shoeybi *et al.*, “Megatron-LM: Training Multi-Billion Parameter Language Models,” *arXiv:1909.08053*, 2019.  
- Narang *et al.*, “How To Fit a Bigger Model and Train It Faster,” Hugging Face Docs, 2022.  
- Peng *et al.*, “FP8-LM: Training FP8 Large Language Models,” *arXiv:2310.18313*, 2023.

┌─────────────────────────────────────────┐
│  Temporal Downsampling Effect           │
│                                         │
│  Teacher Sequence (1500 frames)         │
│  ┌─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┐  │
│  │▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│▓│  │
│  └─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┘  │
│          ↓ 3:1 compression              │
│  Student Sequence (499 frames)          │
│  ┌─────┬─────┬─────┬─────┬─────┬─────┐  │
│  │ ▓▓▓ │ ▓▓▓ │ ▓▓▓ │ ▓▓▓ │ ▓▓▓ │ ▓▓▓ │  │
│  └─────┴─────┴─────┴─────┴─────┴─────┘  │
│     ↑                                   │
│  Information "leaks" to adjacent windows│
└─────────────────────────────────────────┘

 Audio Signal Characteristics:
├── High Continuity
│   └── Speech changes relatively slowly
├── Local Similarity  
│   └── Adjacent 10-30ms audio content is similar
└── Perceptual Redundancy
    └── Human ear is insensitive to small temporal differences