summaries: This paper proposes to learn a joint model consisting of a VAE and a latent space diffusion. The training process is augmented with a contrastive loss.  The authors show that this approach makes possible a few-shot generation of images with desired properties. To achieve this,  one needs to train a classifier in the latent space and after that sample latent codes with high classifier scores in any way.  Also, the case of semantic editing is evaluated. The manipulation is implemented using a Langevin-like approach in the latent space.

summaries: The paper describes two methods: 1) how to incorporate self-supervised learning in VAE to have a meaningful encoded representation, 2) how to use diffusion model over later to bridge the gap between posterior and prior for unconditional generation, 3) how to use the given setup to do few-shot conditional generation. Empirically, authors show highly improved representation quality, which can be attributed to the incorporation of self-supervised learning. Additionally, experiments show improved FID scores for face datasets, but not for CIFAR. Authors don’t show results on ImageNet. The authors show improved FID scores for the conditional generation task. There is no FID comparison against StyleGAN2 (or any other GAN-based model), even if the authors claim to come very close in sample quality to such models. The qualitative depiction of image manipulation tasks is very impressive.

summaries: The authors present Diffusion Decoding Models with Contrastive Representations (D2C) - a special VAE model with diffusion modeling in the latent space combined with contrastive learning based training of the inference network. The authors claim that D2C helps avoid the prior hole issue with VAEs due to diffusion modeling of the latent space, while contrastive learning helps learn meaningful representations allowing Few-Shot conditional generation. They further demonstrate that D2C outperforms NVAE and DDIM in few-shot conditional generation while remaining comparable to DDIM in unconditional generation.

summaries: D2C proposes three main contributions to diffusion models, a recently popularized class of likelihood-based generative models based on parameterized Markov chains, particularly successful when viewed from a denoising score matching lens. First, this paper learns a prior for the *latent space* of an image VAE (NVAE), while past work mostly applied diffusion to the data space. Second, D2C adds a SimCLR-style contrastive loss to image latents in order to improve the quality of latent representations for classification. Third, binary attribute classifiers are learned conditioned on these latents and used to guide prior sampling with rejection sampling or Langevin dynamics. Experimental results qualitatively and quantitatively show some success in conditional image generation and manipulation.

abstractText: Conditional generative models of high-dimensional images have many applications, but supervision signals from conditions to images can be expensive to acquire. This paper describes Diffusion-Decoding models with Contrastive representations (D2C), a paradigm for training unconditional variational autoencoders (VAEs) for few-shot conditional image generation. D2C uses a learned diffusion-based prior over the latent representations to improve generation and contrastive selfsupervised learning to improve representation quality. D2C can adapt to novel generation tasks conditioned on labels or manipulation constraints, by learning from as few as 100 labeled examples. On conditional generation from new labels, D2C achieves superior performance over state-of-the-art VAEs and diffusion models. On conditional image manipulation, D2C generations are two orders of magnitude faster to produce over StyleGAN2 ones and are preferred by 50% − 60% of the human evaluators in a double-blind study. We release our code at https: //github.com/jiamings/d2c.

authors:
Alexander A Alemi
Ben Poole
Ian Fischer
Joshua V Dillon
Rif A Saurous
Kevin Murphy

title: Fixing a broken ELBO

venue: arXiv preprint arXiv:1711.00464,

authors:
Jyoti Aneja
Harsh Agrawal
Dhruv Batra
Alexander Schwing

title: Sequential latent spaces for modeling the intention during diverse image captioning

venue: In Proceedings of the IEEE/CVF International Conference on Computer Vision,

authors:
Jyoti Aneja
Alexander Schwing
Jan Kautz
Arash Vahdat

title: NCP-VAE: Variational autoencoders with noise contrastive priors

venue: arXiv preprint arXiv:2010.02917,

authors:
Muhammad Asim
Ali Ahmed
Paul Hand

title: Invertible generative models for inverse problems: mitigating representation error and dataset bias

venue: arXiv preprint arXiv:1905.11672,

authors:
Sergey Bartunov
Dmitry Vetrov

title: Few-shot generative modelling with generative matching networks

venue: Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics,

authors:
David Bau
Jun-Yan Zhu
Jonas Wulff
William Peebles
Hendrik Strobelt
Bolei Zhou
Antonio Torralba

title: Seeing what a gan cannot generate

authors:
Andrew Brock
Jeff Donahue
Karen Simonyan

title: Large scale GAN training for high fidelity natural image synthesis

venue: arXiv preprint arXiv:1809.11096,

authors:
Andrew Brock
Theodore Lim
James M Ritchie
Nick Weston

title: Neural photo editing with introspective adversarial networks

venue: arXiv preprint arXiv:1609.07093,

authors:
Tom B Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell
Sandhini Agarwal
Ariel Herbert-Voss
Gretchen Krueger
Tom Henighan
Rewon Child
Aditya Ramesh
Daniel M Ziegler
Jeffrey Wu
Clemens Winter
Christopher Hesse
Mark Chen
Eric Sigler
Mateusz Litwin
Scott Gray
Benjamin Chess
Jack Clark
Christopher Berner
Sam McCandlish
Alec Radford
Ilya Sutskever
Dario Amodei

title: Language models are Few-Shot learners

venue: arXiv preprint arXiv:2005.14165,

authors:
Emmanuel J Candes
Justin K Romberg
Terence Tao

title: Stable signal recovery from incomplete and inaccurate measurements

venue: Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences,

authors:
Qifeng Chen
Vladlen Koltun

title: Photographic image synthesis with cascaded refinement networks

venue: In ICCV,

authors:
Ricky T Q Chen
Yulia Rubanova
Jesse Bettencourt
David Duvenaud

title: Neural ordinary differential equations

venue: arXiv preprint arXiv:1806.07366,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: arXiv preprint arXiv:2002.05709,

authors:
Xi Chen
Diederik P Kingma
Tim Salimans
Yan Duan
Prafulla Dhariwal
John Schulman
Ilya Sutskever
Pieter Abbeel

title: Variational lossy autoencoder

venue: arXiv preprint arXiv:1611.02731,

authors:
Xinlei Chen
Haoqi Fan
Ross Girshick
Kaiming He

title: Improved baselines with momentum contrastive learning

venue: arXiv preprint arXiv:2003.04297,

authors:
Xinlei Chen
Kaiming He

title: Exploring simple siamese representation learning

venue: arXiv preprint arXiv:2011.10566,

authors:
Xinlei Chen
Saining Xie
Kaiming He

title: An empirical study of training Self-Supervised vision transformers

venue: arXiv preprint arXiv:2104.02057,

title: Very deep vaes generalize autoregressive models and can outperform them on images

venue: arXiv preprint arXiv:2011.10650,

authors:
Kristy Choi
Aditya Grover
Trisha Singh
Rui Shu
Stefano Ermon

title: Fair generative modeling via weak supervision

authors:
Gordon Christie
Neil Fendley
James Wilson
Ryan Mukherjee

title: Functional map of the world

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Louis Clouâtre
Marc Demers

title: Figr: Few-shot image generation with reptile

venue: arXiv preprint arXiv:1901.02199,

authors:
Bin Dai
David Wipf

title: Diagnosing and enhancing VAE models

venue: arXiv preprint arXiv:1903.05789,

authors:
Giannis Daras
Joseph Dean
Ajil Jalal
Alexandros G Dimakis

title: Intermediate layer optimization for inverse problems using deep generative models

venue: arXiv preprint arXiv:2102.07364,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

venue: arXiv preprint arXiv:1810.04805,

authors:
Prafulla Dhariwal
Heewoo Jun
Christine Payne
Jong Wook Kim
Alec Radford
Ilya Sutskever

title: Jukebox: A generative model for music

venue: arXiv preprint arXiv:2005.00341,

authors:
Prafulla Dhariwal
Alex Nichol

title: Diffusion models beat GANs on image synthesis

venue: arXiv preprint arXiv:2105.05233,

authors:
L Dinh
D Krueger
Y Bengio

title: NICE: Non-linear independent components estimation

venue: arXiv preprint arXiv:1410.8516,

authors:
Laurent Dinh
Jascha Sohl-Dickstein
Samy Bengio

title: Density estimation using real NVP

venue: arXiv preprint arXiv:1605.08803,

authors:
Jeff Donahue
Philipp Krähenbühl
Trevor Darrell

title: Adversarial feature learning

venue: arXiv preprint arXiv:1605.09782,

authors:
Yilun Du
Igor Mordatch

title: Implicit generation and generalization in Energy-Based models

venue: arXiv preprint arXiv:1903.08689,

authors:
Vincent Dumoulin
Ishmael Belghazi
Ben Poole
Olivier Mastropietro
Alex Lamb
Martin Arjovsky
Aaron Courville

title: Adversarially learned inference

venue: arXiv preprint arXiv:1606.00704,

authors:
Charles Elkan
Keith Noto

title: Learning classifiers from only positive and unlabeled data

venue: In 14th ACM SIGKDD,

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

venue: Advances in Neural Information Processing Systems

authors:
Jean-Bastien Grill
Florian Strub
Florent Altché
Corentin Tallec
Pierre H Richemond
Elena Buchatskaya
Carl Doersch
Bernardo Avila Pires
Zhaohan Daniel Guo
Mohammad Gheshlaghi Azar
Bilal Piot
Koray Kavukcuoglu
Rémi Munos
Michal Valko

title: Bootstrap your own latent: A new approach to Self-Supervised learning

venue: arXiv preprint arXiv:2006.07733,

authors:
Aditya Grover
Aaron Zweig
Stefano Ermon

title: Graphite: Iterative generative modeling of graphs

venue: In International conference on machine learning,

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: arXiv preprint arXiv:1911.05722,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: arXiv preprint arXiv:1512.03385,

title: Data-efficient image recognition with contrastive predictive coding

authors:
Martin Heusel
Hubert Ramsauer
Thomas Unterthiner
Bernhard Nessler
Sepp Hochreiter

title: GANs trained by a two Time-Scale update rule converge to a local nash equilibrium

venue: arXiv preprint arXiv:1706.08500,

title: Training products of experts by minimizing contrastive divergence

venue: Neural computation,

authors:
Jonathan Ho
Ajay Jain
Pieter Abbeel

title: Denoising diffusion probabilistic models

venue: arXiv preprint arXiv:2006.11239,

authors:
Matthew D Hoffman
Matthew J Johnson

title: Elbo surgery: yet another way to carve up the variational evidence lower bound

venue: In Workshop in Advances in Approximate Bayesian Inference, NIPS,

authors:
Drew A Hudson
C Lawrence Zitnick

title: Generative adversarial transformers

venue: arXiv preprint arXiv:2103.01209,

authors:
Phillip Isola
Jun-Yan Zhu
Tinghui Zhou
Alexei A Efros

title: Image-to-image translation with conditional adversarial networks

authors:
Erik Jones
Shiori Sagawa
Pang Wei Koh
Ananya Kumar
Percy Liang

title: Selective classification can magnify disparities across groups

venue: arXiv preprint arXiv:2010.14134,

authors:
Zahra Kadkhodaie
Eero P Simoncelli

title: Solving linear inverse problems using the prior implicit in a denoiser

venue: arXiv preprint arXiv:2007.13640,

authors:
Tero Karras
Timo Aila
Samuli Laine
Jaakko Lehtinen

title: Progressive growing of GANs for improved quality, stability, and variation

venue: arXiv preprint arXiv:1710.10196,

authors:
Tero Karras
Samuli Laine
Timo Aila

title: A Style-Based generator architecture for generative adversarial networks

venue: arXiv preprint arXiv:1812.04948,

authors:
Tero Karras
Samuli Laine
Miika Aittala
Janne Hellsten
Jaakko Lehtinen
Timo Aila

title: Analyzing and improving the image quality of stylegan

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Diederik P Kingma
Tim Salimans
Rafal Jozefowicz
Xi Chen
Ilya Sutskever
Max Welling

title: Improved variational inference with inverse autoregressive flow

authors:
Diederik P Kingma
Max Welling

title: Auto-Encoding variational bayes

venue: arXiv preprint arXiv:1312.6114v10,

authors:
Alex Krizhevsky
Ilya Sutskever
Geoffrey E Hinton

title: ImageNet classification with deep convolutional neural networks

authors:
Anders Boesen Lindbo Larsen
Søren Kaae Sønderby
Hugo Larochelle
Ole Winther

title: Autoencoding beyond pixels using a learned similarity metric

authors:
C Li
H Liu
C Chen
Y Pu
L Chen
others

title: Alice: Towards understanding adversarial learning for joint distribution matching

venue: Advances in neural information processing systems,

authors:
Chunyuan Li
Xiang Gao
Yuan Li
Baolin Peng
Xiujun Li
Yizhe Zhang
Jianfeng Gao

title: Optimus: Organizing sentences via pre-trained modeling of a latent space

authors:
Tsung-Yi Lin
Michael Maire
Serge Belongie
James Hays
Pietro Perona
Deva Ramanan
Piotr Dollár
C Lawrence Zitnick

title: Microsoft coco: Common objects in context

venue: In European conference on computer vision,

authors:
Ziwei Liu
Ping Luo
Xiaogang Wang
Xiaoou Tang

title: Deep learning face attributes in the wild

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Elman Mansimov
Emilio Parisotto
Jimmy Lei Ba
Ruslan Salakhutdinov

title: Generating images from captions with attention

venue: arXiv preprint arXiv:1511.02793,

authors:
Mehdi Mirza
Simon Osindero

title: Conditional generative adversarial nets

venue: arXiv preprint arXiv:1411.1784,

authors:
Gautam Mittal
Jesse Engel
Curtis Hawthorne
Ian Simon

title: Symbolic music generation with diffusion models

venue: arXiv preprint arXiv:2103.16091,

authors:
Takeru Miyato
Toshiki Kataoka
Masanori Koyama
Yuichi Yoshida

title: Spectral normalization for generative adversarial networks

venue: arXiv preprint arXiv:1802.05957,

authors:
Shakir Mohamed
Balaji Lakshminarayanan

title: Learning in implicit generative models

venue: arXiv preprint arXiv:1610.03483,

title: Racial Discrimination in Face Recognition Technology, 2020

authors:
Anh Nguyen
Jeff Clune
Yoshua Bengio
Alexey Dosovitskiy
Jason Yosinski

title: Plug & play generative networks: Conditional iterative generation of images in latent space

authors:
Chenhao Niu
Yang Song
Jiaming Song
Shengjia Zhao
Aditya Grover
Stefano Ermon

title: Permutation invariant graph generation via score-based generative modeling

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Mehdi Noroozi
Paolo Favaro

title: Unsupervised learning of visual representations by solving jigsaw puzzles

authors:
Taesung Park
Jun-Yan Zhu
Oliver Wang
Jingwan Lu
Eli Shechtman
Alexei A Efros
Richard Zhang

title: Swapping autoencoder for deep image manipulation

venue: arXiv preprint arXiv:2007.00653,

authors:
Or Patashnik
Zongze Wu
Eli Shechtman
Daniel Cohen-Or
Dani Lischinski

title: StyleCLIP: Text-Driven manipulation of StyleGAN imagery

venue: arXiv preprint arXiv:2103.17249,

authors:
Wei Ping
Kainan Peng
Kexin Zhao
Zhao Song

title: Waveflow: A compact flow-based model for raw audio

authors:
Tiziano Portenier
Qiyang Hu
Attila Szabó
Siavash Arjomand Bigdeli
Paolo Favaro
Matthias Zwicker

title: Faceshop: Deep sketch-based face image editing

venue: ACM Transactions on Graphics,

authors:
Alec Radford
Jong Wook Kim
Chris Hallacy
Aditya Ramesh
Gabriel Goh
Sandhini Agarwal
Girish Sastry
Amanda Askell
Pamela Mishkin
Jack Clark
Gretchen Krueger
Ilya Sutskever

title: Learning transferable visual models from natural language supervision

venue: arXiv preprint arXiv:2103.00020,

authors:
Aditya Ramesh
Mikhail Pavlov
Gabriel Goh
Scott Gray
Chelsea Voss
Alec Radford
Mark Chen
Ilya Sutskever

title: Zero-Shot Text-to-Image generation

venue: arXiv preprint arXiv:2102.12092,

authors:
Ali Razavi
Aaron van den Oord
Oriol Vinyals

title: Generating diverse high-fidelity images with vq-vae-2

venue: arXiv preprint arXiv:1906.00446,

authors:
Scott Reed
Zeynep Akata
Xinchen Yan
Lajanugen Logeswaran
Bernt Schiele
Honglak Lee

title: Generative adversarial text to image synthesis

authors:
Danilo Jimenez Rezende
Shakir Mohamed

title: Variational inference with normalizing flows

venue: arXiv preprint arXiv:1505.05770,

authors:
Gareth O Roberts
Jeffrey S Rosenthal

title: Optimal scaling of discrete approximations to langevin diffusions

venue: Journal of the Royal Statistical Society: Series B (Statistical Methodology),

authors:
Mihaela Rosca
Balaji Lakshminarayanan
Shakir Mohamed

title: Distribution matching in variational inference

venue: arXiv preprint arXiv:1802.06847,

authors:
Yujun Shen
Jinjin Gu
Xiaoou Tang
Bolei Zhou

title: Interpreting the latent space of gans for semantic face editing

authors:
Jascha Sohl-Dickstein
Eric A Weiss
Niru Maheswaranathan
Surya Ganguli

title: Deep unsupervised learning using nonequilibrium thermodynamics

venue: arXiv preprint arXiv:1503.03585,

authors:
Jiaming Song
Stefano Ermon

title: Multi-label contrastive predictive coding

venue: arXiv preprint arXiv:2007.09852,

authors:
Jiaming Song
Pratyusha Kalluri
Aditya Grover
Shengjia Zhao
Stefano Ermon

title: Learning controllable fair representations

venue: arXiv preprint arXiv:1812.04218,

authors:
Jiaming Song
Chenlin Meng
Stefano Ermon

title: Denoising diffusion implicit models

venue: arXiv preprint arXiv:2010.02502,

authors:
Yang Song
Stefano Ermon

title: Generative modeling by estimating gradients of the data distribution

venue: arXiv preprint arXiv:1907.05600,

authors:
Yang Song
Jascha Sohl-Dickstein
Diederik P Kingma
Abhishek Kumar
Stefano Ermon
Ben Poole

title: Score-based generative modeling through stochastic differential equations

venue: arXiv preprint arXiv:2011.13456,

authors:
Hiroshi Takahashi
Tomoharu Iwata
Yuki Yamanaka
Masanori Yamada
Satoshi Yagi

title: Variational autoencoder with implicit optimal priors

venue: Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Jakub M Tomczak
Max Welling

title: Improving variational auto-encoders using householder flow

venue: arXiv preprint arXiv:1611.09630,

title: VAE with a VampPrior

venue: arXiv preprint arXiv:1705.07120,

authors:
Arash Vahdat
Jan Kautz

title: NVAE: A deep hierarchical variational autoencoder

venue: arXiv preprint arXiv:2007.03898,

authors:
Arash Vahdat
Karsten Kreis
Jan Kautz

title: Score-based generative modeling in latent space

venue: arXiv preprint arXiv:2106.05931,

authors:
Aaron van den Oord
Sander Dieleman
Heiga Zen
Karen Simonyan
Oriol Vinyals
Alex Graves
Nal Kalchbrenner
Andrew Senior
Koray Kavukcuoglu

title: WaveNet: A generative model for raw audio

venue: arXiv preprint arXiv:1609.03499,

authors:
Aaron van den Oord
Nal Kalchbrenner
Koray Kavukcuoglu

title: Pixel recurrent neural networks

venue: arXiv preprint arXiv:1601.06759,

authors:
Aaron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive coding

venue: arXiv preprint arXiv:1807.03748,

authors:
Aaron van den Oord
Oriol Vinyals
Koray Kavukcuoglu

title: Neural discrete representation learning

venue: arXiv preprint arXiv:1711.00937,

authors:
Ting-Chun Wang
Ming-Yu Liu
Jun-Yan Zhu
Andrew Tao
Jan Kautz
Bryan Catanzaro

title: High-resolution image synthesis and semantic manipulation with conditional gans

authors:
Weihao Xia
Yujiu Yang
Jing-Hao Xue
Baoyuan Wu

title: TediGAN: Text-Guided diverse face image generation and manipulation

authors:
Zhe Xie
Chengxuan Liu
Yichi Zhang
Hongtao Lu
Dong Wang
Yue Ding

title: Adversarial and contrastive variational autoencoder for sequential recommendation

venue: arXiv preprint arXiv:2103.10693,

authors:
Yinghao Xu
Yujun Shen
Jiapeng Zhu
Ceyuan Yang
Bolei Zhou

title: Generative hierarchical features from synthesizing images

authors:
Jiaxuan You
Bowen Liu
Rex Ying
Vijay Pande
Jure Leskovec

title: Graph convolutional policy network for Goal-Directed molecular graph generation

venue: arXiv preprint arXiv:1806.02473,

authors:
Lantao Yu
Weinan Zhang
Jun Wang
Yong Yu

title: Seqgan: Sequence generative adversarial nets with policy gradient

venue: In Proceedings of the AAAI conference on artificial intelligence,

authors:
Shengjia Zhao
Jiaming Song
Stefano Ermon

title: InfoVAE: Information maximizing variational autoencoders

venue: arXiv preprint arXiv:1706.02262,

title: Towards deeper understanding of variational autoencoding models

title: A lagrangian perspective on latent variable generative models

venue: In Proc. 34th Conference on Uncertainty in Artificial Intelligence,

authors:
Jiapeng Zhu
Yujun Shen
Deli Zhao
Bolei Zhou

title: In-domain gan inversion for real image editing

venue: In European Conference on Computer Vision,

authors:
Jiapeng Zhu
Deli Zhao
Bolei Zhou
Bo Zhang

title: Lia: Latently invertible autoencoder with adversarial learning

heading: 1 Introduction

text:
Generative models trained on large amounts of unlabeled data have achieved great success in various domains including images [8, 50, 75, 42], text [56, 2], audio [26, 71, 93, 62], and graphs [36, 67]. However, downstream applications of generative models are often based on various conditioning signals, such as labels [61], text descriptions [60], reward values [101], or similarity with existing data [45]. While it is possible to directly train conditional models, this often requires large amounts of paired data [57, 74] that are costly to acquire. Hence, it would be desirable to learn conditional generative models using large amounts of unlabeled data and as little paired data as possible.
Contrastive self-supervised learning (SSL) methods can greatly reduce the need for labeled data in discriminative tasks by learning effective representations from unlabeled data [95, 37, 35], and have also been shown to improve few-shot learning [39]. It is therefore natural to ask if they can also be used to improve few-shot generation. Latent variable generative models (LVGM) are a natural candidate for this, since they already involve a low-dimensional, structured latent representation of the data they generate. However, generative adversarial networks (GANs, [34, 50]) and diffusion
∗Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
models [42, 84], lack explicit tractable functions to map inputs to representations, making it difficult to optimize latent variables with SSL. Variational autoencoders (VAEs, [52, 77]), on the other hand, can naturally adopt SSL through their encoder model, but they typically have worse sample quality.
In this paper, we propose Diffusion-Decoding models with Contrastive representations (D2C), a special VAE that is suitable for conditional few-shot generation. D2C uses contrastive self-supervised learning methods to obtain a latent space that inherits the transferrability and few-shot capabilities of self-supervised representations. Unlike other VAEs, D2C learns a diffusion model over the latent representations. This latent diffusion model ensures that D2C uses the same latent distribution for both training and generation. We provide a formal argument to explain why this approach leads to better sample quality than existing hierarchical VAEs. We further discuss how to apply D2C to few-shot conditional generation where the conditions are defined through labeled examples or manipulation constraints. Our approach combines a discriminative model providing conditioning signal and generative diffusion model over the latent space, and is computationally more efficient than methods that act directly over the image space (Figure 1).
We evaluate and compare D2C with several state-of-the-art generative models over 6 datasets. On unconditional generation, D2C outperforms state-of-the-art VAEs and is competitive with diffusion models under similar computational budgets. On conditional generation with 100 labeled examples, D2C significantly outperforms state-of-the-art VAE [91] and diffusion models [84]. D2C can also learn to perform certain image manipulation tasks from as few as 100 labeled examples. Notably, for manipulating images, D2C is two orders of magnitude faster than StyleGAN2 [106] and preferred by 50%− 60% of human evaluations, which to our best knowledge is the first for any VAE model.

heading: 2 Background

text:
Latent variable generative models A latent variable generative model (LVGM) is posed as a conditional distribution pθ : Z → P(X ) from a latent variable z to a generated sample x, parametrized by θ. To acquire new samples, LVGMs draw random latent variables z from some distribution p(z) and map them to image samples through x ∼ pθ(x|z). Most LVGMs are built on top of four paradigms: variational autoencoders (VAEs, [52, 77]), Normalizing Flows (NFs, [28, 29]), Generative Adversarial Networks (GANs, [34]), and diffusion / score-based generative models [42, 85].
Particularly, VAEs use an inference model from x to z for training. Denoting the inference distribution from x to z as qφ(z|x), the generative distribution from z to x as pθ(x|z), VAEs are trained by minimizing the following upper bound of negative log-likelihood:
LVAE = Ex∼pdata [Ez∼qφ(z|x)[− log p(x|z)] +DKL(qφ(z|x)‖p(z))] (1)
where pdata is the data distribution and DKL is the KL-divergence.
Diffusion models Diffusion models [81, 42, 84] produce samples by reversing a Gaussian diffusion process. We use the index α ∈ [0, 1] to denote the particular noise level of an noisy observation x(α) = √ αx + √ 1− α , where x is the clean observation and ∼ N (0, I) is a standard Gaussian distribution; as α→ 0, the distribution of x(α) converges to N (0, I). Diffusion models are typically parametrized as reverse noise models θ(x(α), α) that predict the noise component of x(α) given a noise level α, and trained to minimize ‖ θ(x(α), α) − ‖22, the mean squared error loss between the true noise and predicted noise. Given any non-increasing series {αi}Ti=0 between 0 and 1, the diffusion objective for a clean sample from the data x is:
`diff(x;w, θ) := T∑ i=1 w(αi)E ∼N (0,I)[‖ − θ(x(αi), αi)‖ 2 2], x (αi) := √ αix + √ 1− αi (2)
where w : {αi}Ti=1 → R+ controls the loss weights for each α. When w(α) = 1 for all α, we recover the denoising score matching objective for training score-based generative models [85].
Given an initial sample x0 ∼ N (0, I), diffusion models acquires clean samples (i.e., samples of x1) through a gradual denoising process, where samples with reducing noise levels α are produced (e.g., x0 → x0.3 → x0.7 → x1). In particular, Denoising Diffusion Implicit Models (DDIMs, [84]) uses an Euler discretization of some neural ODE [13] to produce samples (Figure 2, left).
We provide a more detailed description for training diffusion models in Appendix A.1 and sampling from DDIM in Appendix A.2. For conciseness, we use the notation p(α)(x(α)) to denote the marginal distribution of x(α) under the diffusion model, and p(α1,α2)(x(α2) | x(α1)) to denote the diffusion sampling process from x(α1) to x(α2) (assuming α1 < α2).
Self-supervised learning of representations In self-supervised learning (SSL), representations are learned by completing certain pretext tasks that do not require extra manual labeling [68, 25]; these representations can then be applied to other downstream tasks, often in few-shot or zero-shot scenarios. In particular, contrastive representation learning encourages representations to be closer between “positive” pairs and further between “negative” pairs; contrastive predictive coding (CPC, [95]), based on multi-class classification, have been commonly used in state-of-the-art methods [37, 16, 18, 14, 82]. Other non-contrastive methods exist, such as BYOL [35] and SimSiam [17], but they usually require additional care to prevent the representation network from collapsing.

heading: 3 Problem Statement

text:
Few-shot conditional generation Our goal is to learn an unconditional generative model pθ(x) such that it is suitable for conditional generation. Let C(x, c, f) describe an event that “f(x) = c”, where c is a property value and f(x) is a property function that is unknown at training. In conditional generation, our goal is to sample x such that the event C(x, c, f) occurs for a chosen c. If we have access to some “ground-truth” model that gives us p(C|x) := p(f(x) = c|x), then the conditional model can be derived from Bayes’ rule: pθ(x|C) ∝ p(C|x)pθ(x). These properties c include (but are not limited to2) labels [61], text descriptions [60, 76], noisy or partial observations [11, 5, 47, 24], and manipulation constraints [69]. In many cases, we do not have direct access to the true f(x), so we need to learn an accurate model from labeled data [6] (e.g., (c,x) pairs).
Desiderata Many existing methods are optimized for some known condition (e.g., labels in conditional GANs [8]) or assume abundant pairs between images and conditions that can be used for pretraining (e.g., DALL-E [74] and CLIP [73] over image-text pairs). Neither is the case in this paper, as we do not expect to train over paired data.
While high-quality latent representations are not essential to unconditional image generation (e.g., autoregressive [94], energy-based [31], and some diffusion models [42]), they can be beneficial when we wish to specify certain conditions with limited supervision signals, similar to how SSL representations can reduce labeling efforts in downstream tasks. A compelling use case is detecting and removing biases in datasets via image manipulation, where we should not only address wellknown biases a-priori but also address other hard-to-anticipate biases, adapting to societal needs [65].
2When C refers to an event that is always true, we recover unconditioned generation.
Therefore, a desirable generative model should not only have high sample quality but also contain informative latent representations. While VAEs are ideal for learning rich latent representations due to being able to incorporate SSL within the encoder, they generally do not achieve the same level of sample quality as GANs and diffusion models.

heading: 4 Diffusion-Decoding Generative Models with Contrastive Learning

text:
In this paper, we present Diffusion-Decoding generative models with Contrastive Learning (D2C), an extension to VAEs with high-quality samples and high-quality latent representations, and are thus well suited to few-shot conditional generation. Moreover, unlike GAN-based methods, D2C does not involve unstable adversarial training (Table 1).
As its name suggests, the generative model for D2C has two components – diffusion and decoding; the diffusion component operates over the latent space and the decoding component maps from latent representations to images. Let us use the α index notation for diffusion random variables: z(0) ∼ p(0)(z(0)) := N (0, I) is the “noisy” latent variable with α = 0, and z(1) is the “clean” latent variable with α = 1. The generative process of D2C, which we denote pθ(x|z(0)), is then defined as:
z(0) ∼ p(0)(z(0)), z(1) ∼ p(0,1)θ (z (1)|z(0))︸ ︷︷ ︸
diffusion
, x ∼ pθ(x|z(1))︸ ︷︷ ︸ decoding , (3)
where p(0)(z(0)) = N (0, I) is the prior distribution for the diffusion model, p(0,1)θ (z(1)|z(0)) is the diffusion process from z(0) to z(1), and pθ(x|z(1)) is the decoder from z(1) to x. Intuitively, D2C models produce samples by drawing z(1) from a diffusion process and then decoding x from z(1).
In order to train a D2C model, we use an inference model qφ(z(1)|x) that predicts proper z(1) latent variables from x; this can directly incorporate SSL methods [99], leading to the following objective:
LD2C(θ, φ;w) := LD2(θ, φ;w) + λLC(qφ), (4)
LD2(θ, φ;w) := Ex∼pdata,z(1)∼qφ(z(1)|x)[− log pθ(x|z (1)) + `diff(z (1);w, θ)], (5)
where `diff is defined as in Eq.(2), LC(qφ) denotes any contrastive predictive coding objective [95] with rich data augmentations [37, 16, 18, 14, 82] (details in Appendix A.3) and λ > 0 is a weight hyperparameter. We illustrate D2C in Figure 2, and its training procedure in Appendix A.4.

heading: 4.1 Relationship to maximum likelihood

text:
The D2 objective (LD2) appears similar to the original VAE objective (LVAE). Here, we make an informal statement that the D2 objective function is deeply connected to the variational lower bound of log-likelihood; we present the full statement and proof in Appendix B.1. Theorem 1. (informal) For any valid {αi}Ti=0, there exists some weights ŵ : {αi}Ti=1 → R+ for the diffusion objective such that −LD2 is a variational lower bound to the log-likelihood, i.e.,
−LD2(θ, φ; ŵ) ≤ Epdata [log pθ(x)], (6) where pθ(x) := Ez0∼p(0)(z(0))[pθ(x|z (0))] is the marginal probability of x under the D2C model.
16
Proof. (sketch) The diffusion term `diff upper bounds the KL divergence between qφ(z1|x) and p
(1) θ (z (1)) for suitable weights [42, 84], which recovers a VAE objective.
Moreover, since some contrastive losses are lower bounds to mutual information [82], D2C can be treated as an implementation of the InfoVAE approach [103], where the latent variables are encouraged to have higher mutual information with the observations. Since the evidence lower bound (ELBO) for VAEs do not necessarily promote informative latent representations [15], D2C comes with the benefit of being able to learn more informative latents despite not exactly optimizing an ELBO.

heading: 4.2 D2C models address latent posterior mismatch in VAEs

text:
While D2C is a special case of VAE, we argue that D2C is non-trivial in the sense that it addresses a long-standing problem in VAE methods [90, 87], namely the mismatch between the prior distribution pθ(z) and the aggregate (approximate) posterior distribution qφ(z) := Epdata(x)[qφ(z|x)]. A mismatch could create “holes” [79, 43, 3] in the prior that the aggregate posterior fails to cover during training, resulting in worse sample quality, as many latent variables used during generation are likely to never have been trained on. We formalize this notion in the following definition.
Definition 1 (Prior hole). Let p(z), q(z) be two distributions with supp(q) ⊆ supp(p). We say that q has an ( , δ)-prior hole with respect to (the prior) p for , δ ∈ (0, 1), δ > , if there exists a set S ∈ supp(P ), such that ∫ S p(z)dz ≥ δ and ∫ S q(z)dz ≤ .
Intuitively, if qφ(z) has a prior hole with large δ and small (e.g., inversely proportional to the number of training samples), then it is very likely that latent variables within the hole are never seen during training (small ), yet frequently used to produce samples (large δ). Most existing methods address this problem by optimizing certain statistical divergences between qφ(z) and pθ(z), such as the KL divergence or Wasserstein distance [88]. However, we argue in the following statement that prior holes might not be eliminated even if we optimize certain divergence values to be reasonably low, especially when qφ(z) is very flexible. We present the formal statement and proof in Appendix B.2.
p(z) q(z)Theorem 2. (informal) Let pθ(z) = N (0, 1). For any > 0, there exists a qφ(z) with an ( , 0.49)-prior hole, such that DKL(qφ‖pθ) ≤ log 23and W2(qφ, pθ) < γ for any γ > 0, where W2 is the 2-Wasserstein distance.
3This is reasonably low for realistic VAE models (NVAE [91] reports a KL divergence of around 2810 nats).
Proof. (sketch) We construct a qφ that satisfies these properties (top-right figure). First, we truncate the Gaussian and divide them into regions with same probability mass; then we support qφ over half of these regions (so δ > 0.49); finally, we show that the divergences are small enough.
In contrast to addressing prior holes by optimization, diffusion models eliminate prior holes by construction, since the diffusion process from z(1) to z(0) is constructed such that the distribution of z(α) always converges to a standard Gaussian as α → 0. As a result, the distribution of latent variables used during training is arbitrarily close to that used in generation4, which is also the case in GANs. Therefore, our argument provides an explanation as to why we observe better sample quality results from GANs and diffusion models than VAEs and NFs.

heading: 5 Few-shot Conditional Generation with D2C

text:
In this section, we discuss how D2C can be used to learn to perform conditional generation from few-shot supervision. We note that D2C is only trained on images and not with any other data modalities (e.g., image-text pairs [74]) or supervision techniques (e.g., meta-learning [22, 6]).
Algorithm 1 Conditional generation with D2C 1: Input n examples {(xi, ci)}ni=1, property c. 2: Acquire latents z(1)i ∼ qφ(z
(1)|xi) for i ∈ [n]; 3: Train model rψ(c|z(1)) over {(z(1)i , ci)} n i=1 4: Sample latents with ẑ(1) ∼ rψ(c|z(1)) · p(1)θ (z (1))
(unnormalized); 5: Decode x̂ ∼ pθ(x|ẑ(1)). 6: Output x̂. Algorithm We describe the general algorithm for conditional generation from a few images in Algorithm 1, and detailed implementations in Appendix C. With a model over the latent space (denoted as rψ(c|z(1))), we draw conditional latents from an unnormalized distribution with the diffusion prior (line 4). This can be implemented in many ways such as rejection sampling or Langevin dynamics [66, 86, 27].
Conditions from labeled examples Given a few labeled examples, we wish to produce diverse samples with a certain label. For labeled examples we can directly train a classifier over the latent space, which we denote as rψ(c|z(1)) with c being the class label and z(1) being the latent representation of x from qφ(z(1)|x). If these examples do not have labels (i.e., we merely want to generate new samples similar to given ones), we can train a positive-unlabeled (PU) classifier [33] where we assign “positive” to the new examples and “unlabeled” to training data. Then we use the classifier with the diffusion model pθ(z(1)|z(0)) to produce suitable values of z(1), such as by rejecting samples from the diffusion model that has a small rψ(c|z(1)).
Conditions from manipulation constraints Given a few labeled examples, here we wish to learn how to manipulate images. Specifically, we condition over the event that “x has label c but is similar to image x̄”. Here rψ(c|z(1)) is the unnormalized product between the classifier conditional probability and closeness to the latent z̄(1) of x̄ (e.g., measured with RBF kernel). We implement line 4 of Alg. 1 with a Lanvegin-like procedure where we take a gradient step with respect to the classifier probability and then correct this gradient step with the diffusion model. Unlike many GAN-based methods [12, 72, 97, 45, 98], D2C does not need to optimize an inversion procedure at evaluation time, and thus the latent value is much faster to compute; D2C is also better at retaining fine-grained features of the original image due to the reconstruction loss.

heading: 6 Related Work

text:
Latent variable generative models Most deep generative models explicitly define a latent representation, except for some energy-based models [41, 31] and autoregressive models [94, 93, 10]. Unlike VAEs and NFs, GANs do not explicitly define an inference model and instead optimize a two-player game. In terms of sample quality, GANs currently achieve superior performance over VAEs and NFs, but they can be difficult to invert even with additional optimization [48, 100, 7]. This can be partially addressed by training reconstruction-based losses with GANs [54, 55]. Moreover,
4We expand this argument in Appendix B.2.
the GAN training procedure can be unstable [9, 8, 63], lack a informative objective for measuring progress [4], and struggle with discrete data [102]. Diffusion models [27] achieves high sample quality without adversarial training, but its latent dimension must be equal to the image dimension. A concurrent work [92] applied score-based generative modeling on the latent space of an VAE, and achieved similar improvements compared to regular VAEs.
Addressing posterior mismatch in VAEs Most methods address this mismatch problem by improving inference models [64, 51, 89], prior models [90, 3, 87], or objective functions [103, 104, 105, 1, 59]; all these approaches optimize the posterior model to be close to the prior. In Section 4.2, we explain why these approaches do not necessarily remove large “prior holes”, so their sample qualities remain relatively poor even after many layers [91, 19]. Other methods adopt a “two-stage” approach [23], which fits a generative model over the latent space of autoencoders [96, 75, 26, 74]. [62] have applied diffusion models directly on latent spaces learned by a music VAE.
Conditional generation with unconditional models To perform conditional generation over an unconditional LVGM, most methods assume access to a discriminative model (e.g., a classifier); the latent space of the generator is then modified to change the outputs of the discriminative model. The disciminative model can operate on either the image space [66, 70, 27] or the latent space [80, 98]. For image space discriminative models, plug-and-play generative networks [66] control the attributes of generated images via Langevin dynamics [78]; these ideas are also explored in diffusion models [86]. Image manipulation methods are based on GANs often operate with latent space discriminators [80, 98]. However, these methods have some trouble manipulating real images because of imperfect reconstruction [107, 7]. This is not a problem in D2C since a reconstruction objective is optimized.

heading: 7 Experiments

text:
We examine the conditional and unconditional generation qualities of D2C over CIFAR-10 [53], CIFAR-100 [53], fMoW [21], CelebA-64 [58], CelebA-HQ-256 [48], and FFHQ-256 [49]. Our D2C implementation is based on the state-of-the-art NVAE [91] autoencoder structure, the U-Net diffusion model [42], and the MoCo-v2 contrastive representation learning method [16]. We also consider the D2 objective where we do not applying the contrastive loss in the NVAE autoencoder; we also tried adding contrastive loss directly to the VAE objective, but we were unable to achieve satisfactory generation results (reconstruction MSE remains high). This is possibly due to the many regularizations needed for NVAE to work well, which could conflict with contrastive learning5.
We keep the diffusion series hyperparameter {αi}Ti=1 identical to ensure a fair comparison with different diffusion models. For the contrastive weight λ in Equation (4), we consider the value of λ = 10−4 based on the relative scale between the LC and LD2; we find that the results are relatively insensitive to λ. We use 100 diffusion steps for DDIM and D2C unless mentioned otherwise, as running with longer steps is not computationally economical despite tiny gains in FID [84]. We include additional training details, such as architectures, optimizers and learning rates in Appendix C.

heading: 7.1 Unconditional generation

text:
For unconditional generation, we measure the sample quality of images using the Frechet Inception Distance (FID, [40]) with 50,000 images. In particular, we extensively evaluate NVAE [91] and
5See https://github.com/NVlabs/NVAE#known-issues for a detailed description.
DDIM [84], a competitive VAE model and a competitive diffusion model as baselines because we can directly obtain features from them without additional optimization steps6. For them, we report mean-squared reconstruction error (MSE, summed over all pixels, pixels normalized to [0, 1]) and linear classification accuracy (Acc., measured in percentage) over z1 features for the test set.
We report sample quality results7 in Tables 2, and 3. For FID, we outperform NVAE in all datasets and outperform DDIM on CelebA-64 and CelebA-HQ-256, which suggests our results are competitive with state-of-the-art non-adversarial generative models. In Table 2, we additionally compare NVAE, DDIM and D2C in terms of reconstruction and linear classification accuracy. As all three methods contain reconstruction losses, the MSE values are low and comparable. However, D2C enjoys much better linear classification accuracy than the other two thanks to the contrastive SSL component. We further note that training the same contrastive SSL method without LD2 achieves slightly higher 78.3% accuracy on CIFAR-10. We tried improving this via ResNet [38] encoders, but this significantly increased reconstruction error, possibly due to loss of information in average pooling layers.

heading: 7.2 Few-shot conditional generation from examples

text:
We demonstrate the advantage of D2C representations by performing few-shot conditional generation over labels. We consider two types of labeled examples: one has binary labels for which we train a binary classifier; the other is positive-only labeled (e.g., images of female class) for which we train a PU classifier. Our goal here is to generate a diverse group of images with a certain label. We evaluate and compare three models: D2C, NVAE and DDIM. We train a classifier rψ(c|z) over the latent space of these models; we also train a image space classifier and use it with DDIM (denoted as DDIM-I). We run Algorithm 1 for these models, where line 4 is implemented via rejection sampling. As our goal is to compare different models, we leave more sophisticated methods [27] as future work.
First, we consider performing 8 conditional generation tasks over CelebA-64 with 2 binary classifiers (trained over 100 samples, 50 for each class) and 4 PU classifiers (trained over 100 positively labeled and 10k unlabeled samples). We also report a “naive” approach where we use all the training images (regardless of labels) and compute its FID with the corresponding subset of images (e.g., all images versus blond images). In Table 4, we report the FID score between generated images (5k samples) and real images of the corresponding label.
Next, we perform a similar experiment on CIFAR-10, but with 50 labels for each of the 10 classes. For each label, we evaluate the FID between the real images and generated images under that label
6For DDIM, the latent representations x(0) are obtained by reversing the neural ODE process. 7Due to space limits, we place additional CIFAR-10 results in Appendix D.
and take the average FID over the 10 labels. The average FID for D2C, DDIM, and NVAE are 32.57, 82.78, and 95.45 respectively (note that FIDs are expected to be higher because it is evaluated with 5k real samples per class instead of the commonly used 50k). This shows that D2C performs significantly better than DDIM and NVAE on the few-shot conditional generation task for both datasets, illustrating the advantage of contrastive representations for few-shot conditional generation. We demonstrate some figures in the appendix.
Fair generative modeling with D2C We illustrate some conditional generation results with D2C (from 100 “blond” labels) in Figure 4 (left), which exhibits a biased gender ratio (3 examples out of 36). This kind of dataset bias in CelebA is well-documented in previous work [46]; among all the images with the “blond” attribute, only 5.6 are labeled as “male”. In fact, we observe that the D2C samples have a similar gender bias as the original dataset since D2C has learned the data distribution quite well. To mitigate the presence of dataset bias in D2C’s samples, we leverage a conditional generation approach. This is in contrast to previous methods that try to learn a fair generative model via a carefully designed representative dataset [20]. We give a concrete instantiation of our method in the case of gender bias that can be generalized to all attributes. Suppose our goal is to achieve a balanced gender ratio among all the generated images with “blond” attribute; we can do so by generating the same number of samples conditioned on the (1) “blond, male” attributes and the (2) “blond, female” attributes. Figure 4 (right) demonstrates that D2C can be used for fair generative modeling by simply specifying our desired constraints with a handful of labels.
17

heading: 7.3 Few-shot conditional generation from manipulation constraints

text:
Finally, we consider image manipulation where we use binary classifiers that are learned over 50 labeled instances for each class. To manipulate an image, we first acquire its latent code, move its latent code towards the decision boundary of the classifier and then use the diffusion prior to make the
sample more realistic. We perform Amazon Mechanical Turk (AMT) evaluations over two attributes in the CelebA-256 dataset, blond and red lipstick, over D2C, DDIM, NVAE and StyleGAN2 [50] (see Figure 5). The evaluation is double-blinded: neither we nor the evaluators know the correspondence between generated image and underlying model during the study. We include more details (algorithm, setup and human evaluation) in Appendix C and additional qualitative results (such as beard and gender attributes) in Appendix D.
In Figure 8, we show the percentage of manipulations preferred by AMT evaluators for each model; D2C slightly outperforms StyleGAN2 for blond and significantly outperforms StyleGAN2 for red lipstick. When we compare D2C with only StyleGAN2, D2C is preferred over 51.5% for blond and 60.8% for red lipstick. An additional advantage of D2C is that the manipulation is much faster than StyleGAN2, since the latter requires additional optimization over the latent space to improve reconstruction [106]. On the same Nvidia 1080Ti GPU, it takes 0.013 seconds to obtain the latent code in D2C, while the same takes 8 seconds [106] for StyleGAN2 (615× slower). As decoding is very fast for both models, D2C generations are around two orders of magnitude faster to produce.

heading: 8 Discussions and Limitations

text: We introduced D2C, a VAE-based generative model with a latent space suitable for few-shot conditional generation. To our best knowledge, our model is the first unconditional VAE to demonstrate superior image manipulation performance than StyleGAN2, which is surprising given our use of a regular NVAE architecture. We believe that with better architectures, such as designs from StyleGAN2 or Transformers [44], D2C can achieve even better performance. It is also interesting to formally investigate the integration between D2C and other types of conditions on the latent space, as well as training D2C in conjunction with other domains and data modalities, such as text [74], in a fashion that is similar to semi-supervised learning. Nevertheless, we note that our model have to be used properly in order to mitigate potential negative societal impacts, such as deep fakes.

heading: Acknowledgements

text:
We thank the anonymous reviewers and Kristy Choi for insightful discussions and feedback.
Funding transparency statement. This research was supported by NSF (#1651565, #1522054, #1733686), ONR (N00014-19-1-2145), AFOSR (FA9550-19-1-0024), ARO (W911NF-21-1-0125), Sloan Fellowship, Amazon AWS, Stanford Institute for Human-Centered Artificial Intelligence (HAI), and Google Cloud.

title: D2C: Diffusion-Decoding Models for Few-Shot Conditional Generation


summaries: This paper theoretically analyzes contrastive learning relying on a concept of population augmentation graph. In the graph, two augmented samples are connected if they can be augmented from the same original instance. Based on this setup, - This paper shows that spectral decomposition can be reformulated as a contrastive learning objective. Furthermore, this paper shows that minimizing the objective can guarantee a high linear probe accuracy under some mild assumptions. - This paper shows the guarantee can be extended to the finite-sample regime and learning the linear probe with labeled data. - This paper demonstrates the efficacy of the proposed spectral contrastive objective empirically in the standard setup: the proposed objective achieves similar performance as other state-of-the-art SSL methods.

summaries: The paper studies how does contrastive learning works from a theoretical perspective focusing on the spectral properties of the graph of the data induced by augmentations. They propose a new loss that is close up to a constant to the matrix factorization of the normalized adjacency matrix of the augmentations graph, under some assumptions they show that minimizing this loss provides downstream classification guarantees. Then, using standard Rademacher analysis they move from population guarantees to finite sample guarantees. Finally, they show empirically that their loss gives results close to state of the art.

summaries: This paper derives generalization bounds for learning using a contrastive objective under realistic assumptions on the positive pair generation process. The approach taken is to consider the so-called augmentation graph - the set of nodes comprises all augmented views of population distribution, with edges connecting nodes that correspond to different views of the same input datapoint. The core assumption made is that this graph cannot be split into a large number of disconnected subgraphs. This set-up aligns well with the intuition that in order to generalize, the contrastive notion of “similarity” must extent beyond the purely single-instance-level, and must somehow connect distinct inputs points.  My view is that this work is the most significant advance in the statistical learning theory of contrastive learning since the paper “A theoretical analysis of contrastive unsupervised representation learning” by Arora et al. This paper under review addresses the primary weakness in the work by Arora et al., which used an unrealistic assumption on positive pair generation. I will strongly advocate for acceptance of this paper.

summaries: The paper studies the theoretical foundations of contrastive self-supervised learning. To do so, the authors study contrastive learning through the lens of the augmentation graph - the graph of augmented examples, with edges connecting "similar" examples. The authors suggest a novel contrastive loss function, s.t. minimizing this loss is equivalent to spectral clustering on the augmentation graph. Using this, the authors prove that if the augmentation graph has few edges that cross between sub-graphs of different labels, then using a linear readout on the learned representation, a small loss can be achieved w.r.t the supervised objective. The authors also analyze the sample complexity of learning the representation and learning the linear probe from sampled data. Finally, the authors show empirically that the suggested contrastive loss performs similarly to some baseline contrastive loss functions.

abstractText: Recent works in self-supervised learning have advanced the state-of-the-art by relying on the contrastive learning paradigm, which learns representations by pushing positive pairs, or similar examples from the same class, closer together while keeping negative pairs far apart. Despite the empirical successes, theoretical foundations are limited – prior analyses assume conditional independence of the positive pairs given the same class label, but recent empirical applications use heavily correlated positive pairs (i.e., data augmentations of the same image). Our work analyzes contrastive learning without assuming conditional independence of positive pairs using a novel concept of the augmentation graph on data. Edges in this graph connect augmentations of the same datapoint, and ground-truth classes naturally form connected sub-graphs. We propose a loss that performs spectral decomposition on the population augmentation graph and can be succinctly written as a contrastive learning objective on neural net representations. Minimizing this objective leads to features with provable accuracy guarantees under linear probe evaluation. By standard generalization bounds, these accuracy guarantees also hold when minimizing the training contrastive loss. Empirically, the features learned by our objective can match or outperform several strong baselines on benchmark vision datasets. In all, this work provides the first provable analysis for contrastive learning where guarantees for linear probe evaluation can apply to realistic empirical settings.

title: Community detection and stochastic block models: recent developments, 2017

authors:
S. Arora
S. Rao
U. Vazirani

title: Expander flows, geometric embeddings and graph partitioning

venue: Journal of the ACM (JACM),

authors:
S. Arora
H. Khandeparkar
M. Khodak
O. Plevrakis
N. Saunshi

authors:
P. Bachman
R.D. Hjelm
W. Buchwalter

title: Learning representations by maximizing mutual information across views

venue: arXiv preprint arXiv:1906.00910,

authors:
Y. Bansal
G. Kaplun
B. Barak

title: For self-supervised learning, rationality implies generalization, provably

venue: arXiv preprint arXiv:2010.08508,

authors:
A. Bardes
J. Ponce
Y. LeCun

title: Vicreg: Variance-invariance-covariance regularization for self-supervised learning

venue: arXiv preprint arXiv:2105.04906,

title: An isoperimetric inequality on the discrete cube, and an elementary proof of the isoperimetric inequality in gauss space

venue: The Annals of Probability,

authors:
J. Bromley
I. Guyon
Y. LeCun
E. Säckinger
R. Shah

title: Signature verification using a" siamese" time delay neural network

venue: Advances in neural information processing systems,

authors:
T. Cai
R. Gao
J.D. Lee
Q. Lei

title: A theory of label propagation for subpopulation shift

venue: arXiv preprint arXiv:2102.11203,

authors:
M. Caron
I. Misra
J. Mairal
P. Goyal
P. Bojanowski
A. Joulin

title: Unsupervised learning of visual features by contrasting cluster assignments

venue: arXiv preprint arXiv:2006.09882,

title: A lower bound for the smallest eigenvalue of the laplacian

venue: In Proceedings of the Princeton conference in honor of Professor S. Bochner,

authors:
T. Chen
S. Kornblith
M. Norouzi
G. Hinton

title: A simple framework for contrastive learning of visual representations

venue: In International conference on machine learning,

authors:
T. Chen
S. Kornblith
K. Swersky
M. Norouzi
G. Hinton

title: Big self-supervised models are strong semi-supervised learners

venue: arXiv preprint arXiv:2006.10029,

authors:
X. Chen
K. He

title: Exploring simple siamese representation learning

venue: arXiv preprint arXiv:2011.10566,

authors:
X. Chen
H. Fan
R. Girshick
K. He

title: Improved baselines with momentum contrastive learning

venue: arXiv preprint arXiv:2003.04297,

authors:
C.-Y. Chuang
J. Robinson
L. Yen-Chen
A. Torralba
S. Jegelka

title: Debiased contrastive learning

venue: arXiv preprint arXiv:2007.00224,

authors:
F.R. Chung
F.C. Graham

title: Spectral graph theory

venue: Number 92. American Mathematical Soc.,

authors:
J. Deng
W. Dong
R. Socher
L.-J. Li
K. Li
L. Fei-Fei

title: ImageNet: A Large-Scale Hierarchical Image Database

venue: In CVPR09,

authors:
C. Eckart
G. Young

title: The approximation of one matrix by another of lower rank

authors:
N. Golowich
A. Rakhlin
O. Shamir

title: Size-independent sample complexity of neural networks

venue: In Conference On Learning Theory,

authors:
J.-B. Grill
F. Strub
F. Altché
C. Tallec
P.H. Richemond
E. Buchatskaya
C. Doersch
B.A. Pires
Z.D. Guo
M.G. Azar

title: Bootstrap your own latent: A new approach to self-supervised learning

venue: arXiv preprint arXiv:2006.07733,

authors:
K. He
H. Fan
Y. Wu
S. Xie
R. Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

title: Data-efficient image recognition with contrastive predictive coding

venue: In International Conference on Machine Learning,

authors:
R.D. Hjelm
A. Fedorov
S. Lavoie-Marchildon
K. Grewal
P. Bachman
A. Trischler
Y. Bengio

title: Learning deep representations by mutual information estimation and maximization

venue: In International Conference on Learning Representations,

authors:
R. Kannan
S. Vempala
A. Vetta

title: On clusterings: Good, bad and spectral

authors:
A. Krizhevsky
G. Hinton

title: Learning multiple layers of features from tiny images

authors:
Y. Le
X. Yang

title: Tiny imagenet visual recognition challenge

authors:
J.D. Lee
Q. Lei
N. Saunshi
J. Zhuo

title: Predicting what you already know helps: Provable self-supervised learning

venue: arXiv preprint arXiv:2008.01064,

authors:
J.R. Lee
S.O. Gharan
L. Trevisan

title: Multiway spectral partitioning and higher-order cheeger inequalities

authors:
J. Lei
A. Rinaldo

title: Consistency of spectral clustering in stochastic block models

venue: Annals of Statistics,

authors:
T. Leighton
S. Rao

title: Multicommodity max-flow min-cut theorems and their use in designing approximation algorithms

authors:
A. Louis
K. Makarychev

title: Approximation algorithm for sparsest k-partitioning

venue: In Proceedings of the twenty-fifth annual ACM-SIAM symposium on Discrete algorithms,

authors:
A. Louis
P. Raghavendra
P. Tetali
S. Vempala

title: Algorithmic extensions of cheeger’s inequality to higher eigenvalues and partitions. In Approximation, Randomization, and Combinatorial Optimization

venue: Algorithms and Techniques,

title: Spectral partitioning of random graphs

venue: In Proceedings 42nd IEEE Symposium on Foundations of Computer Science,

authors:
I. Misra
L. v. d. Maaten

title: Self-supervised learning of pretext-invariant representations

authors:
M. Mohri
A. Rostamizadeh
A. Talwalkar

title: Foundations of machine learning

venue: MIT press,

authors:
A. Ng
M. Jordan
Y. Weiss

title: On spectral clustering: Analysis and an algorithm

authors:
A. v. d. Oord
Y. Li
O. Vinyals

authors:
J. Shi
J. Malik

title: Normalized cuts and image segmentation

venue: IEEE Transactions on pattern analysis and machine intelligence,

title: Improved deep metric learning with multi-class n-pair loss objective

venue: In Proceedings of the 30th International Conference on Neural Information Processing Systems,

authors:
Y. Tian
D. Krishnan
P. Isola

title: Contrastive multiview coding

venue: arXiv preprint arXiv:1906.05849,

authors:
Y. Tian
C. Sun
B. Poole
D. Krishnan
C. Schmid
P. Isola

title: What makes for good views for contrastive learning

venue: arXiv preprint arXiv:2005.10243,

authors:
Y. Tian
L. Yu
X. Chen
S. Ganguli

title: Understanding self-supervised learning with dual deep networks

venue: arXiv preprint arXiv:2010.00578,

authors:
C. Tosh
A. Krishnamurthy
D. Hsu

title: Contrastive estimation reveals topic posterior information to linear models

title: Contrastive learning, multi-view redundancy, and linear models

venue: In Algorithmic Learning Theory,

authors:
T.W. Tsai
C. Li
J. Zhu

title: Mice: Mixture of contrastive experts for unsupervised image clustering

venue: arXiv preprint arXiv:2105.01899,

authors:
Y.-H.H. Tsai
Y. Wu
R. Salakhutdinov
L.-P. Morency

title: Self-supervised learning from a multi-view perspective

venue: arXiv preprint arXiv:2006.05576,

authors:
T. Wang
P. Isola

title: Understanding contrastive representation learning through alignment and uniformity on the hypersphere

authors:
C. Wei
K. Shen
Y. Chen
T. Ma

title: Theoretical analysis of self-training with deep networks on unlabeled data

venue: arXiv preprint arXiv:2010.03622,

authors:
Z. Wu
Y. Xiong
S.X. Yu
D. Lin

title: Unsupervised feature learning via non-parametric instance discrimination

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
M. Ye
X. Zhang
P.C. Yuen
S.-F. Chang

title: Unsupervised embedding learning via invariant and spreading instance feature

authors:
J. Zbontar
L. Jing
I. Misra
Y. LeCun
S. Deny. Barlow

title: twins: Self-supervised learning via redundancy reduction

venue: arXiv preprint arXiv:2103.03230,

heading: 1 Introduction

text:
Recent empirical breakthroughs have demonstrated the effectiveness of self-supervised learning, which trains representations on unlabeled data with surrogate losses and self-defined supervision signals [4, 6, 10, 14, 23, 24, 35, 38, 41, 42, 50–52]. Self-supervision signals in computer vision are often defined by using data augmentation to produce multiple views of the same image. For example, the recent contrastive learning objectives [3, 12, 13, 15, 22] encourage closer representations for augmentations (views) of the same natural data than for randomly sampled pairs of data.
Despite the empirical successes, there is a limited theoretical understanding of why self-supervised losses learn representations that can be adapted to downstream tasks, for example, using linear heads. Recent mathematical analyses by Arora et al. [3], Lee et al. [28], Tosh et al. [44, 45] provide guarantees under the assumption that two views are somewhat independent conditioned on the label. However, the pair of augmented examples used in practical algorithms usually exhibit a strong correlation, even conditioned on the label. For instance, two augmentations of the same dog image
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
share much more similarity than augmentations of two different random dog images. Thus the existing theory does not explain the practical success of self-supervised learning.
This paper presents a theoretical framework for self-supervised learning without requiring conditional independence. We design a principled, practical loss function for learning neural net representations that resembles state-of-the-art contrastive learning methods. We prove that, under a simple and realistic data assumption, linear classification using representations learned on a polynomial number of unlabeled data samples can recover the ground-truth labels of the data with high accuracy.
The fundamental data property that we leverage is a notion of continuity of the population data within the same class. Though a random pair of examples from the same class can be far apart, the pair is often connected by (many) sequences of examples, where consecutive examples in the sequences are close neighbors within the same class. This property is more salient when the neighborhood of an example includes many different types of augmentations. Prior work [49] empirically demonstrates this type of connectivity property and uses it in the analysis of pseudolabeling algorithms.
More formally, we define the population augmentation graph, whose vertices are all the augmented data in the population distribution, which can be an exponentially large or infinite set. Two vertices are connected with an edge if they are augmentations of the same natural example. Our main assumption is that for some proper m 2 Z+, the sparsest m-partition (Definition 3.4) is large. This intuitively states that we can’t split the augmentation graph into too many disconnected sub-graphs by only removing a sparse set of edges. This assumption can be seen as a graph-theoretic version of the continuity assumption on population data. We also assume that there are very few edges across different ground-truth classes (Assumption 3.5). Figure 1 (left) illustrates a realistic scenario where dog and cat are the ground-truth categories, between which edges are very rare. Each breed forms a sub-graph that has sufficient inner connectivity and thus cannot be further partitioned.
Our assumption fundamentally does not require conditional independence and can allow disconnected sub-graphs within a class. The classes in the downstream task can be also somewhat flexible as long as they are disconnected in the augmentation graph. For example, when the augmentation graph consists of m disconnected sub-graphs corresponding to fine-grained classes, our assumptions allow the downstream task to have any r  m coarse-grained classes containing these fine-grained classes as a sub-partition. Prior work [49] on pseudolabeling algorithms essentially requires an exact alignment between sub-graphs and downstream classes (i.e., r = m). They face this limitation because their analysis requires fitting discrete pseudolabels on the unlabeled data. We avoid this difficulty because we consider directly learning continuous representations on the unlabeled data.
We apply spectral decomposition—a classical approach for graph partitioning, also known as spectral clustering [37, 39] in machine learning—to the adjacency matrix defined on the population augmentation graph. We form a matrix where the top-k eigenvectors are the columns and interpret each row of the matrix as the representation (in Rk) of an example. Somewhat surprisingly, we show that this feature extractor can be also recovered (up to some linear transformation) by minimizing the following population objective which is similar to the standard contrastive loss (Section 3.2):
L(f) = 2 · Ex,x+ ⇥ f(x)>f(x+) ⇤ + Ex,x0 ⇥ f(x)>f(x0) 2 ⇤ ,
where (x, x+) is a pair of augmentations of the same data, (x, x0) is a pair of independently random augmented data, and f is a parameterized function from augmented data to Rk. Figure 1 (right) illustrates the relationship between the eigenvector matrix and the learned representations. We call this loss the population spectral contrastive loss.
We analyze the linear classification performance of the representations learned by minimizing the population spectral contrastive loss. Our main result (Theorem 3.7) shows that when the representation dimension exceeds the maximum number of disconnected sub-graphs, linear classification with learned representations is guaranteed to have a small error. Our theorem reveals a trend that a larger representation dimension is needed when there are a larger number of disconnected sub-graphs. Our analysis relies on novel techniques tailored to linear probe performance, which have not been studied in the spectral graph theory community to the best of our knowledge.
The spectral contrastive loss also works on empirical data. Since our approach optimizes parametric loss functions, guarantees involving the population loss can be converted to finite sample results using off-the-shelf generalization bounds. The sample complexity is polynomial in the Rademacher complexity of the model family and other relevant parameters (Theorem 4.1 and Theorem 4.2).
In summary, our main theoretical contributions are: 1) we propose a simple contrastive loss motivated by spectral decomposition of the population data graph, 2) under simple and realistic assumptions, we provide downstream classification guarantees for the representation learned by minimizing this loss on population data, and 3) our analysis is easily applicable to deep networks with polynomial unlabeled samples via off-the-shelf generalization bounds.
In addition, we implement and test the proposed spectral contrastive loss on standard vision benchmark datasets. We demonstrate that the features learned by our algorithm can match or outperform several strong baselines [12, 14, 15, 21] when evaluated using a linear probe.

heading: 2 Additional related works

text:
Empirical works on self-supervised learning. Self-supervised learning algorithms have been shown to successfully learn representations that benefit downstream tasks [4, 6, 10, 12, 13, 15, 22– 24, 35, 38, 41, 42, 50–52]. Many recent self-supervised learning algorithms learn features with siamese networks [8], where two neural networks of shared weights are applied to pairs of augmented data. Introducing asymmetry to siamese networks either with a momentum encoder like BYOL [21] or by stopping gradient propagation for one branch of the siamese network like SimSiam [14] has been shown to effectively avoid collapsing. Contrastive methods [12, 15, 22] minimize the InfoNCE loss [38], where two views of the same data are attracted while views from different data are repulsed.
Theoretical works on self-supervised learning. In addition to works [3, 28, 44, 45] discussed in the introduction, several other works [5, 43, 47, 48] also theoretically study self-supervised learning. The work Tsai et al. [47] prove that self-supervised learning methods can extract task-relevant information and discard task-irrelevant information, but lacks guarantees for solving downstream tasks efficiently with simple (e.g., linear) models. Tian et al. [43] study why non-contrastive selfsupervised learning methods can avoid feature collapse. Cai et al. [9] analyze domain adaptation algorithms for subpopulation shift with a similar expansion condition as [49] while also allowing disconnected parts within each class, but require access to ground-truth labels during training. In contrast, our algorithm doesn’t need labels during pre-training.

heading: 3 Spectral contrastive learning on population data

text:
In this section, we introduce our theoretical framework, the spectral contrastive loss, and the main analysis of the performance of the representations learned on population data.
We use X to denote the set of all natural data (raw inputs without augmentation). We assume that each x̄ 2 X belongs to one of r classes, and let y : X ! [r] denote the ground-truth (deterministic) labeling function. Let PX be the population distribution over X from which we draw training data and test our final performance. For the ease of exposition, we assume X to be a finite but exponentially large set (e.g., all real vectors in Rd with bounded precision).1
We next formulate data augmentations. Given a natural data sample x̄ 2 X , we use A(·|x̄) to denote the distribution of its augmentations. For instance, when x̄ represents an image, A(·|x̄) can be the distribution of common augmentations [12] that includes Gaussian blur, color distortion and random cropping. We use X to denote the set of all augmented data, which is the union of supports of all A(·|x̄) for x̄ 2 X . As with X , we also assume that X is a finite but exponentially large set, and denote N = |X |.
We will learn an embedding function f : X ! Rk, and then evaluate its quality by the minimum error achieved with a linear probe. Concretely, a linear classifier has weights B 2 Rk⇥r and predicts gf,B(x) = argmaxi2[r](f(x) > B)i for an augmented data x (argmax breaks tie arbitrarily). Then, given a natural data sample x̄, we ensemble the predictions on augmented data and predict: ḡf,B(x̄) := argmax
i2[r] Pr x⇠A(·|x̄) [gf,B(x) = i] .
Define the linear probe error as the error of the best possible linear classifier on the representations: E(f) := min
B2Rk⇥r Pr x̄⇠PX [y(x̄) 6= ḡf,B(x̄)] (1)

heading: 3.1 Augmentation graph and spectral decomposition

text:
Our approach is based on the central concept of population augmentation graph, denoted by G(X , w), where the vertex set is all augmentation data X and w denotes the edge weights defined below. For any two augmented data x, x0 2 X , define the weight wxx0 as the marginal probability of generating the pair x and x0 from a random natural data x̄ ⇠ PX :
wxx0 := Ex̄⇠PX [A(x|x̄)A(x 0|x̄)] (2)
Therefore, the weights sum to 1 because the total probability mass is 1: P
x,x02X wxx0 = 1. The relative magnitude intuitively captures the closeness between x and x0 with respect to the augmentation transformation. For most of the unrelated x and x0, the value wxx0 will be significantly smaller than the average value. For example, when x and x0 are random croppings of a cat and a dog respectively, wxx0 will be essentially zero because no natural data can be augmented into both x and x0. On the other hand, when x and x0 are very close in `2-distance or very close in `2-distance up to color distortion, wxx0 is nonzero because they may be augmentations of the same image with Gaussian blur and color distortion. We say that x and x0 are connected with an edge if wxx0 > 0. See Figure 1 (left) for more illustrations.
Given the structure of the population augmentation graph, we apply spectral decomposition to the population graph to construct principled embeddings. The eigenvalue problems are closely related to graph partitioning as shown in spectral graph theory [17] for both worst-case graphs [11, 25, 29, 33] and random graphs [1, 30, 34]. In machine learning, spectral clustering [37, 39] is a classical algorithm that learns embeddings by eigendecomposition on an empirical distance graph and invoking k-means on the embeddings.
We will apply eigendecomposition to the population augmentation graph (and then later use linear probe for classification). Let wx = P x02X wxx0 be the total weights associated to x, which is often viewed as an analog of the degree of x in weighted graph. A central object in spectral graph theory is the so-called normalized adjacency matrix:
A := D 1/2AD 1/2 (3) 1This allows us to use sums instead of integrals and avoid non-essential nuances related to calculus.
where A 2 RN⇥N is adjacency matrix with entires Axx0 = wxx0 and D 2 RN⇥N is a diagonal matrix with Dxx = wx.2
Standard spectral graph theory approaches produce vertex embeddings as follows. Let 1, 2, · · · , k be the k largest eigenvalues of A, and v1, v2, · · · , vk be the corresponding unit-norm eigenvectors. Let F ? = [v1, v2, · · · , vk] 2 RN⇥k be the matrix that collects these eigenvectors in columns, and we refer to it as the eigenvector matrix. Let u⇤x 2 Rk be the x-th row of the matrix F ?. It turns out that u ⇤ x’s can serve as desirable embeddings of x’s because they exhibit clustering structure in Euclidean space that resembles the clustering structure of the graph G(X , w).

heading: 3.2 From spectral decomposition to spectral contrastive learning

text:
The embeddings u⇤x obtained by eigendecomposition are nonparametric—a k-dimensional parameter is needed for every x—and therefore cannot be learned with a realistic amount of data. The embedding matrix F ? cannot be even stored efficiently. Therefore, we will instead parameterize the rows of the eigenvector matrix F ? as a neural net function, and assume embeddings u⇤x can be represented by f(x) for some f 2 F , where F is the hypothesis class containing neural networks. As we’ll show in Section 4, this allows us to leverage the extrapolation power of neural networks and learn the representation on a finite dataset.
Next, we design a proper loss function for the feature extractor f , such that minimizing this loss could recover F ? up to some linear transformation. As we will show in Section 4, the resulting population loss function on f also admits an unbiased estimator with finite training samples. Let F be an embedding matrix with ux on the x-th row, we will first design a loss function of F that can be decomposed into parts about individual rows of F .
We employ the following matrix factorization based formulation for eigenvectors. Consider the objective
min F2RN⇥k
Lmf(F ) := A FF> 2 F . (4)
By the classical theory on low-rank approximation (Eckart–Young–Mirsky theorem [19]), any minimizer bF of Lmf(F ) contains scaling of the largest eigenvectors of A up to a right transformation— for some orthonormal matrix R 2 Rk⇥k, we have bF = F ? · diag([p 1, . . . , p k])Q. Fortunately, multiplying the embedding matrix by any matrix on the right and any diagonal matrix on the left does not change its linear probe performance, which is formalized by the following lemma. Lemma 3.1. Consider an embedding matrix F 2 RN⇥k and a linear classifier B 2 Rk⇥r. Let D 2 RN⇥N be a diagonal matrix with positive diagonal entries and Q 2 Rk⇥k be an invertible matrix. Then, for any embedding matrix eF = D · F ·Q, the linear classifier B̃ = Q 1B on eF has the same prediction as B on F . As a consequence, we have
E(F ) = E( eF ). (5)
where E(F ) denotes the linear probe performance when the rows of F are used as embeddings.
The proof can be found in Section C.1.
The main benefit of objective Lmf(F ) is that it’s based on the rows of F . Recall that vectors ux are the rows of F . Each entry of FF> is of the form u>x xx0 , and thus Lmf(F ) can be decomposed into a sum of N2 terms involving terms u>x ux0 . Interestingly, if we reparameterize each row ux by w 1/2 x f(x), we obtain a very similar loss function for f that resembles the contrastive learning loss used in practice [12] as shown below in Lemma 3.2. See Figure 1 (right) for an illustration of the relationship between the eigenvector matrix and the representations learned by minimizing this loss.
We formally define the positive and negative pairs to introduce the loss. Let x̄ ⇠ PX be a random natural data and draw x ⇠ A(·|x̄) and x+ ⇠ A(·|x̄) independently to form a positive pair (x, x+). Draw x̄0 ⇠ PX and x0 ⇠ A(·|x̄0) independently with x̄, x, x+. We call (x, x0) a negative pair.3
2We index the matrix A, D by (x, x0) 2 X ⇥ X . Generally we index N -dimensional axis by x 2 X . 3Though x and x0 are simply two independent draws, we call them negative pairs following the literature [3].
Lemma 3.2 (Spectral contrastive loss). Recall that ux is the x-th row of F . Let ux = w 1/2 x f(x) for some function f . Then, the loss function Lmf(F ) is equivalent to the following loss function for f , called spectral contrastive loss, up to a additive constant:
Lmf(F ) = L(f) + const
where L(f) , 2 · Ex,x+ ⇥ f(x)>f(x+) ⇤ + Ex,x0 h f(x)>f(x0) 2i (6)
The proof can be found in Section C.1.
We note that spectral contrastive loss is similar to many popular contrastive losses [12, 38, 40, 50]. For instance, the contrastive loss in SimCLR [12] can be rewritten as (with simple algebraic manipulation)
f(x)>f(x+) + log exp f(x)>f(x+) +
nX
i=1
exp f(x)>f(xi)
!
.
Here x and x+ are a positive pair and x1, · · · , xn are augmentations of other data. Spectral contrastive loss can be seen as removing f(x)>f(x+) from the second term, and replacing the log sum of exponential terms with the average of the squares of f(x)>f(xi). We will show in Section 6 that our loss has a similar empirical performance as SimCLR without requiring a large batch size.

heading: 3.3 Theoretical guarantees for spectral contrastive loss on population data

text:
In this section, we introduce the main assumptions on the data and state our main theoretical guarantee for spectral contrastive learning on population data.
To formalize the idea that G cannot be partitioned into too many disconnected sub-graphs, we introduce the notions of Dirichlet conductance and sparsest m-partition, which are standard in spectral graph theory. Dirichlet conductance represents the fraction of edges from S to its complement: Definition 3.3 (Dirichlet conductance). For a graph G = (X , w) and a subset S ✓ X , we define the Dirichlet conductance of S as
G(S) :=
P x2S,x0 /2S wxx0P
x2S wx .
We note that when S is a singleton, there is G(S) = 1 due to the definition of wx. We introduce the sparsest m-partition to represent the number of edges between m disjoint subsets. Definition 3.4 (Sparsest m-partition). Let G = (X , w) be the augmentation graph. For an integer m 2 [2, |X |], we define the sparsest m-partition as
⇢m := min S1,··· ,Sm max{ G(S1), . . . , G(Sm)}
where S1, · · · , Sm are non-empty sets that form a partition of X .
When r is the number of underlying classes, we might expect ⇢r ⇡ 0 since the augmentations from different classes almost compose a disjoint r-way partition of X . However, for m > r, we can expect ⇢m to be much larger. For instance, in the extreme case when m = |X | = N , every set Si is a singleton, which implies that ⇢N = 1.
Next, we formalize the assumption that very few edges cross different ground-truth classes. It turns out that it suffices to assume that the labels are recoverable from the augmentations (which is also equivalent to that two examples in different classes can rarely be augmented into the same point). Assumption 3.5 (Labels are recoverable from augmentations). Let x̄ ⇠ PX and y(x̄) be its label. Let the augmentation x ⇠ A(·|x̄). We assume that there exists a classifier g that can predict y(x̄) given x with error at most ↵. That is, g(x) = y(x̄) with probability at least 1 ↵.
We also introduce the following assumption which states that some universal minimizer of the population spectral contrastive loss can be realized by the hypothesis class. Assumption 3.6 (Realizability). Let F be a hypothesis class containing functions from X to Rk. We assume that at least one of the global minima of L(f) belongs to F .
Our main theorem bound from above the linear probe error of the features learned by minimizing the population spectral contrastive loss. Theorem 3.7. Assume the representation dimension k 2r and Assumption 3.5 holds for ↵ > 0. Let F be a hypothesis class that satisfies Assumption 3.6 and let f⇤pop 2 F be a minimizer of L(f). Then, we have
E(f⇤pop)  eO ⇣ ↵/⇢ 2 bk/2c ⌘ .
Here we use eO(·) to hide universal constant factors and logarithmic factor in k. We note that ↵ = 0 when augmentations from different classes are perfectly disconnected in the augmentation graph, in which case the above theorem guarantees the exact recovery of the ground truth. Generally, we expect ↵ to be an extremely small constant independent of k, whereas ⇢bk/2c increases with k and can be much larger than ↵ when k is reasonably large. For instance, when there are t sub-graphs that have sufficient inner connections, we expect ⇢t+1 to be on the order of a constant because any t+ 1 partition needs to break one sub-graph into two pieces and incur a large conductance. We characterize the ⇢k’s growth on more concrete distributions in the next subsection.
Previous works on graph partitioning [2, 29, 31] often analyze the so-called rounding algorithms that conduct clustering based on the representations of unlabeled data and do not analyze the performance of linear probe (which has access to labeled data). These results provide guarantees on the approximation ratio—the ratio between the conductance of the obtained partition to the best partition—which may depend on graph size [2] that can be exponentially large in our setting. The approximation ratio guarantee does not lead to a guarantee on the representations’ performance on downstream tasks. Our guarantees are on the linear probe accuracy on the downstream tasks and independent of the graph size. We rely on the formulation of the downstream task’s labeling function (Assumption 3.5) as well as a novel analysis technique that characterizes the linear structure of the representations. In Section C, we provide the proof of Theorem 3.7 as well as its more generalized version where k/2 is relaxed to be any constant fraction of k.

heading: 3.4 Provable instantiation of Theorem 3.7 to mixture of manifold data

text:
In this section, we exemplify Theorem 3.7 on an example where the natural data distribution is a mixture of manifolds, and the augmentation transformation is adding Gaussian noise. Example 3.8 (Mixture of manifolds). Suppose PX is mixture of r  d distributions P1, · · · , Pr, where each Pi is generated by some -bi-Lipschitz 4 generator Q : Rd0 ! Rd on some latent variable z 2 Rd0 with d0  d which as a mixture of Gaussian distribution:
x ⇠ Pi () x = Q(z), z ⇠ N (µi, 1
d0 · Id0⇥d0).
Let the data augmentation of a natural data sample x̄ is x̄+ ⇠ where ⇠ ⇠ N (0, 2
d · Id⇥d) is isotropic Gaussian noise with 0 < . 1p
d . We also assume mini 6=j kµi µjk2 & · p log dp d0 .
Let the ground-truth label be the most likely mixture index i that generates x: y(x) := argmaxi Pi(x). We note that the intra-class distance in the latent space is on the scale of ⌦(1), which can be much larger than the distance between class means which is assumed to be & ·
p log dp d0
. Therefore, distancebased clustering algorithms do not apply. We apply Theorem 3.7 and get the following theorem: Theorem 3.9. When k 2r + 2, Example 3.8 satisfies Assumption 3.5 with ↵  1poly(d) , and has
⇢bk/2c & pd . As a consequence, the error bound is E(f ⇤ pop)  Õ
⇣ 2
2·poly(d)
⌘ .
The theorem above guarantees small error even when is polynomially small. In this case, the augmentation noise has a much smaller scale than the data (which is at least on the order of 1/). This suggests that contrastive learning can non-trivially leverage the structure of the underlying data and learn good representations with relatively weak augmentation. The proof can be found in Section D.
4A  bi-Lipschitz function satisfies 1 kf(x) f(y)k2  kx yk2   kf(x) f(y)k2.

heading: 4 Finite-sample generalization bounds

text:
In Section 3, we provide guarantees for spectral contrastive learning on population data. In this section, we show that these guarantees can be naturally extended to the finite-sample regime with standard concentration bounds. In particular, given a training dataset {x̄1, x̄2, · · · , x̄n} with x̄i ⇠ PX , we learn a feature extractor by minimizing the following empirical spectral contrastive loss:
bLn(f) := 2
n
nX
i=1
E x⇠A(·|x̄i) x+⇠A(·|x̄i)
⇥ f(x)>f(x+) ⇤ +
1 n(n 1) X
i 6=j E x⇠A(·|x̄i) x0⇠A(·|x̄j)
h f(x)>f(x0) 2i .
It is worth noting that bLn(f) is an unbiased estimator of the population spectral contrastive loss L(f). (See Claim E.2 for a proof.) Therefore, we can derive generalization bounds via off-the-shelf concentration inequalities. Let F be a hypothesis class containing feature extractors from X to Rk. We extend Rademacher complexity to function classes with highdimensional outputs and define the Rademacher complexity of F on n data as bRn(F) := maxx1,··· ,xn2X E h supf2F,i2[k] 1 n ⇣Pn j=1 jfi(xj) ⌘i ,
where is a uniform random vector in { 1, 1}n and fi(z) is the i-th dimension of f(z). Recall that f⇤pop 2 F is a minimizer of L(f). The following theorem with proofs in Section E.1 bounds the population loss of a feature extractor trained with finite data: Theorem 4.1. For some  > 0, assume kf(x)k1   for all f 2 F and x 2 X . Let f⇤pop 2 F be a minimizer of the population loss L(f). Given a random dataset of size n, let f̂emp 2 F be a minimizer of empirical loss bLn(f). Then, with probability at least 1 over the randomness of data, we have
L(f̂emp)  L(f⇤pop) + c1 · bRn/2(F) + c2 · r log 2/
n +
! ,
where constants c1 . k22 + k and c2 . k2 + k24.
We can apply Theorem 4.1 to any hypothesis class F of interest (e.g., deep neural networks) and plug in off-the-shelf Rademacher complexity bounds. For instance, in Section E.2 we give a corollary of Theorem 4.1 when F contains deep neural networks with ReLU activation. The theorem above shows that we can achieve near-optimal population loss by minimizing empirical loss up to some small excess loss. The following theorem characterizes how the error propagates to the linear probe performance mildly under some spectral gap conditions. Theorem 4.2. Assume representation dimension k 4r + 2, Assumption 3.5 holds for ↵ > 0 and Assumption 3.6 holds. Recall i be the i-th largest eigenvalue of the normalized adjacency matrix. Then, for any ✏ < 2 k and f̂emp 2 F such that L(f̂emp) < L(f⇤pop) + ✏, we have:
E(f̂emp) . ↵ ⇢2bk/2c · log k + k✏ 2 ,
where := b3k/4c k is the eigenvalue gap between the b3k/4c-th and the k-th eigenvalue.
This theorem shows that the error on the downstream task only grows linearly with the error ✏ during pretraining. We can relax Assumption 3.6 to approximate realizability in the sense that F contains some sub-optimal feature extractor under the population spectral loss and pay an additional error term in the linear probe error bound. The proof of Theorem 4.2 can be found in Section E.3.

heading: 5 Guarantee for learning linear probe with labeled data

text:
In this section, we provide theoretical guarantees for learning a linear probe with labeled data. Theorem 3.7 guarantees the existence of a linear probe that achieves a small downstream classification error. However, a priori it is unclear how large the margin of the linear classifier can be, so it is hard to apply margin theory to provide generalization bounds for 0-1 loss. We could in principle control the margin of the linear head, but using capped quadratic loss turns out to suffice and mathematically more
convenient. We learn a linear head with the following capped quadratic loss: given a tuple (z, y(x̄)) where z 2 Rk is a representation of augmented data x ⇠ A(·|x̄) and y(x̄) 2 [r] is the label of x̄, for a linear probe B 2 Rk⇥r we define loss `((z, y(x̄)), B) := Pr i=1 min B > z ~y(x̄) 2 i , 1 , where ~y(x̄) is the one-hot embedding of y(x̄) as a r-dimensional vector (1 on the y(x̄)-th dimension, 0 on other dimensions). This is a standard modification of quadratic loss in statistical learning theory that ensures the boundedness of the loss for the ease of analysis [36].
The following Theorem 5.1 provides a generalization guarantee for the linear classifier that minimizes capped quadratic loss on a labeled dataset. The key challenge of the proof is showing the existence of a small-norm linear head B that gives small population quadratic loss, which is not obvious from Theorem 3.7 where only small 0-1 error is guaranteed. Recall i is the i-th largest eigenvalue of the the normalized adjacency matrix. Given a labeled dataset {(x̄i, y(x̄i))}ni=1 where x̄i ⇠ PX and y(x̄i) is its label, we sample xi ⇠ A(·|x̄i) for i 2 [n]. Theorem 5.1. In the setting of Theorem 3.7, assume k C for some C > 0. Learn a linear probe bB 2 argminkBkF1/C Pn i=1 `((f ⇤ pop(xi), y(x̄i)), B) by minimizing the capped quadratic loss subject to a norm constraint. Then, with probability at least 1 over random data, we have
Pr x̄⇠PX
⇣ ḡf⇤pop, bB(x̄) 6= y(x̄) ⌘ . ↵
⇢2bk/2c · log k + r C · r k n + r log 1/ n .
Here the first term is the population error from Theorem 3.7. The last two terms are the generalization gap from standard concentration inequalities for linear classification and are small when the number of labeled data n is polynomial in the feature dimension k. We note that this result reveals a tradeoff when choosing the feature dimension k: when n is fixed, a larger k decreases the population contrastive loss while increases the generalization gap for downstream linear classification. The proof of Theorem 5.1 is in Section F.

heading: 6 Experiments

text:
We test spectral contrastive learning on benchmark vision datasets. We minimize the empirical spectral contrastive loss with an encoder network f and sample fresh augmentation in each iteration. The pseudo-code for the algorithm and more implementation details can be found in Section A.
Encoder / feature extractor. The encoder f contains three components: a backbone network, a projection MLP and a projection function. The backbone network is a standard ResNet architecture. The projection MLP is a fully connected network with BN applied to each layer, and ReLU activation applied to each except for the last layer. The projection function takes a vector and projects it to a sphere ball with radius pµ, where µ > 0 is a hyperparameter that we tune in experiments. We find that using a projection MLP and a projection function improves the performance.
Linear evaluation protocol. Given the pre-trained encoder network, we follow the standard linear evaluation protocol [14] and train a supervised linear classifier on frozen representations, which are from the ResNet’s global average pooling layer.
Results. We report the accuracy on CIFAR-10/100 [26] and Tiny-ImageNet [27] in Table 1. Our empirical results show that spectral contrastive learning achieves better performance than two popular baseline algorithms SimCLR [12] and SimSiam [14]. In Table 2 we report results on ImageNet [18] dataset, and show that our algorithm achieves similar performance as other state-of-the-art methods. We note that our algorithm is much more principled than previous methods and doesn’t rely on large batch sizes (SimCLR [12]), momentum encoders (BYOL [21] and MoCo [22]) or additional tricks such as stop-gradient (SimSiam [14]).

heading: 7 Conclusion

text: In this paper, we present a novel theoretical framework of self-supervised learning and provide provable guarantees for the learned representations on downstream linear classification. We hope our study can facilitate future theoretical analyses of self-supervised learning and inspire new practical algorithms. For instance, one interesting future direction is to test the topology of the augmentation graph on empirical data distributions and design algorithms using tools from graph theory.

heading: Acknowledgements

text: We thank Margalit Glasgow, Ananya Kumar, Jason D. Lee, Sang Michael Xie, and Guodong Zhang for helpful discussions. CW acknowledges support from an NSF Graduate Research Fellowship. TM acknowledges support of Google Faculty Award and NSF IIS 2045685. We also acknowledge the support of HAI and the Google Cloud. Toyota Research Institute ("TRI") provided funds to assist the authors with their research but this article solely reflects the opinions and conclusions of its authors and not TRI or any other Toyota entity.

title: Provable Guarantees for Self-Supervised Deep Learning with Spectral Contrastive Loss