Disentangled and Distilled Encoder for Out-of-Distribution Reasoning with Rademacher Guarantees

Recently, the disentangled latent space of a variational autoencoder (VAE) has been used to reason about multi-label out-of-distribution (OOD) test samples that are derived from different distributions than training samples. Disentangled latent space means having one-to-many maps between latent dimensions and generative factors or important characteristics of an image. This paper proposes a disentangled distilled encoder (DDE) framework to decrease the OOD reasoner size for deployment on resource-constrained devices while preserving disentanglement. DDE formalizes student-teacher distillation for model compression as a constrained optimization problem while preserving disentanglement with disentanglement constraints. Theoretical guarantees for disentanglement during distillation based on Rademacher complexity are established. The approach is evaluated empirically by deploying the compressed model on an NVIDIA

💡 Research Summary

This paper addresses a critical challenge in deploying machine learning for safety-critical cyber-physical systems (CPS) like autonomous vehicles: performing reliable reasoning on Out-of-Distribution (OOD) data using models that are compact enough for resource-constrained devices. While prior work has shown that the disentangled latent space of a Variational Autoencoder (VAE) can effectively identify the generative factors causing OOD behavior, these OOD reasoner models are often too large for embedded deployment.

The core contribution is the Disentangled Distilled Encoder (DDE) framework, which compresses a large “teacher” OOD reasoner into a smaller “student” model while rigorously preserving the disentanglement property crucial for interpretable OOD analysis. Disentanglement refers to a latent space where individual dimensions primarily correspond to specific generative factors (e.g., brightness, object type). The authors formulate the knowledge distillation process as a constrained optimization problem. The primary objective is to minimize a distillation loss (Jensen-Shannon divergence) between the latent distributions of the teacher and student encoders. To preserve disentanglement, two key constraints are enforced: 1) Adaptability: Information about a change in a specific generative factor must be transferred to the corresponding representative dimensions in the student’s latent space. 2) Isolation: Changes in a factor should be isolated to its representative dimensions and not leak into others. The paper provides differentiable surrogate losses for these constraints based on a mutual information estimator.

A significant theoretical advancement is the analysis of the optimization problem’s solution quality. The authors adapt concepts from constrained learning theory to account for the non-convexity of neural networks (“parameterization gap”) and finite training data (“empirical gap”). They employ Rademacher complexity to derive generalization bounds for the expected loss functions, providing theoretical guarantees for the preservation of disentanglement during distillation.

Empirically, the method is evaluated on the CARLA autonomous driving dataset. The student model, trained using DDE, demonstrates comparable OOD reasoning performance to the larger teacher model while achieving significant reductions in model size and inference latency. Practical deployment on an NVIDIA Jetson Nano showcases the framework’s viability for real-time CPS applications. Notably, DDE operates under weak supervision using only “match-paired” data (groups of samples sharing a factor value) and does not require prior knowledge of OOD samples during training, enhancing its practicality.