🎧 Audio Model Distillation for Hearing Aids

Acoustic Scene Classification & Model Distillation Notes

During my master’s thesis research, I worked on developing an Acoustic Scene Classification (ASC) system designed to run on wireless hearing aids for noise suppression. While deep learning models perform exceptionally well on audio classification tasks, porting them to embedded hearing aid hardware presents extreme computational challenges.

I built this repository to implement the complete data processing, feature extraction, and knowledge distillation framework developed during that work.

You can check out the complete repository here: rubzip/model-distillation and read the thesis paper here: Acoustic Scene Classification in Hearing Aids.


Key Ideas: The Challenge of Edge Audio

In hearing aids, the device continuously records the acoustic environment and must decide in real time whether and how to apply noise reduction filters (e.g., distinguishing between clean speech, traffic, wind, or quiet indoor environments).

From a hardware perspective, hearing aids run on micro-batteries and ultra-low-power DSPs. The technical requirement for our production model was strict: under 20 MFLOPs (Million Floating Point Operations per Second) and a macro recall above 80%.

From an algorithmic perspective, modern state-of-the-art audio classifiers rely on deep convolutional backbones. A standard baseline feature extractor like Google’s VGGish requires ~864 MFLOPs and weighs 275 MB (over 70 million parameters). Running VGGish directly on a hearing aid chip is physically impossible—it would instantly drain the battery and exceed on-chip memory.

To bridge this gap, I designed a pipeline focused on aggressive model miniaturization:

  1. Knowledge Distillation: Transferring the rich semantic representations of a pretrained heavy teacher (VGGish) into an ultra-compact student network.
  2. Dimensionality Reduction: Reducing the teacher’s 128-dimensional embedding space to 32 dimensions via Incremental PCA.
  3. Depthwise Separable Convolutions: Replacing standard 2D convolutions with factorized depthwise and pointwise layers to cut computation by nearly 90%.

The result was a production model running at just 1.62 MFLOPs and 0.45 MB, while matching the teacher’s recall (86.5%).


Core Concepts Explained

1. Depthwise Separable Convolutions

Standard 2D convolutions apply spatial filtering and channel mixing simultaneously across all input channels. In mobile architectures (like MobileNet), this operation is split into two lighter steps:

  • Depthwise Convolution: Applies a single spatial filter per input channel independently.
  • Pointwise Convolution: Applies a 1×11 \times 1 convolution to linearly combine the channels.

The theoretical reduction in computational cost compared to a standard convolution is given by:

NDepthwiseNStandard=1M+1Dk2\frac{N_{\text{Depthwise}}}{N_{\text{Standard}}} = \frac{1}{M} + \frac{1}{D_k^2}

Where MM is the number of output channels and DkD_k is the kernel size. For a standard 3×33 \times 3 filter (Dk=3D_k = 3), this slashes computational overhead by roughly 88–90%.

Here is the depthwise block implementation in PyTorch:

def depthwise_conv_block(in_c: int, out_c: int) -> nn.Sequential:
    return nn.Sequential(
        # Depthwise spatial convolution
        nn.Conv2d(in_c, in_c, kernel_size=3, padding=1, groups=in_c),
        # Pointwise channel mixing
        nn.Conv2d(in_c, out_c, kernel_size=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )

2. Simultaneous Distillation Loss

Instead of training the student solely on class labels or solely on mimicking embeddings, the model is trained jointly using a weighted multi-task loss:

L=αeLe+αcLc\mathcal{L} = \alpha_e \mathcal{L}_e + \alpha_c \mathcal{L}_c

  • Le\mathcal{L}_e (Embedding Loss): Root Mean Square Error (RMSE) between the predicted 32-dimensional student embedding and the compressed teacher embedding.
  • Lc\mathcal{L}_c (Classification Loss): Weighted Categorical Cross-Entropy across target scene categories.
  • Weights: In our experiments, αe=1.5\alpha_e = 1.5 and αc=1.0\alpha_c = 1.0 produced the best balance, forcing the convolutional layers to retain rich acoustic features while the final classification head optimizes decision boundaries.

3. Why HDF5 (.h5) Files for Audio Data?

Working with hundreds of hours of raw audio introduces two major bottlenecks:

  1. RAM Limits: Storing raw audio slices, log-mel spectrogram matrices, and embedding vectors in memory will easily trigger Out-Of-Memory (OOM) crashes.
  2. I/O File Handle Overhead: Saving each 1-second segment as an individual .wav or .npy file leads to hundreds of thousands of tiny files. Operating system overhead for opening, reading, and closing descriptors causes severe training starvation.

The Solution: We store intermediate segments, mel spectrograms, and precomputed embeddings in structured HDF5 (.h5) files. During training, the PyTorch Dataset loads tensor slices lazily from disk on-demand for each batch:

with h5py.File(self.features_h5_path, "r") as f:
    feature_slice = torch.from_numpy(f["data"][start_idx:end_idx]).float()

This guarantees constant, low RAM usage regardless of how large the dataset grows.


Production Model Architecture

The distilled production network combines 5 depthwise separable stages with a lightweight projection and classifier:

LayerType / FiltersKernel ShapeActivation
Conv 1Depthwise Conv (8)(3,3)(3, 3)ReLU + MaxPool
Conv 2Depthwise Conv (16)(3,3)(3, 3)ReLU + MaxPool
Conv 3Depthwise Conv (32)(3,3)(3, 3)ReLU + MaxPool
Conv 4Depthwise Conv (64)(3,3)(3, 3)ReLU + MaxPool
Conv 5Depthwise Conv (64)(3,3)(3, 3)ReLU + MaxPool
FC 1Dense (256)-ReLU
FC 2Dense (32)-Linear (Embedding Output)
FC 3Dense (32)-ReLU
FC 4Dense (32)-ReLU
FC 5Dense (KK classes)-Logits / Softmax

Performance Comparison

MetricProduction ModelBaseline VGGishTarget Specification
Operations1.62 MFLOPs864 MFLOPs<20< 20 MFLOPs
Model Size0.45 MB275 MB-
Macro Recall86.5%86.5%>80.0%> 80.0\%
Macro Precision79.5%76.0%-
Macro F1-Score82.2%79.7%-

Project Structure

The repository is modularized into data processing and training pipelines:

.
├── config.py                 # Paths, audio hyperparameters, and label definitions
├── mock_data/                # Synthetic test dataset (raw audio and HDF5s)
├── scripts/
│   ├── data/                 # Data preparation, feature extraction, and validation
│   │   ├── generate_mock_data.py
│   │   ├── process_data.py
│   │   └── validate_dataset.py
│   └── training/             # Model training entry points
│       └── train_production_model.py
├── src/
│   ├── data/                 # Chunked H5 processors, metadata, and PyTorch dataset
│   ├── models/               # VGGish, MLP, and Production architectures
│   ├── processing/           # Mel spectrogram extraction, VGGish, and PCA reducers
│   └── training/             # Loss functions, metrics, and training loops
└── tests/                    # Unit and integration test suite

Running the Project

The project uses uv for reproducible environment and dependency management:

# 1. Generate synthetic mock data
uv run python scripts/data/generate_mock_data.py

# 2. Run the end-to-end processing pipeline (segments -> Mel -> embeddings -> PCA)
uv run python scripts/data/process_data.py

# 3. Validate dataset integrity
uv run python scripts/data/validate_dataset.py --mock

# 4. Train the production student model
uv run python scripts/training/train_production_model.py

# 5. Run the test suite
./test.sh

Future Improvements

  1. Quantization (INT8): Post-training quantization or quantization-aware training (QAT) to compress the model from 32-bit floats to 8-bit integers, reducing size to ~0.11 MB.
  2. On-Device ONNX / TFLite Export: Adding automated export scripts with validation benchmarks on actual embedded ARM Cortex-M DSP toolchains.
  3. Streaming Frame Inference: Adapting the framing from 1-second independent windows to causal ring-buffer streaming to reduce real-time inference latency.