A multi-resolution framework for U-nets with applications to hierarchical VAEs

Conference item

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.

Authors: Falck, F; Williams, C; Danks, D; Deligiannidis, G; Yau, C

• Publication status: Published
• Peer review status: Peer reviewed
• Publisher: Curran Associates
• Host title: Advances in Neural Information Processing Systems 35 (NeurIPS 2022)
• Volume: 21
• Pages: 15529-15544
• Publication date: 2023-04-01
• Acceptance date: 2022-09-14
• Event title: 36th Conference on Neural Information Processing Systems (NeurIPS 2022)
• Event location: New Orleans, USA
• Event website: https://nips.cc/Conferences/2022
• Event start date: 2022-11-28
• Event end date: 2022-12-09
• ISSN: 1049-5258
• EISBN: 9781713873129
• ISBN: 9781713871088
• Language: English
• Keywords: FFR