Fast Ensembling with Diffusion Schrödinger Bridge
This addresses the computational bottleneck in ensemble methods for deep learning practitioners, offering a more efficient inference solution.
The paper tackles the high computational cost of Deep Ensembles during inference by proposing the Diffusion Bridge Network (DBN), which simulates an SDE to approximate ensemble predictions without running multiple models, achieving reduced computational cost while maintaining accuracy and uncertainty scores on datasets like CIFAR-10, CIFAR-100, and TinyImageNet.
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ödinger 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.