A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs

Fabian Falck, Christopher Williams, Dominic Danks, George Deligiannidis, Christopher Yau, Christopher C. Holmes, Arnaud Doucet, Matthew Willetts

Keywords: U-Net, multi-resolution analysis, hierarchical variational autoencoders

TL;DR: We provide a multi-resolution framework which theoretically characterises the regularisation in U-Nets with average pooling and apply it to hierarchical VAEs.

Abstract: U-Net architectures are ubiquitous in state-of-the-art deep learning, however their regularisation properties and relationship to wavelets are understudied. In this paper, we formulate a multi-resolution framework which identifies U-Nets as finite-dimensional truncations of models on an infinite-dimensional function space. We provide theoretical results which prove that average pooling corresponds to projection within the space of square-integrable functions and show that U-Nets with average pooling implicitly learn a Haar wavelet basis representation of the data. We then leverage our framework to identify state-of-the-art hierarchical VAEs (HVAEs), which have a U-Net architecture, as a type of two-step forward Euler discretisation of multi-resolution diffusion processes which flow from a point mass, introducing sampling instabilities. We also demonstrate that HVAEs learn a representation of time which allows for improved parameter efficiency through weight-sharing. We use this observation to achieve state-of-the-art HVAE performance with half the number of parameters of existing models, exploiting the properties of our continuous-time formulation.