Fast Ensembling with Diffusion Schr\"odinger Bridge

TL;DR

This paper introduces Diffusion Bridge Network (DBN), a novel method based on Schr"odinger bridge theory, enabling efficient ensemble predictions by simulating an SDE, significantly reducing inference costs while maintaining accuracy.

Contribution

The paper presents a new diffusion-based approach for ensemble prediction that bypasses multiple forward passes, reducing computational overhead in deep neural network ensembles.

Findings

Achieved comparable accuracy to traditional ensembles on CIFAR-10, CIFAR-100, TinyImageNet.

Reduced inference computational cost significantly compared to standard deep ensembles.

Maintained uncertainty estimation quality with the proposed method.

Abstract

Deep Ensemble (DE) approach is a straightforward technique used to enhance the performance of deep neural networks by training them from different initial points, converging towards various local optima. However, a limitation of this methodology lies in its high computational overhead for inference, arising from the necessity to store numerous learned parameters and execute individual forward passes for each parameter during the inference stage. We propose a novel approach called Diffusion Bridge Network (DBN) to address this challenge. Based on the theory of the Schr\"odinger bridge, this method directly learns to simulate an Stochastic Differential Equation (SDE) that connects the output distribution of a single ensemble member to the output distribution of the ensembled model, allowing us to obtain ensemble prediction without having to invoke forward pass through all the ensemble models. By substituting the heavy ensembles with this lightweight neural network constructing DBN, we achieved inference with reduced computational cost while maintaining accuracy and uncertainty scores on benchmark datasets such as CIFAR-10, CIFAR-100, and TinyImageNet. Our implementation is available at https://github.com/kim-hyunsu/dbn.