Failure-Aware Self-Diagnosing Speech Model

Introduction

30 Second Summary

Voice assistants handle billions of requests daily. When they misunderstand you, the worst outcome is a confident wrong answer instead of simply saying "I didn't catch that."

In this project, you will build a self-diagnosing Spoken Language Understanding model that classifies speech intents and predicts its own failure probability. When the model is uncertain, it abstains instead of guessing wrong.

What You'll Build

A dual-head transformer model that takes raw speech audio, converts it to Mel spectrograms, classifies 31 intents, and outputs a confidence score indicating whether it should predict or abstain.

By the end of this project, you'll have:

• A dual-head transformer model (~1.2M parameters) with prediction, selection, and auxiliary heads trained on the Fluent Speech Commands dataset.
• A SelectiveNet training pipeline with a joint loss function that optimizes classification accuracy and coverage simultaneously.
• A complete selective prediction evaluation suite with coverage-accuracy curves, AUROC for failure detection, and Expected Calibration Error metrics.
• A noise robustness analysis proving the model abstains more under noisy conditions instead of silently failing.
• Secret Mission: Add temperature scaling calibration and compare it against the learned selection head for selective prediction.

Are there any prerequisites?

This project assumes experience building and training models in PyTorch, an understanding of transformer architecture, and familiarity with classification tasks and loss functions. You will also need a Google account for Google Colab GPU access.

Before We Start

Before We Start...

Set Up the Environment and Dataset

Before you can build a self-diagnosing speech model, you need raw audio data and a GPU to train on. In this step, you will set up a Google Colab notebook with GPU acceleration, load the Fluent Speech Commands dataset, and build a data pipeline that converts raw audio waveforms into Mel spectrograms ready for your transformer.

The Fluent Speech Commands dataset contains 30,043 utterances from 97 speakers, each labeled with an intent made up of three slots: action, object, and location. These combine into 31 unique intents that your model will learn to classify.

In this step, get ready to:

• Set up a Google Colab notebook with T4 GPU runtime and verify PyTorch is available.
• Download and load the Fluent Speech Commands dataset with torchaudio.
• Build a custom Dataset class that extracts Mel spectrograms and create DataLoaders.

Create a Colab notebook with GPU runtime

Training a transformer model on audio data requires GPU acceleration. Google Colab provides free access to NVIDIA T4 GPUs, which is more than enough for our ~1.2M parameter model.

• Open Google Colab in your browser.
• Click New notebook to create a fresh notebook.
• Click Runtime in the top menu bar.
• Select Change runtime type.
• Under Hardware accelerator, select T4 GPU, then click Save.

Now verify that PyTorch and torchaudio are available on your runtime.

• Run the following version check in your first notebook cell:

import torch
import torchaudio

print(f"PyTorch: {torch.__version__}")
print(f"torchaudio: {torchaudio.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")

What does this cell do?

This imports PyTorch and torchaudio, then prints their versions and confirms GPU access. You should see CUDA available: True and a device name containing "T4" in the output.

CUDA not available?

If you see CUDA available: False, go back to Runtime > Change runtime type and confirm T4 GPU is selected. You may need to click Connect in the top-right to start the runtime.

help me set up GPU runtime in Google Colab

Download and load the Fluent Speech Commands dataset

The Fluent Speech Commands dataset is a collection of 30,043 spoken commands from 97 speakers, each labeled with an intent. You need to download it from Fluent.ai and upload it to your Colab environment.

Why a manual download?

The torchaudio.datasets.FluentSpeechCommands class does not auto-download the audio files. It expects the dataset to already exist at the path you provide. The dataset is licensed under CC-BY-NC-ND 4.0 for academic research, so Fluent.ai requires you to agree to terms before downloading.

• Visit the Fluent Speech Commands dataset page and follow the instructions to download the dataset archive.
• Upload the downloaded archive to your Colab environment by clicking the folder icon in the left sidebar, then the upload button.
• Extract the dataset into a ./data directory by running this cell:

# Extract the dataset archive to ./data
# Adjust the filename if your download has a different name
!mkdir -p ./data
!unzip -q fluent_speech_commands_dataset.zip -d ./data

What directory structure does torchaudio expect?

The FluentSpeechCommands class looks for data at root/fluent_speech_commands_dataset/data/{subset}_data.csv. After extraction, your ./data folder should contain a fluent_speech_commands_dataset subfolder with the CSV files and audio.

Now verify that torchaudio can load the dataset correctly.

• Load one split and print a sample tuple by running this cell:

# Verify the dataset loads correctly
dataset = torchaudio.datasets.FluentSpeechCommands(root="./data", subset="train")

# Print a sample tuple
sample = dataset[0]
print(f"Waveform shape: {sample[0].shape}")
print(f"Sample rate: {sample[1]}")
print(f"Filename: {sample[2]}")
print(f"Speaker ID: {sample[3]}")
print(f"Transcription: {sample[4]}")
print(f"Action: {sample[5]}")
print(f"Object: {sample[6]}")
print(f"Location: {sample[7]}")

What should I see?

You should see a waveform tensor shape (1-channel audio), a sample rate of 16000, and text fields showing the filename, speaker ID, transcription, and the three intent slots (action, object, location). This confirms torchaudio can read the audio files and metadata correctly.

Seeing RuntimeError: Dataset not found?

This means the extracted folder structure does not match what torchaudio expects. Run !ls ./data/ and confirm you see a fluent_speech_commands_dataset folder. If the dataset extracted into a nested folder, move it so the path is ./data/fluent_speech_commands_dataset/.

help me fix the dataset path for FluentSpeechCommands in torchaudio

Build the data pipeline with Mel spectrograms

Raw waveforms are too long and unstructured for a transformer to process directly. A Mel spectrogram converts audio into a compact time-frequency representation that captures the features humans use to distinguish speech sounds.

You will now create the failure_aware_slu.py file with a custom Dataset class that wraps the raw dataset and outputs fixed-size Mel spectrograms paired with integer intent labels.

• Create a new code cell and paste the following imports and dataset class:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label

What does this class do?

The FluentSLUDataset class does three things for each audio sample:

• Converts the raw waveform into a Mel spectrogram using 80 Mel bands.
• Pads short spectrograms or truncates long ones to exactly 100 time frames. This ensures every sample has the same shape for batching.
• Maps the three-slot intent label (action + object + location) to a single integer using INTENT_MAP, a class-level dictionary built once from the full dataset.

The INTENT_MAP is a class variable (shared across all instances) so it only needs to scan the dataset metadata once, regardless of how many splits you load.

Now create the Mel spectrogram transform and instantiate the DataLoaders.

• In a new cell, set up the transform and create train, validation, and test loaders:

# Configure the Mel spectrogram transform for 16kHz audio
mel_transform = T.MelSpectrogram(
    sample_rate=16000,
    n_fft=512,
    hop_length=160,
    n_mels=80,
)

# Load all three splits
train_dataset = FluentSLUDataset("./data", "train", mel_transform)
val_dataset = FluentSLUDataset("./data", "valid", mel_transform)
test_dataset = FluentSLUDataset("./data", "test", mel_transform)

num_intents = len(FluentSLUDataset.INTENT_MAP)
print(f"Number of intents: {num_intents}")

# Create DataLoaders with batch_size=64
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=2)

What do these parameters mean?

The MelSpectrogram transform parameters are tuned for 16kHz speech:

• n_fft=512: the FFT window size (32ms at 16kHz). This determines frequency resolution.
• hop_length=160: the step between windows (10ms). This determines time resolution.
• n_mels=80: the number of Mel frequency bands. 80 is standard for speech models.

With max_frames=100 and hop_length=160 at 16kHz, each sample covers about 1 second of audio. Most Fluent Speech Commands utterances are under 3 seconds, so truncation keeps the most informative portion.

Finally, verify that batches have the correct shape.

• Print the shape of one batch by running this cell:

# Verify batch shapes
mel_batch, label_batch = next(iter(train_loader))
print(f"Mel spectrogram batch shape: {mel_batch.shape}")
print(f"Label batch shape: {label_batch.shape}")

• Confirm the output shows Mel spectrogram batch shape: torch.Size([64, 80, 100]) and Label batch shape: torch.Size([64]).

What do these dimensions mean?

The shape [64, 80, 100] means: 64 samples per batch, 80 Mel frequency bands, and 100 time frames. The label shape [64] is one integer intent label per sample (values 0 through 30 for the 31 intents).

Batch shape does not match?

If you see a different first dimension (less than 64), your dataset may be smaller than expected. Confirm the dataset loaded all utterances by checking len(train_dataset). You should see 23,132 training samples.

If the Mel dimension is wrong (not 80) or time dimension is wrong (not 100), double-check that n_mels=80 and max_frames=100 are set correctly in your code.

help me debug my Mel spectrogram batch shapes

✔️ Awesome, I've got everything!

Great. Double check you have saved all your cells in order.

ⓧ I'd like to double check the full code

Here is the complete failure_aware_slu.py code you should have at this point:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label

Your data pipeline is ready. Next up, you will build the dual-head transformer model that processes these Mel spectrograms and outputs both intent predictions and confidence scores.

Build the Dual-Head Transformer Model

Your dataset pipeline is ready and producing batches of Mel spectrograms. Now you need a model that can process those spectrograms and do two things at once: classify the intent AND predict whether it should trust its own answer.

In this step, you will build a SelectiveNet-style transformer model with a shared backbone that branches into three output heads. This dual-head architecture is what gives the model its self-diagnosing ability.

In this step, get ready to:

• Build a sinusoidal positional encoding module for the transformer backbone.
• Define the SelectiveSLUModel with prediction, selection, and auxiliary heads.
• Implement the forward pass and verify the model produces correct output shapes.

Build the shared transformer backbone

A Transformer encoder processes sequences but has no inherent sense of order. Without positional information, the model treats the first and last frames of a spectrogram identically. Sinusoidal positional encoding solves this by injecting a unique position signal into each time frame.

• In failure_aware_slu.py, find the end of the FluentSLUDataset class.
• Add the following PositionalEncoding class below it:

# === Model ===

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer."""

    def __init__(self, d_model, max_len=500):
        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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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

What does this code do?

Sinusoidal positional encoding creates a unique pattern for each position using sine and cosine waves at different frequencies. The model learns to use these patterns to understand temporal order in the spectrogram frames.

The register_buffer call stores the encoding as part of the model state without treating it as a trainable parameter. This means it moves to the GPU automatically but does not get updated during backpropagation.

• Save failure_aware_slu.py.
• You won't see a visible change yet. The next chunk uses this class inside the main model.

Define the SelectiveSLUModel with prediction, selection, and auxiliary heads

The SelectiveNet architecture uses three output heads that all share a single transformer backbone. The prediction head classifies intents. The selection head outputs a confidence score indicating whether the model should answer or abstain. The auxiliary head regularizes the shared features during training.

• Add the SelectiveSLUModel class below the PositionalEncoding class by copying this code:

class SelectiveSLUModel(nn.Module):
    """Dual-head transformer for selective spoken language understanding."""

    def __init__(self, n_mels=80, d_model=128, nhead=4, num_layers=4,
                 dim_feedforward=256, dropout=0.1, num_intents=31):
        super().__init__()

        # Shared backbone: project mel bands to model dimension
        self.input_proj = nn.Linear(n_mels, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout)

        # 4-layer transformer encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        # Prediction head f(x): classifies intent
        self.prediction_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, num_intents),
        )

        # Selection head g(x): confidence score for abstention
        self.selection_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        # Auxiliary head h(x): regularizes shared features during training
        self.auxiliary_head = nn.Sequential(
            nn.Linear(d_model, num_intents),
        )

What do the three heads do?

• Prediction head f(x): outputs logits for 31 intent classes. This is the standard classification output.
• Selection head g(x): outputs a single value between 0 and 1 via Sigmoid. A score near 1 means the model is confident. A score near 0 means the model is unsure and should abstain.
• Auxiliary head h(x): a simple linear head used only during training. It forces the backbone to learn features useful for classification even when the selection head would prefer to abstain on hard examples.

Why batch_first=True?

PyTorch's TransformerEncoderLayer defaults to (sequence, batch, feature) input format. Setting batch_first=True switches it to (batch, sequence, feature), which matches how your Mel spectrograms are shaped after transposing.

• Save failure_aware_slu.py.
• You won't see output yet. The next chunk adds the forward pass that connects everything.

Implement the SelectiveNet forward pass

The forward method routes data through the shared backbone, applies global average pooling to collapse the time dimension, then branches to all three heads independently.

• Add the forward method inside the SelectiveSLUModel class, below the __init__ method:

    def forward(self, mel):
        # Transpose from (batch, mels, time) to (batch, time, mels)
        x = mel.transpose(1, 2)
        # Project mel bands to model dimension
        x = self.input_proj(x)
        x = self.pos_encoder(x)
        x = self.dropout(x)

        # Pass through transformer encoder
        x = self.transformer_encoder(x)

        # Global average pooling over time dimension
        x = x.mean(dim=1)

        # Branch to three independent heads
        pred_logits = self.prediction_head(x)
        selection_score = self.selection_head(x)
        aux_logits = self.auxiliary_head(x)

        return pred_logits, selection_score, aux_logits

What's happening here?

• mel.transpose(1, 2) reshapes from [batch, 80, 100] to [batch, 100, 80] so each time frame becomes a sequence element with 80 features.
• input_proj projects each 80-dim frame up to the 128-dim model space.
• x.mean(dim=1) performs global average pooling, collapsing 100 time steps into a single 128-dim vector that represents the entire utterance.
• The three heads then independently process this shared representation, returning classification logits, a confidence score, and auxiliary logits.

• Save failure_aware_slu.py.
• Verify the model works by running this test in a new cell:

# Quick shape verification
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SelectiveSLUModel(n_mels=80, d_model=128, nhead=4, num_layers=4,
                          dim_feedforward=256, dropout=0.1, num_intents=31).to(device)

# Pass a random batch through the model
test_input = torch.randn(64, 80, 100).to(device)
pred_logits, sel_scores, aux_logits = model(test_input)

print(f"Input shape: {test_input.shape}")
print(f"Prediction logits: {pred_logits.shape}")
print(f"Selection scores: {sel_scores.shape}")
print(f"Auxiliary logits: {aux_logits.shape}")
print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")

You should see output confirming the shapes match:

• Input shape: torch.Size([64, 80, 100]).
• Prediction logits: torch.Size([64, 31]).
• Selection scores: torch.Size([64, 1]).
• Auxiliary logits: torch.Size([64, 31]).
• Total parameters around 1.2M.

Shapes don't match or seeing an error?

• If you see a dimension mismatch error, check that the forward method is indented inside the SelectiveSLUModel class, not at the module level.
• If selection scores are not between 0 and 1, make sure nn.Sigmoid() is the last layer in your selection_head Sequential.
• If you get a NameError for PositionalEncoding, make sure that class is defined above SelectiveSLUModel in the same cell or file.

Help me debug my SelectiveSLUModel output shapes.

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before moving on.

ⓧ I'd like to double check the full code

Here is the complete failure_aware_slu.py file so far:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label


# === Model ===

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer."""

    def __init__(self, d_model, max_len=500):
        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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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


class SelectiveSLUModel(nn.Module):
    """Dual-head transformer for selective spoken language understanding."""

    def __init__(self, n_mels=80, d_model=128, nhead=4, num_layers=4,
                 dim_feedforward=256, dropout=0.1, num_intents=31):
        super().__init__()

        self.input_proj = nn.Linear(n_mels, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        self.prediction_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, num_intents),
        )

        self.selection_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        self.auxiliary_head = nn.Sequential(
            nn.Linear(d_model, num_intents),
        )

    def forward(self, mel):
        x = mel.transpose(1, 2)
        x = self.input_proj(x)
        x = self.pos_encoder(x)
        x = self.dropout(x)

        x = self.transformer_encoder(x)

        x = x.mean(dim=1)

        pred_logits = self.prediction_head(x)
        selection_score = self.selection_head(x)
        aux_logits = self.auxiliary_head(x)

        return pred_logits, selection_score, aux_logits

Your dual-head transformer is built and verified. Next up, you will implement the SelectiveNet loss function and train this model to jointly optimize intent accuracy and abstention coverage.

Implement SelectiveNet Loss and Train

Your dual-head model can output predictions and confidence scores, but it doesn't yet know how to learn when to abstain. The secret sauce is the loss function that teaches the selection head to output low scores for uncertain samples.

In this step, you'll implement the SelectiveNet loss that jointly trains classification accuracy and coverage. Then you'll train the model for 30 epochs and save the best checkpoint.

In this step, get ready to:

• Implement the three-component SelectiveNet loss function.
• Build evaluation and training functions with Adam and cosine annealing.
• Train the model and save the best checkpoint based on selective accuracy.

Implement the SelectiveNet loss function

Standard classification loss treats every sample equally. SelectiveNet's loss is different. It weights each sample's loss by that sample's selection score, so the model learns to "give up" on hard examples rather than being forced to classify them.

• Add the SelectiveNetLoss class in a new cell by copying the following code:

# === Loss ===

class SelectiveNetLoss(nn.Module):
    """Joint loss for SelectiveNet: selective CE + coverage penalty + auxiliary."""

    def __init__(self, target_coverage=0.8, lam=32.0, alpha=0.5):
        super().__init__()
        self.target_coverage = target_coverage
        self.lam = lam
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss(reduction="none")

    def forward(self, pred_logits, selection_scores, aux_logits, targets):
        # Per-sample cross-entropy (unreduced)
        ce = self.ce_loss(pred_logits, targets)

        # Flatten selection scores and compute empirical coverage
        selection_scores_flat = selection_scores.squeeze(-1)
        empirical_coverage = selection_scores_flat.mean()

        # Selective CE: weight each loss by its selection score, normalize by coverage
        selective_ce = (ce * selection_scores_flat).sum() / (
            empirical_coverage * len(targets) + 1e-8
        )

        # Quadratic penalty if coverage drops below target
        coverage_penalty = self.lam * torch.clamp(
            self.target_coverage - empirical_coverage, min=0.0
        ) ** 2

        # Standard CE on auxiliary head to regularize features
        aux_loss = self.ce_loss(aux_logits, targets).mean()

        # Combine all three components
        total_loss = selective_ce + coverage_penalty + self.alpha * aux_loss

        return total_loss, empirical_coverage

What does this loss do?

The loss has three components working together:

• Selective cross-entropy multiplies each sample's loss by its selection score g(x). Samples the model wants to reject (low g(x)) contribute less loss, so the model isn't punished for abstaining on hard examples.
• Coverage penalty prevents the model from rejecting everything. The quadratic penalty kicks in when empirical coverage drops below 80%, forcing the model to accept at least 80% of inputs.
• Auxiliary loss trains the auxiliary head with standard cross-entropy. This regularizes the shared backbone so it learns good features even for samples the selection head might reject.

Why lambda=32 and alpha=0.5?

Lambda (lam=32.0) controls how strongly the model is penalized for under-coverage. A value of 32 is standard from the SelectiveNet paper and provides enough pressure to maintain the 80% target without overwhelming the classification signal.

Alpha (alpha=0.5) balances the auxiliary loss against the selective loss. At 0.5, both contribute equally to feature learning.

• Run the cell to define the class.

Seeing a NameError?

Make sure you ran the earlier cells that import torch and torch.nn. The SelectiveNetLoss class inherits from nn.Module, so PyTorch must be imported first.

help me debug a NameError in my SelectiveNetLoss class

Set up the training loop

Before the training function, you need an evaluation helper that computes selective accuracy. This function only counts predictions where the selection head's score is above a threshold.

• Add the evaluate_selective function in a new cell:

# === Evaluation ===

def evaluate_selective(model, loader, device, threshold=0.5):
    """Evaluate selective accuracy and coverage at a given threshold."""
    model.eval()
    correct = 0
    selected = 0
    total = 0

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            # Only count predictions where selection score >= threshold
            preds = pred_logits.argmax(dim=1)
            mask = sel_scores.squeeze(-1) >= threshold

            selected += mask.sum().item()
            correct += (preds[mask] == labels[mask]).sum().item()
            total += len(labels)

    coverage = selected / total if total > 0 else 0.0
    accuracy = correct / selected if selected > 0 else 0.0
    return accuracy, coverage

What does selective accuracy mean?

Traditional accuracy counts every prediction. Selective accuracy only counts predictions the model is confident about (g(x) >= 0.5). The function returns two values: accuracy on accepted samples, and coverage (what fraction of samples were accepted).

If the model accepts 80% of samples and gets 95% accuracy on those, that's much more useful than 85% accuracy on everything. The rejected 20% can be routed to a human or a fallback system.

• Run the cell to define evaluate_selective.
• Now add the train_model function in the next cell:

# === Training ===

def train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3):
    """Train the SelectiveSLU model."""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    criterion = SelectiveNetLoss(target_coverage=0.8, lam=32.0, alpha=0.5)

    best_val_acc = 0.0
    history = {"train_loss": [], "val_acc": [], "val_coverage": []}

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0

        for mel, labels in train_loader:
            mel, labels = mel.to(device), labels.to(device)

            # Forward pass through all three heads
            pred_logits, sel_scores, aux_logits = model(mel)
            loss, coverage = criterion(pred_logits, sel_scores, aux_logits, labels)

            # Backprop and update
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()

        # Validate after each epoch
        val_acc, val_coverage = evaluate_selective(model, val_loader, device)
        history["train_loss"].append(total_loss / len(train_loader))
        history["val_acc"].append(val_acc)
        history["val_coverage"].append(val_coverage)

        if (epoch + 1) % 5 == 0:
            print(
                f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} "
                f"| Val Acc: {val_acc:.4f} | Coverage: {val_coverage:.4f}"
            )

        # Save best model based on selective accuracy
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_selective_slu.pth")

    return history

What does this training loop do?

• Adam optimizer with lr=1e-3 handles parameter updates.
• Cosine annealing scheduler gradually reduces the learning rate over 30 epochs, helping the model converge smoothly.
• Every 5 epochs, the loop prints loss, validation selective accuracy, and coverage so you can monitor training.
• The best checkpoint is saved based on validation selective accuracy (not overall accuracy). This ensures you keep the model that performs best on samples it chooses to accept.

• Run the cell to define train_model.

Seeing a NameError for evaluate_selective?

The train_model function calls evaluate_selective internally. Make sure you ran the cell containing evaluate_selective before this one.

help me fix a NameError in my training loop

Train the model

Everything is wired up. Time to run 30 epochs of training. On a T4 GPU, this takes approximately 15 minutes with the Fluent Speech Commands dataset.

• Start training and load the best checkpoint by running this cell:

# Train for 30 epochs
print("\n--- Training ---")
history = train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3)

# Load the best checkpoint
model.load_state_dict(torch.load("best_selective_slu.pth"))
print("\nBest model loaded.")

What should I see during training?

Every 5 epochs, you'll see a print line showing loss, validation selective accuracy, and coverage. Over 30 epochs:

• Loss should decrease from around 3-5 down to below 1.
• Validation accuracy should climb above 90%.
• Coverage should settle near the 80% target (between 0.75 and 0.85).

• Wait for training to complete. You should see progress printed at epochs 5, 10, 15, 20, 25, and 30.
• Verify the selection head produces distributed scores by running this check:

# Verify selection head produces a distribution of scores
model.eval()
with torch.no_grad():
    sample_mel, sample_labels = next(iter(val_loader))
    sample_mel = sample_mel.to(device)
    _, sel_scores, _ = model(sample_mel)
    print(f"Selection score stats:")
    print(f"  Min: {sel_scores.min().item():.4f}")
    print(f"  Max: {sel_scores.max().item():.4f}")
    print(f"  Mean: {sel_scores.mean().item():.4f}")
    print(f"  Std: {sel_scores.std().item():.4f}")

• Confirm that the selection scores have a standard deviation greater than 0. If min and max are different values, the selection head is successfully learning to differentiate between confident and uncertain predictions.

Selection scores all stuck at 1.0?

If every selection score is 1.0 (or very close), the coverage penalty may be dominating the loss. This means the model learned it's safer to accept everything than to risk the penalty. Try reducing lam from 32.0 to 16.0 and retraining.

Also check that your model is on the GPU. Training on CPU is extremely slow and may not converge in 30 epochs.

help me debug my selection head not learning

✔️ Awesome, I've got everything!

Great. Double check you've run all your cells and the best model is loaded.

ⓧ I'd like to double check the full code

Here is the complete failure_aware_slu.py up to this point:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label


# === Model ===

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer."""

    def __init__(self, d_model, max_len=500):
        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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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


class SelectiveSLUModel(nn.Module):
    """Dual-head transformer for selective spoken language understanding."""

    def __init__(self, n_mels=80, d_model=128, nhead=4, num_layers=4,
                 dim_feedforward=256, dropout=0.1, num_intents=31):
        super().__init__()

        self.input_proj = nn.Linear(n_mels, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        self.prediction_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, num_intents),
        )

        self.selection_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        self.auxiliary_head = nn.Sequential(
            nn.Linear(d_model, num_intents),
        )

    def forward(self, mel):
        x = mel.transpose(1, 2)
        x = self.input_proj(x)
        x = self.pos_encoder(x)
        x = self.dropout(x)

        x = self.transformer_encoder(x)

        x = x.mean(dim=1)

        pred_logits = self.prediction_head(x)
        selection_score = self.selection_head(x)
        aux_logits = self.auxiliary_head(x)

        return pred_logits, selection_score, aux_logits


# === Loss ===

class SelectiveNetLoss(nn.Module):
    """Joint loss for SelectiveNet: selective CE + coverage penalty + auxiliary."""

    def __init__(self, target_coverage=0.8, lam=32.0, alpha=0.5):
        super().__init__()
        self.target_coverage = target_coverage
        self.lam = lam
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss(reduction="none")

    def forward(self, pred_logits, selection_scores, aux_logits, targets):
        ce = self.ce_loss(pred_logits, targets)

        selection_scores_flat = selection_scores.squeeze(-1)
        empirical_coverage = selection_scores_flat.mean()

        selective_ce = (ce * selection_scores_flat).sum() / (
            empirical_coverage * len(targets) + 1e-8
        )

        coverage_penalty = self.lam * torch.clamp(
            self.target_coverage - empirical_coverage, min=0.0
        ) ** 2

        aux_loss = self.ce_loss(aux_logits, targets).mean()

        total_loss = selective_ce + coverage_penalty + self.alpha * aux_loss

        return total_loss, empirical_coverage


# === Training ===

def train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3):
    """Train the SelectiveSLU model."""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    criterion = SelectiveNetLoss(target_coverage=0.8, lam=32.0, alpha=0.5)

    best_val_acc = 0.0
    history = {"train_loss": [], "val_acc": [], "val_coverage": []}

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0

        for mel, labels in train_loader:
            mel, labels = mel.to(device), labels.to(device)

            pred_logits, sel_scores, aux_logits = model(mel)
            loss, coverage = criterion(pred_logits, sel_scores, aux_logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()

        val_acc, val_coverage = evaluate_selective(model, val_loader, device)
        history["train_loss"].append(total_loss / len(train_loader))
        history["val_acc"].append(val_acc)
        history["val_coverage"].append(val_coverage)

        if (epoch + 1) % 5 == 0:
            print(
                f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} "
                f"| Val Acc: {val_acc:.4f} | Coverage: {val_coverage:.4f}"
            )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_selective_slu.pth")

    return history


# === Evaluation ===

def evaluate_selective(model, loader, device, threshold=0.5):
    """Evaluate selective accuracy and coverage at a given threshold."""
    model.eval()
    correct = 0
    selected = 0
    total = 0

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            mask = sel_scores.squeeze(-1) >= threshold

            selected += mask.sum().item()
            correct += (preds[mask] == labels[mask]).sum().item()
            total += len(labels)

    coverage = selected / total if total > 0 else 0.0
    accuracy = correct / selected if selected > 0 else 0.0
    return accuracy, coverage

Your model now knows when to speak up and when to stay quiet. Next, you'll put that ability to the test with rigorous evaluation metrics that quantify exactly how well the selection head separates correct from incorrect predictions.

Evaluate with Selective Prediction Metrics

Your model is trained and the selection head is producing a distribution of confidence scores. Now comes the critical question: does the selection head actually separate correct predictions from incorrect ones?

In this step, you will build a complete evaluation suite with production-relevant metrics. You will sweep confidence thresholds to plot coverage-accuracy curves, compute AUROC for failure detection, and measure calibration error. These are the same metrics Amazon uses to validate their FPI system before deploying to live traffic.

In this step, get ready to:

• Implement and plot a coverage-accuracy curve showing the accuracy-abstention trade-off.
• Compute AUROC and Expected Calibration Error to measure selection quality.
• Generate per-intent analysis to identify which intents the model finds hardest.

Implement coverage-accuracy curve evaluation

A coverage-accuracy curve answers the question: if the model is allowed to abstain on its least confident predictions, how accurate is it on the rest? By sweeping the threshold from 0.0 to 1.0, you trace the trade-off between how many samples the model covers and how accurate it is on those it keeps.

• In a new cell in your Colab notebook, add the compute_coverage_accuracy_curve function by pasting this code:

def compute_coverage_accuracy_curve(model, loader, device):
    """Sweep thresholds to build coverage-accuracy curve."""
    model.eval()
    all_preds = []
    all_labels = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            all_preds.append(preds.cpu())
            all_labels.append(labels.cpu())
            all_scores.append(sel_scores.squeeze(-1).cpu())

    all_preds = torch.cat(all_preds)
    all_labels = torch.cat(all_labels)
    all_scores = torch.cat(all_scores)

    thresholds = np.arange(0.0, 1.01, 0.05)
    coverages = []
    accuracies = []

    for tau in thresholds:
        mask = all_scores >= tau
        coverage = mask.float().mean().item()
        if mask.sum() > 0:
            acc = (all_preds[mask] == all_labels[mask]).float().mean().item()
        else:
            acc = 0.0
        coverages.append(coverage)
        accuracies.append(acc)

    return coverages, accuracies, thresholds

What does this code do?

The function collects all predictions, labels, and selection scores from the test set in one pass. It then sweeps a threshold tau from 0.0 to 1.0 in steps of 0.05.

At each threshold, it keeps only samples where the selection score is at or above tau, then measures accuracy on those kept samples. Higher thresholds mean fewer samples kept (lower coverage) but higher accuracy on the remaining ones.

• Now call the function and plot the result in a new cell:

# Run the coverage-accuracy sweep on the test set
coverages, accuracies, thresholds = compute_coverage_accuracy_curve(
    model, test_loader, device
)

# Plot and save the coverage-accuracy curve
plt.figure(figsize=(8, 5))
plt.plot(coverages, accuracies, "b-o", markersize=4)
plt.xlabel("Coverage")
plt.ylabel("Selective Accuracy")
plt.title("Coverage vs. Accuracy Curve")
plt.grid(True)
plt.savefig("coverage_accuracy_curve.png", dpi=150, bbox_inches="tight")
plt.show()

What should I see?

You should see a curve that starts at the bottom-right (100% coverage, baseline accuracy) and moves toward the top-left (lower coverage, higher accuracy). The curve demonstrates that as the model rejects its least confident predictions, accuracy on the remaining ones climbs.

• Run the cell. You should see a plot with coverage on the x-axis and selective accuracy on the y-axis, trending upward as coverage decreases.

Curve looks flat or noisy?

If the curve is completely flat, your selection head may have collapsed to outputting a constant score. Check that your trained model's selection scores have variance by printing all_scores.std(). If it is near zero, re-train with a higher lambda value.

Make sure you loaded the best checkpoint with model.load_state_dict(torch.load("best_selective_slu.pth")) before running evaluation.

Help me debug my flat coverage-accuracy curve

Compute AUROC and Expected Calibration Error

The coverage-accuracy curve shows the trade-off visually. Now you need single-number metrics that quantify how well the selection head works. AUROC measures how well the selection score separates correct predictions from incorrect ones. Expected Calibration Error (ECE) measures whether the confidence scores match actual accuracy.

• Add the compute_auroc_failure_detection function in a new cell:

def compute_auroc_failure_detection(model, loader, device):
    """Compute AUROC: can selection score separate correct from incorrect?"""
    model.eval()
    all_correct = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            correct = (preds == labels).cpu().numpy()
            scores = sel_scores.squeeze(-1).cpu().numpy()

            all_correct.append(correct)
            all_scores.append(scores)

    all_correct = np.concatenate(all_correct)
    all_scores = np.concatenate(all_scores)

    auroc = roc_auc_score(all_correct, all_scores)
    return auroc, all_correct, all_scores

What does this code do?

This function treats the selection score as a binary classifier for the question: "will this prediction be correct?" It collects all selection scores and whether each prediction was actually correct, then uses roc_auc_score from scikit-learn to measure how well the scores separate correct from incorrect.

An AUROC of 0.5 means random separation (the head is useless). An AUROC of 1.0 means perfect separation. You want above 0.75 to confirm the selection head has learned meaningful confidence patterns.

• Add the compute_ece function in the next cell:

def compute_ece(confidences, accuracies_bin, n_bins=10):
    """Compute Expected Calibration Error."""
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    total_samples = len(confidences)

    for i in range(n_bins):
        mask = (confidences >= bin_boundaries[i]) & (confidences < bin_boundaries[i + 1])
        if mask.sum() > 0:
            bin_conf = confidences[mask].mean()
            bin_acc = accuracies_bin[mask].mean()
            ece += (mask.sum() / total_samples) * abs(bin_acc - bin_conf)

    return ece

What does ECE measure?

ECE divides the confidence range [0, 1] into 10 bins. For each bin, it compares the mean confidence of samples in that bin to the actual accuracy of those samples. A perfectly calibrated model has ECE = 0, meaning a confidence of 0.8 corresponds to 80% accuracy.

Lower ECE means the model's confidence scores are trustworthy. This matters in production because downstream systems need to make routing decisions based on these scores.

• Call both functions and print the results in a new cell:

# Compute AUROC for failure detection
auroc, all_correct, all_scores = compute_auroc_failure_detection(
    model, test_loader, device
)
print(f"AUROC (failure detection): {auroc:.4f}")

# Compute Expected Calibration Error
ece = compute_ece(all_scores, all_correct.astype(float))
print(f"Expected Calibration Error: {ece:.4f}")

• Run the cell. You should see an AUROC value above 0.75 and an ECE value printed below it.

What do these numbers mean?

An AUROC above 0.75 confirms your selection head separates correct from incorrect predictions significantly better than random chance. The ECE tells you how well-calibrated those confidence scores are. Typical values range from 0.02 (well-calibrated) to 0.15 (poorly calibrated).

AUROC below 0.75?

If your AUROC is near 0.5, the selection head is not learning to distinguish correct from incorrect predictions. Try increasing the training epochs to 40 or increasing the coverage penalty lambda from 32.0 to 64.0 to give the selection head a stronger training signal.

Help me improve my selection head's AUROC

✔️ Awesome, I've got everything!

Great. Double check you have all three evaluation functions defined in your notebook cells.

ⓧ I'd like to double check the full code

Here is the complete code including all functions defined so far:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label


# === Model ===

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer."""

    def __init__(self, d_model, max_len=500):
        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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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


class SelectiveSLUModel(nn.Module):
    """Dual-head transformer for selective spoken language understanding."""

    def __init__(self, n_mels=80, d_model=128, nhead=4, num_layers=4,
                 dim_feedforward=256, dropout=0.1, num_intents=31):
        super().__init__()

        self.input_proj = nn.Linear(n_mels, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        self.prediction_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, num_intents),
        )

        self.selection_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        self.auxiliary_head = nn.Sequential(
            nn.Linear(d_model, num_intents),
        )

    def forward(self, mel):
        x = mel.transpose(1, 2)
        x = self.input_proj(x)
        x = self.pos_encoder(x)
        x = self.dropout(x)

        x = self.transformer_encoder(x)

        x = x.mean(dim=1)

        pred_logits = self.prediction_head(x)
        selection_score = self.selection_head(x)
        aux_logits = self.auxiliary_head(x)

        return pred_logits, selection_score, aux_logits


# === Loss ===

class SelectiveNetLoss(nn.Module):
    """Joint loss for SelectiveNet: selective CE + coverage penalty + auxiliary."""

    def __init__(self, target_coverage=0.8, lam=32.0, alpha=0.5):
        super().__init__()
        self.target_coverage = target_coverage
        self.lam = lam
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss(reduction="none")

    def forward(self, pred_logits, selection_scores, aux_logits, targets):
        ce = self.ce_loss(pred_logits, targets)

        selection_scores_flat = selection_scores.squeeze(-1)
        empirical_coverage = selection_scores_flat.mean()

        selective_ce = (ce * selection_scores_flat).sum() / (
            empirical_coverage * len(targets) + 1e-8
        )

        coverage_penalty = self.lam * torch.clamp(
            self.target_coverage - empirical_coverage, min=0.0
        ) ** 2

        aux_loss = self.ce_loss(aux_logits, targets).mean()

        total_loss = selective_ce + coverage_penalty + self.alpha * aux_loss

        return total_loss, empirical_coverage


# === Training ===

def train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3):
    """Train the SelectiveSLU model."""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    criterion = SelectiveNetLoss(target_coverage=0.8, lam=32.0, alpha=0.5)

    best_val_acc = 0.0
    history = {"train_loss": [], "val_acc": [], "val_coverage": []}

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0

        for mel, labels in train_loader:
            mel, labels = mel.to(device), labels.to(device)

            pred_logits, sel_scores, aux_logits = model(mel)
            loss, coverage = criterion(pred_logits, sel_scores, aux_logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()

        val_acc, val_coverage = evaluate_selective(model, val_loader, device)
        history["train_loss"].append(total_loss / len(train_loader))
        history["val_acc"].append(val_acc)
        history["val_coverage"].append(val_coverage)

        if (epoch + 1) % 5 == 0:
            print(
                f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} "
                f"| Val Acc: {val_acc:.4f} | Coverage: {val_coverage:.4f}"
            )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_selective_slu.pth")

    return history


# === Evaluation ===

def evaluate_selective(model, loader, device, threshold=0.5):
    """Evaluate selective accuracy and coverage at a given threshold."""
    model.eval()
    correct = 0
    selected = 0
    total = 0

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            mask = sel_scores.squeeze(-1) >= threshold

            selected += mask.sum().item()
            correct += (preds[mask] == labels[mask]).sum().item()
            total += len(labels)

    coverage = selected / total if total > 0 else 0.0
    accuracy = correct / selected if selected > 0 else 0.0
    return accuracy, coverage


def compute_coverage_accuracy_curve(model, loader, device):
    """Sweep thresholds to build coverage-accuracy curve."""
    model.eval()
    all_preds = []
    all_labels = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            all_preds.append(preds.cpu())
            all_labels.append(labels.cpu())
            all_scores.append(sel_scores.squeeze(-1).cpu())

    all_preds = torch.cat(all_preds)
    all_labels = torch.cat(all_labels)
    all_scores = torch.cat(all_scores)

    thresholds = np.arange(0.0, 1.01, 0.05)
    coverages = []
    accuracies = []

    for tau in thresholds:
        mask = all_scores >= tau
        coverage = mask.float().mean().item()
        if mask.sum() > 0:
            acc = (all_preds[mask] == all_labels[mask]).float().mean().item()
        else:
            acc = 0.0
        coverages.append(coverage)
        accuracies.append(acc)

    return coverages, accuracies, thresholds


def compute_auroc_failure_detection(model, loader, device):
    """Compute AUROC: can selection score separate correct from incorrect?"""
    model.eval()
    all_correct = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            correct = (preds == labels).cpu().numpy()
            scores = sel_scores.squeeze(-1).cpu().numpy()

            all_correct.append(correct)
            all_scores.append(scores)

    all_correct = np.concatenate(all_correct)
    all_scores = np.concatenate(all_scores)

    auroc = roc_auc_score(all_correct, all_scores)
    return auroc, all_correct, all_scores


def compute_ece(confidences, accuracies_bin, n_bins=10):
    """Compute Expected Calibration Error."""
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    total_samples = len(confidences)

    for i in range(n_bins):
        mask = (confidences >= bin_boundaries[i]) & (confidences < bin_boundaries[i + 1])
        if mask.sum() > 0:
            bin_conf = confidences[mask].mean()
            bin_acc = accuracies_bin[mask].mean()
            ece += (mask.sum() / total_samples) * abs(bin_acc - bin_conf)

    return ece

Generate per-intent analysis

Aggregate metrics tell you the model works overall. Per-intent analysis reveals which specific intents the model struggles with. Some intents may sound similar (like "increase volume" vs "increase heat"), causing confusion even when the model is confident.

• Generate a confusion matrix on accepted predictions by running this in a new cell:

# Gather predictions and selection scores for per-intent analysis
model.eval()
all_preds = []
all_labels = []
all_sel_scores = []

with torch.no_grad():
    for mel, labels in test_loader:
        mel, labels = mel.to(device), labels.to(device)
        pred_logits, sel_scores, _ = model(mel)
        all_preds.append(pred_logits.argmax(dim=1).cpu())
        all_labels.append(labels.cpu())
        all_sel_scores.append(sel_scores.squeeze(-1).cpu())

all_preds = torch.cat(all_preds).numpy()
all_labels = torch.cat(all_labels).numpy()
all_sel_scores = torch.cat(all_sel_scores).numpy()

# Confusion matrix on accepted predictions only (selection score >= 0.5)
accepted_mask = all_sel_scores >= 0.5
cm = confusion_matrix(all_labels[accepted_mask], all_preds[accepted_mask])

plt.figure(figsize=(10, 8))
plt.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
plt.title("Confusion Matrix (Accepted Predictions Only)")
plt.colorbar()
plt.xlabel("Predicted Intent")
plt.ylabel("True Intent")
plt.tight_layout()
plt.savefig("confusion_matrix_accepted.png", dpi=150, bbox_inches="tight")
plt.show()

What does this show?

The confusion matrix only includes samples the model chose to keep (selection score >= 0.5). A strong diagonal means the model is accurate on the predictions it is confident about. Off-diagonal clusters reveal which intent pairs get confused even after abstention filtering.

• Now generate a bar chart of abstention rate per intent in a new cell:

# Compute abstention rate per intent
intent_names = sorted(FluentSLUDataset.INTENT_MAP.keys())
num_intents = len(intent_names)
abstention_rates = []

for intent_idx in range(num_intents):
    # Find all test samples belonging to this intent
    intent_mask = all_labels == intent_idx
    if intent_mask.sum() > 0:
        # Fraction of samples where model abstains (score < 0.5)
        abstained = (all_sel_scores[intent_mask] < 0.5).sum()
        rate = abstained / intent_mask.sum()
    else:
        rate = 0.0
    abstention_rates.append(rate)

# Plot bar chart showing which intents are hardest
plt.figure(figsize=(14, 5))
plt.bar(range(num_intents), abstention_rates)
plt.xlabel("Intent Index")
plt.ylabel("Abstention Rate")
plt.title("Abstention Rate per Intent")
plt.tight_layout()
plt.savefig("abstention_per_intent.png", dpi=150, bbox_inches="tight")
plt.show()

What does the abstention chart reveal?

Tall bars indicate intents the model finds most confusing. These are the intents where the selection head says "I don't know" most often. In a production system like Alexa, these would be the intents you prioritize for additional training data or specialized handling.

• Run the cell. You should see a bar chart where some intents have noticeably higher abstention rates than others.

All bars the same height?

If every intent has roughly the same abstention rate, the selection head may not be learning intent-specific uncertainty. This can happen if the coverage penalty dominates and forces a uniform selection distribution. Check that your trained model's overall accuracy is above 90%. If it is much lower, the model has not learned the task well enough for the selection head to differentiate.

Help me understand why my abstention rates are uniform

You now have a complete evaluation suite proving your model can separate correct from incorrect predictions. Next up, you will stress-test this capability by injecting noise and measuring whether the model abstains more when conditions get harder.

Analyze Noise Robustness

Your evaluation metrics prove the selection head separates correct from incorrect predictions on clean audio. But real-world speech systems encounter background noise, poor microphones, and environmental interference constantly.

In this step, you will inject controlled noise into test audio and measure whether the model abstains more when conditions get harder. A model that gives wrong answers confidently under noise is dangerous. A model that says "I don't know" more often under noise is robust.

In this step, get ready to:

• Implement noise injection at multiple SNR levels.
• Evaluate the model on progressively noisier audio.
• Prove graceful degradation through selective gain under noise.

Implement noise injection utilities

Signal-to-Noise Ratio (SNR) measures how much louder the speech signal is compared to background noise. At 20dB SNR, the signal dominates. At 0dB, noise power equals signal power.

• In failure_aware_slu.py, find the compute_ece function you wrote in the previous step.
• Below it, add a noise robustness section with the noise injection function and noisy dataset wrapper by copying the following code:

# === Noise Robustness ===

def add_noise_snr(waveform, snr_db):
    """Add Gaussian noise at a specified SNR level."""
    signal_power = waveform.pow(2).mean()
    snr_linear = 10 ** (snr_db / 10)
    noise_power = signal_power / snr_linear
    noise = torch.randn_like(waveform) * torch.sqrt(noise_power)
    return waveform + noise


class NoisyFluentSLUDataset(Dataset):
    """FluentSLU dataset with noise injection at a specified SNR."""

    def __init__(self, base_dataset, snr_db):
        self.base_dataset = base_dataset
        self.snr_db = snr_db

    def __len__(self):
        return len(self.base_dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = (
            self.base_dataset.dataset[idx]
        )

        noisy_waveform = add_noise_snr(waveform, self.snr_db)

        mel = self.base_dataset.mel_transform(noisy_waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.base_dataset.max_frames:
            mel = mel[:, :self.base_dataset.max_frames]
        elif mel.shape[1] < self.base_dataset.max_frames:
            pad_size = self.base_dataset.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label

What does this code do?

add_noise_snr converts an SNR value in decibels into a noise level, generates Gaussian noise at that power, and adds it to the waveform. Higher SNR means less noise relative to the signal.

NoisyFluentSLUDataset wraps your existing FluentSLUDataset and injects noise at a fixed SNR before computing the Mel spectrogram. This means the noise goes through the same feature pipeline as the signal, simulating realistic microphone noise.

Why inject noise before the Mel spectrogram?

Real-world noise arrives as raw audio, not as spectrogram artifacts. By adding noise to the waveform first, the Mel spectrogram naturally reflects how noise actually corrupts speech features. This produces a more realistic evaluation than adding noise to spectrograms directly.

• Save failure_aware_slu.py.

Seeing a shape mismatch error?

Make sure NoisyFluentSLUDataset accesses the raw dataset through self.base_dataset.dataset[idx] (the underlying torchaudio dataset), not self.base_dataset[idx] (which would return an already-processed Mel spectrogram).

Help me debug a shape mismatch in my NoisyFluentSLUDataset class

Build the main function to run the full pipeline

The main() function orchestrates the entire pipeline: loading data, training, evaluating on clean audio, and testing under noise. It ties together every function you have written across all previous steps.

• Below the NoisyFluentSLUDataset class, add the main execution section by copying this code:

# === Main Execution ===

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    mel_transform = T.MelSpectrogram(
        sample_rate=16000,
        n_fft=512,
        hop_length=160,
        n_mels=80,
    )

    print("Loading Fluent Speech Commands dataset...")
    train_dataset = FluentSLUDataset("./data", "train", mel_transform)
    val_dataset = FluentSLUDataset("./data", "valid", mel_transform)
    test_dataset = FluentSLUDataset("./data", "test", mel_transform)

    num_intents = len(FluentSLUDataset.INTENT_MAP)
    print(f"Number of intents: {num_intents}")

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=2)

    model = SelectiveSLUModel(
        n_mels=80, d_model=128, nhead=4, num_layers=4,
        dim_feedforward=256, dropout=0.1, num_intents=num_intents
    ).to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params:,}")

    print("\n--- Training ---")
    history = train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3)

    model.load_state_dict(torch.load("best_selective_slu.pth"))

    print("\n--- Evaluation ---")
    coverages, accuracies, thresholds = compute_coverage_accuracy_curve(
        model, test_loader, device
    )

    plt.figure(figsize=(8, 5))
    plt.plot(coverages, accuracies, "b-o", markersize=4)
    plt.xlabel("Coverage")
    plt.ylabel("Selective Accuracy")
    plt.title("Coverage vs. Accuracy Curve")
    plt.grid(True)
    plt.savefig("coverage_accuracy_curve.png", dpi=150, bbox_inches="tight")
    plt.show()

    auroc, all_correct, all_scores = compute_auroc_failure_detection(
        model, test_loader, device
    )
    print(f"AUROC (failure detection): {auroc:.4f}")

    ece = compute_ece(all_scores, all_correct.astype(float))
    print(f"Expected Calibration Error: {ece:.4f}")

    print("\n--- Noise Robustness ---")
    snr_levels = [20, 10, 5, 0]
    overall_accs = []
    selective_accs = []

    for snr in snr_levels:
        noisy_dataset = NoisyFluentSLUDataset(test_dataset, snr_db=snr)
        noisy_loader = DataLoader(noisy_dataset, batch_size=64, shuffle=False, num_workers=2)

        sel_acc, coverage = evaluate_selective(model, noisy_loader, device, threshold=0.5)

        model.eval()
        correct_total = 0
        count_total = 0
        with torch.no_grad():
            for mel, labels in noisy_loader:
                mel, labels = mel.to(device), labels.to(device)
                pred_logits, _, _ = model(mel)
                preds = pred_logits.argmax(dim=1)
                correct_total += (preds == labels).sum().item()
                count_total += len(labels)

        overall_acc = correct_total / count_total
        selective_gain = sel_acc - overall_acc

        overall_accs.append(overall_acc)
        selective_accs.append(sel_acc)

        print(
            f"SNR {snr:3d}dB | Overall: {overall_acc:.4f} | "
            f"Selective: {sel_acc:.4f} | Gain: {selective_gain:.4f}"
        )

    plt.figure(figsize=(8, 5))
    plt.plot(snr_levels, overall_accs, "r-o", label="Overall Accuracy")
    plt.plot(snr_levels, selective_accs, "g-o", label="Selective Accuracy")
    plt.xlabel("SNR (dB)")
    plt.ylabel("Accuracy")
    plt.title("Noise Robustness: Overall vs. Selective Accuracy")
    plt.legend()
    plt.grid(True)
    plt.gca().invert_xaxis()
    plt.savefig("noise_robustness.png", dpi=150, bbox_inches="tight")
    plt.show()

    print("\nDone! Model and plots saved.")


if __name__ == "__main__":
    main()

What does this code do?

• The first section sets up the GPU device, Mel spectrogram transform, datasets, DataLoaders, and model. It then trains for 30 epochs and loads the best checkpoint.
• The evaluation section computes the coverage-accuracy curve, AUROC, and ECE on clean test audio.
• The noise robustness section loops through SNR levels [20, 10, 5, 0] dB. For each level, it creates a noisy dataset, computes selective accuracy (predictions where the selection head scores >= 0.5), overall accuracy (all predictions forced), and the selective gain (the difference).
• selective_gain = sel_acc - overall_acc measures how much the model benefits from being allowed to abstain. A positive and increasing gain under worsening noise proves the model degrades gracefully.
• The final plot shows overall accuracy (red) dropping faster than selective accuracy (green) as noise increases, visually proving the model "knows what it doesn't know."

• Save failure_aware_slu.py.

Getting a CUDA out of memory error?

The noise evaluation creates new DataLoaders for each SNR level. If GPU memory is tight, reduce batch_size from 64 to 32 in the noisy DataLoader lines. You can also add torch.cuda.empty_cache() at the start of each SNR iteration.

Help me fix a CUDA out of memory error during the noise robustness evaluation

Run the full pipeline and validate selective gain

Selective gain is the key metric that proves your model is production-worthy. It answers: "Does the model benefit more from abstention as conditions worsen?"

• Run the complete script by executing this in a Colab cell:

main()

What should I see in the noise robustness output?

The output table prints four rows, one per SNR level. Look for these patterns:

• Overall accuracy drops as SNR decreases (noise makes classification harder).
• Selective accuracy stays higher than overall at every SNR level.
• Selective gain (the "Gain" column) is positive at all SNR levels.
• Selective gain increases as noise gets worse (from 20dB down to 0dB).

The increasing gain proves the model benefits MORE from abstention in harder conditions. This is exactly what Amazon's FPI system needs: instead of silently propagating errors, the model flags its own uncertainty.

• After the script finishes, verify that noise_robustness.png was saved. The plot should show the red line (overall accuracy) dropping faster than the green line (selective accuracy) as SNR decreases from left to right.

Why does the x-axis go from high to low SNR?

The plt.gca().invert_xaxis() call flips the x-axis so the plot reads left-to-right as "clean to noisy." This makes the degradation pattern visually intuitive: conditions get harder as you move right, and the gap between the two lines widens.

Selective gain is negative or not increasing?

Check that the model loaded the best checkpoint before the noise evaluation. The line model.load_state_dict(torch.load("best_selective_slu.pth")) must execute before the noise loop. Also confirm training ran for the full 30 epochs. An under-trained selection head produces uninformative scores.

If the selection scores are all very close to 1.0, the coverage penalty may be dominating. This means the selection head learned to always predict rather than learning to discriminate. Try re-running training with a lower lam value (e.g., 16.0 instead of 32.0).

Help me debug why my selective gain is not positive or increasing under noise

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before continuing.

ⓧ I'd like to double check the full code

Here is the complete failure_aware_slu.py file for comparison:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T
import numpy as np
import math
from sklearn.metrics import roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt


# === Dataset ===

class FluentSLUDataset(Dataset):
    """Wraps FluentSpeechCommands with Mel spectrogram extraction."""

    INTENT_MAP = None

    def __init__(self, root, subset, mel_transform, max_frames=100):
        self.dataset = torchaudio.datasets.FluentSpeechCommands(
            root=root, subset=subset
        )
        self.mel_transform = mel_transform
        self.max_frames = max_frames

        if FluentSLUDataset.INTENT_MAP is None:
            intents = set()
            for i in range(len(self.dataset)):
                meta = self.dataset.get_metadata(i)
                intent = f"{meta[5]}_{meta[6]}_{meta[7]}"
                intents.add(intent)
            FluentSLUDataset.INTENT_MAP = {
                intent: idx for idx, intent in enumerate(sorted(intents))
            }

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = self.dataset[idx]

        mel = self.mel_transform(waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.max_frames:
            mel = mel[:, :self.max_frames]
        elif mel.shape[1] < self.max_frames:
            pad_size = self.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label


# === Model ===

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer."""

    def __init__(self, d_model, max_len=500):
        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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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


class SelectiveSLUModel(nn.Module):
    """Dual-head transformer for selective spoken language understanding."""

    def __init__(self, n_mels=80, d_model=128, nhead=4, num_layers=4,
                 dim_feedforward=256, dropout=0.1, num_intents=31):
        super().__init__()

        self.input_proj = nn.Linear(n_mels, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        self.prediction_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, num_intents),
        )

        self.selection_head = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        self.auxiliary_head = nn.Sequential(
            nn.Linear(d_model, num_intents),
        )

    def forward(self, mel):
        x = mel.transpose(1, 2)
        x = self.input_proj(x)
        x = self.pos_encoder(x)
        x = self.dropout(x)

        x = self.transformer_encoder(x)

        x = x.mean(dim=1)

        pred_logits = self.prediction_head(x)
        selection_score = self.selection_head(x)
        aux_logits = self.auxiliary_head(x)

        return pred_logits, selection_score, aux_logits


# === Loss ===

class SelectiveNetLoss(nn.Module):
    """Joint loss for SelectiveNet: selective CE + coverage penalty + auxiliary."""

    def __init__(self, target_coverage=0.8, lam=32.0, alpha=0.5):
        super().__init__()
        self.target_coverage = target_coverage
        self.lam = lam
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss(reduction="none")

    def forward(self, pred_logits, selection_scores, aux_logits, targets):
        ce = self.ce_loss(pred_logits, targets)

        selection_scores_flat = selection_scores.squeeze(-1)
        empirical_coverage = selection_scores_flat.mean()

        selective_ce = (ce * selection_scores_flat).sum() / (
            empirical_coverage * len(targets) + 1e-8
        )

        coverage_penalty = self.lam * torch.clamp(
            self.target_coverage - empirical_coverage, min=0.0
        ) ** 2

        aux_loss = self.ce_loss(aux_logits, targets).mean()

        total_loss = selective_ce + coverage_penalty + self.alpha * aux_loss

        return total_loss, empirical_coverage


# === Training ===

def train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3):
    """Train the SelectiveSLU model."""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    criterion = SelectiveNetLoss(target_coverage=0.8, lam=32.0, alpha=0.5)

    best_val_acc = 0.0
    history = {"train_loss": [], "val_acc": [], "val_coverage": []}

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0

        for mel, labels in train_loader:
            mel, labels = mel.to(device), labels.to(device)

            pred_logits, sel_scores, aux_logits = model(mel)
            loss, coverage = criterion(pred_logits, sel_scores, aux_logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()

        val_acc, val_coverage = evaluate_selective(model, val_loader, device)
        history["train_loss"].append(total_loss / len(train_loader))
        history["val_acc"].append(val_acc)
        history["val_coverage"].append(val_coverage)

        if (epoch + 1) % 5 == 0:
            print(
                f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} "
                f"| Val Acc: {val_acc:.4f} | Coverage: {val_coverage:.4f}"
            )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_selective_slu.pth")

    return history


# === Evaluation ===

def evaluate_selective(model, loader, device, threshold=0.5):
    """Evaluate selective accuracy and coverage at a given threshold."""
    model.eval()
    correct = 0
    selected = 0
    total = 0

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            mask = sel_scores.squeeze(-1) >= threshold

            selected += mask.sum().item()
            correct += (preds[mask] == labels[mask]).sum().item()
            total += len(labels)

    coverage = selected / total if total > 0 else 0.0
    accuracy = correct / selected if selected > 0 else 0.0
    return accuracy, coverage


def compute_coverage_accuracy_curve(model, loader, device):
    """Sweep thresholds to build coverage-accuracy curve."""
    model.eval()
    all_preds = []
    all_labels = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            all_preds.append(preds.cpu())
            all_labels.append(labels.cpu())
            all_scores.append(sel_scores.squeeze(-1).cpu())

    all_preds = torch.cat(all_preds)
    all_labels = torch.cat(all_labels)
    all_scores = torch.cat(all_scores)

    thresholds = np.arange(0.0, 1.01, 0.05)
    coverages = []
    accuracies = []

    for tau in thresholds:
        mask = all_scores >= tau
        coverage = mask.float().mean().item()
        if mask.sum() > 0:
            acc = (all_preds[mask] == all_labels[mask]).float().mean().item()
        else:
            acc = 0.0
        coverages.append(coverage)
        accuracies.append(acc)

    return coverages, accuracies, thresholds


def compute_auroc_failure_detection(model, loader, device):
    """Compute AUROC: can selection score separate correct from incorrect?"""
    model.eval()
    all_correct = []
    all_scores = []

    with torch.no_grad():
        for mel, labels in loader:
            mel, labels = mel.to(device), labels.to(device)
            pred_logits, sel_scores, _ = model(mel)

            preds = pred_logits.argmax(dim=1)
            correct = (preds == labels).cpu().numpy()
            scores = sel_scores.squeeze(-1).cpu().numpy()

            all_correct.append(correct)
            all_scores.append(scores)

    all_correct = np.concatenate(all_correct)
    all_scores = np.concatenate(all_scores)

    auroc = roc_auc_score(all_correct, all_scores)
    return auroc, all_correct, all_scores


def compute_ece(confidences, accuracies_bin, n_bins=10):
    """Compute Expected Calibration Error."""
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    total_samples = len(confidences)

    for i in range(n_bins):
        mask = (confidences >= bin_boundaries[i]) & (confidences < bin_boundaries[i + 1])
        if mask.sum() > 0:
            bin_conf = confidences[mask].mean()
            bin_acc = accuracies_bin[mask].mean()
            ece += (mask.sum() / total_samples) * abs(bin_acc - bin_conf)

    return ece


# === Noise Robustness ===

def add_noise_snr(waveform, snr_db):
    """Add Gaussian noise at a specified SNR level."""
    signal_power = waveform.pow(2).mean()
    snr_linear = 10 ** (snr_db / 10)
    noise_power = signal_power / snr_linear
    noise = torch.randn_like(waveform) * torch.sqrt(noise_power)
    return waveform + noise


class NoisyFluentSLUDataset(Dataset):
    """FluentSLU dataset with noise injection at a specified SNR."""

    def __init__(self, base_dataset, snr_db):
        self.base_dataset = base_dataset
        self.snr_db = snr_db

    def __len__(self):
        return len(self.base_dataset)

    def __getitem__(self, idx):
        waveform, sample_rate, filename, speaker_id, transcription, action, obj, location = (
            self.base_dataset.dataset[idx]
        )

        noisy_waveform = add_noise_snr(waveform, self.snr_db)

        mel = self.base_dataset.mel_transform(noisy_waveform)
        mel = mel.squeeze(0)

        if mel.shape[1] > self.base_dataset.max_frames:
            mel = mel[:, :self.base_dataset.max_frames]
        elif mel.shape[1] < self.base_dataset.max_frames:
            pad_size = self.base_dataset.max_frames - mel.shape[1]
            mel = torch.nn.functional.pad(mel, (0, pad_size))

        intent_key = f"{action}_{obj}_{location}"
        label = FluentSLUDataset.INTENT_MAP[intent_key]

        return mel, label


# === Main Execution ===

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    mel_transform = T.MelSpectrogram(
        sample_rate=16000,
        n_fft=512,
        hop_length=160,
        n_mels=80,
    )

    print("Loading Fluent Speech Commands dataset...")
    train_dataset = FluentSLUDataset("./data", "train", mel_transform)
    val_dataset = FluentSLUDataset("./data", "valid", mel_transform)
    test_dataset = FluentSLUDataset("./data", "test", mel_transform)

    num_intents = len(FluentSLUDataset.INTENT_MAP)
    print(f"Number of intents: {num_intents}")

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=2)

    model = SelectiveSLUModel(
        n_mels=80, d_model=128, nhead=4, num_layers=4,
        dim_feedforward=256, dropout=0.1, num_intents=num_intents
    ).to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params:,}")

    print("\n--- Training ---")
    history = train_model(model, train_loader, val_loader, device, epochs=30, lr=1e-3)

    model.load_state_dict(torch.load("best_selective_slu.pth"))

    print("\n--- Evaluation ---")
    coverages, accuracies, thresholds = compute_coverage_accuracy_curve(
        model, test_loader, device
    )

    plt.figure(figsize=(8, 5))
    plt.plot(coverages, accuracies, "b-o", markersize=4)
    plt.xlabel("Coverage")
    plt.ylabel("Selective Accuracy")
    plt.title("Coverage vs. Accuracy Curve")
    plt.grid(True)
    plt.savefig("coverage_accuracy_curve.png", dpi=150, bbox_inches="tight")
    plt.show()

    auroc, all_correct, all_scores = compute_auroc_failure_detection(
        model, test_loader, device
    )
    print(f"AUROC (failure detection): {auroc:.4f}")

    ece = compute_ece(all_scores, all_correct.astype(float))
    print(f"Expected Calibration Error: {ece:.4f}")

    print("\n--- Noise Robustness ---")
    snr_levels = [20, 10, 5, 0]
    overall_accs = []
    selective_accs = []

    for snr in snr_levels:
        noisy_dataset = NoisyFluentSLUDataset(test_dataset, snr_db=snr)
        noisy_loader = DataLoader(noisy_dataset, batch_size=64, shuffle=False, num_workers=2)

        sel_acc, coverage = evaluate_selective(model, noisy_loader, device, threshold=0.5)

        model.eval()
        correct_total = 0
        count_total = 0
        with torch.no_grad():
            for mel, labels in noisy_loader:
                mel, labels = mel.to(device), labels.to(device)
                pred_logits, _, _ = model(mel)
                preds = pred_logits.argmax(dim=1)
                correct_total += (preds == labels).sum().item()
                count_total += len(labels)

        overall_acc = correct_total / count_total
        selective_gain = sel_acc - overall_acc

        overall_accs.append(overall_acc)
        selective_accs.append(sel_acc)

        print(
            f"SNR {snr:3d}dB | Overall: {overall_acc:.4f} | "
            f"Selective: {sel_acc:.4f} | Gain: {selective_gain:.4f}"
        )

    plt.figure(figsize=(8, 5))
    plt.plot(snr_levels, overall_accs, "r-o", label="Overall Accuracy")
    plt.plot(snr_levels, selective_accs, "g-o", label="Selective Accuracy")
    plt.xlabel("SNR (dB)")
    plt.ylabel("Accuracy")
    plt.title("Noise Robustness: Overall vs. Selective Accuracy")
    plt.legend()
    plt.grid(True)
    plt.gca().invert_xaxis()
    plt.savefig("noise_robustness.png", dpi=150, bbox_inches="tight")
    plt.show()

    print("\nDone! Model and plots saved.")


if __name__ == "__main__":
    main()

Secret mission

Add Temperature Scaling Calibration

Your SelectiveNet head learned to say "I don't know" during training. But there's a simpler post-hoc technique that requires zero retraining: temperature scaling. Can a single learned parameter compete with your dedicated selection head for identifying failures?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Google Colab, so there are no ongoing costs. Colab's ephemeral storage is automatically wiped when your runtime disconnects, meaning the dataset and model files do not persist unless you save them to Google Drive.

Resources you used:

• Google Colab notebook (stored in Google Drive).
• Google Colab GPU runtime (T4 GPU allocation).
• Fluent Speech Commands dataset (downloaded to Colab ephemeral storage).
• Generated files: failure_aware_slu.py, temperature_scaling.py, best_selective_slu.pth, and plot images (in Colab ephemeral storage).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with the model or want to return to the notebook soon.

• Your Colab notebook stays in Google Drive indefinitely.
• The GPU runtime will auto-disconnect after about 90 minutes of inactivity, freeing the GPU for others.
• When you reconnect later, you will need to re-run the cells to re-download the dataset and retrain the model (ephemeral storage does not persist across sessions).

Pause - I'll come back to this later

Disconnect the GPU runtime to free up resources. Your notebook code stays saved in Google Drive.

• In your Colab notebook, click Runtime in the top menu bar.
• Click Disconnect and delete runtime.
• This frees the GPU allocation and deletes all ephemeral files (dataset, model checkpoint, plots).
• When you return, re-run the notebook cells from the top to re-download the dataset and retrain.

Delete - I don't want to use this again

Remove all project resources completely. This deletes the notebook from Google Drive and frees the GPU.

• In your Colab notebook, click Runtime in the top menu bar, then click Disconnect and delete runtime. This removes all ephemeral storage (the dataset, model checkpoint, and plots).
• Open Google Drive at drive.google.com.
• Navigate to Colab Notebooks in your Drive.
• Right-click your project notebook and select Move to trash.
• Empty the trash to permanently delete the file.

Nice Work!

Nice Work!

You just built a self-diagnosing speech model that knows when it does not know. This is the same reliability engineering principle that powers Amazon's billion-request-per-day Alexa pipeline.

Here's what you accomplished:

• Built a transformer-based SLU model from scratch that classifies 31 speech intents from raw audio using Mel spectrograms and a 4-layer encoder.
• Implemented the SelectiveNet architecture with a dual-head design (prediction + selection + auxiliary) trained with a joint loss that optimizes accuracy and coverage simultaneously.
• Evaluated with production-relevant metrics: coverage-accuracy curves, AUROC for failure detection, and Expected Calibration Error. Then proved noise robustness by showing the model abstains more under degraded conditions instead of failing silently.
• Secret Mission: Added temperature scaling calibration and compared the learned selection head against post-hoc calibrated softmax confidence for selective prediction.

Ready to quiz yourself?