On-Device Domain Generalization

TL;DR

This paper investigates domain generalization for tiny neural networks on devices, finding that knowledge distillation and the proposed out-of-distribution knowledge distillation method significantly improve performance without increasing model size.

Contribution

It introduces OKD, a novel distillation approach using synthesized out-of-distribution data, tailored for tiny networks in on-device domain generalization tasks.

Findings

KD outperforms traditional DG methods for tiny networks

OKD improves DG performance without increasing model complexity

Synthesized domain shifts effectively enhance generalization

Abstract

We present a systematic study of domain generalization (DG) for tiny neural networks. This problem is critical to on-device machine learning applications but has been overlooked in the literature where research has been merely focused on large models. Tiny neural networks have much fewer parameters and lower complexity and therefore should not be trained the same way as their large counterparts for DG applications. By conducting extensive experiments, we find that knowledge distillation (KD), a well-known technique for model compression, is much better for tackling the on-device DG problem than conventional DG methods. Another interesting observation is that the teacher-student gap on out-of-distribution data is bigger than that on in-distribution data, which highlights the capacity mismatch issue as well as the shortcoming of KD. We further propose a method called out-of-distribution knowledge distillation (OKD) where the idea is to teach the student how the teacher handles out-of-distribution data synthesized via disruptive data augmentation. Without adding any extra parameter to the model -- hence keeping the deployment cost unchanged -- OKD significantly improves DG performance for tiny neural networks in a variety of on-device DG scenarios for image and speech applications. We also contribute a scalable approach for synthesizing visual domain shifts, along with a new suite of DG datasets to complement existing testbeds.