🎧 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:
- Knowledge Distillation: Transferring the rich semantic representations of a pretrained heavy teacher (VGGish) into an ultra-compact student network.
- Dimensionality Reduction: Reducing the teacher’s 128-dimensional embedding space to 32 dimensions via Incremental PCA.
- 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 convolution to linearly combine the channels.
The theoretical reduction in computational cost compared to a standard convolution is given by:
Where is the number of output channels and is the kernel size. For a standard filter (), 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:
- (Embedding Loss): Root Mean Square Error (RMSE) between the predicted 32-dimensional student embedding and the compressed teacher embedding.
- (Classification Loss): Weighted Categorical Cross-Entropy across target scene categories.
- Weights: In our experiments, and 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:
- RAM Limits: Storing raw audio slices, log-mel spectrogram matrices, and embedding vectors in memory will easily trigger Out-Of-Memory (OOM) crashes.
- I/O File Handle Overhead: Saving each 1-second segment as an individual
.wavor.npyfile 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:
| Layer | Type / Filters | Kernel Shape | Activation |
|---|---|---|---|
| Conv 1 | Depthwise Conv (8) | ReLU + MaxPool | |
| Conv 2 | Depthwise Conv (16) | ReLU + MaxPool | |
| Conv 3 | Depthwise Conv (32) | ReLU + MaxPool | |
| Conv 4 | Depthwise Conv (64) | ReLU + MaxPool | |
| Conv 5 | Depthwise Conv (64) | ReLU + MaxPool | |
| FC 1 | Dense (256) | - | ReLU |
| FC 2 | Dense (32) | - | Linear (Embedding Output) |
| FC 3 | Dense (32) | - | ReLU |
| FC 4 | Dense (32) | - | ReLU |
| FC 5 | Dense ( classes) | - | Logits / Softmax |
Performance Comparison
| Metric | Production Model | Baseline VGGish | Target Specification |
|---|---|---|---|
| Operations | 1.62 MFLOPs | 864 MFLOPs | MFLOPs |
| Model Size | 0.45 MB | 275 MB | - |
| Macro Recall | 86.5% | 86.5% | |
| Macro Precision | 79.5% | 76.0% | - |
| Macro F1-Score | 82.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
- 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.
- On-Device ONNX / TFLite Export: Adding automated export scripts with validation benchmarks on actual embedded ARM Cortex-M DSP toolchains.
- Streaming Frame Inference: Adapting the framing from 1-second independent windows to causal ring-buffer streaming to reduce real-time inference latency.