summaries: The paper considers the task of obtaining additional control over a pre-trained generative model without re-training it. Specifically, it aims to generate conditional samples from an unconditional StyleGAN by only training a classifier with respect to the conditioning information. To achieve this, it considers the conditional distribution of StyleGAN latent codes. By Bayes Theorem, this conditional distribution is proportional to the unconditional distribution of the latent codes (which has a known standard normal distribution) times the distribution of conditioning information given the corresponding sample. The latter is modeled as a Gibbs distribution with an energy derived from the classifier. Two approaches to sample from this unnormalized distribution are evaluated: (i) Langevin sampling and (ii) Performing gradient descent on the classifier-derived energy. Experiments demonstrate good performance of this approach with significant improvements over StyleFlow for combinations of conditionings which are not in the training data.

summaries: This paper studies conditional learning by training an EBM in the latent space of a pre-trained top-down generator such as StyleGANs. The EBM is a joint distribution of data and attributes together, and sampling from it is formulated as solving an ordinary differential equation. Experimental results show that the method outperforms the state-of-the-art in both conditional sampling and sequential editing. The contribution of the paper lies on obtaining state-of-the-art performance of conditional learning by combining existing technologies, such as StyleGAN, latent space energy-based learning, and score-based generative model via stochastic differential equations.

summaries: Using labelled real data, an EBM is trained in the latent space of a pretrained GAN, rather than directly in pixel space. An ODE solver is used to sample from the EBM rather than Langevin Dynamics, which is shown to be more robust to hyperparameter settings. Using properties of the EBM, conditional generation compositional editing can be demonstrated, including in zero-shot scenarios.

summaries: The paper proposes an effective and efficient approach to conditional image generation using GANs. Two main advances are described 1) The use of EBM in latent space that don’t require the retraining of image generators. This allows for very efficient training since the generator can be held fixed. 2) The use of an ODE solver for sampling that is less sensitive to hyperparameters than traditional LD solvers.  Results are shown on CIFAR-10 and FFHQ datasets. This includes results on sequential generation, compositional generation and zero-shot generation.

abstractText: Controllable generation is one of the key requirements for successful adoption of deep generative models in real-world applications, but it still remains as a great challenge. In particular, the compositional ability to generate novel concept combinations is out of reach for most current models. In this work, we use energybased models (EBMs) to handle compositional generation over a set of attributes. To make them scalable to high-resolution image generation, we introduce an EBM in the latent space of a pre-trained generative model such as StyleGAN. We propose a novel EBM formulation representing the joint distribution of data and attributes together, and we show how sampling from it is formulated as solving an ordinary differential equation (ODE). Given a pre-trained generator, all we need for controllable generation is to train an attribute classifier. Sampling with ODEs is done efficiently in the latent space and is robust to hyperparameters. Thus, our method is simple, fast to train, and efficient to sample. Experimental results show that our method outperforms the state-of-the-art in both conditional sampling and sequential editing. In compositional generation, our method excels at zero-shot generation of unseen attribute combinations. Also, by composing energy functions with logical operators, this work is the first to achieve such compositionality in generating photo-realistic images of resolution 1024⇥1024.

authors:
R. Abdal
P. Zhu
N.J. Mitra
P. Wonka

title: Styleflow: Attribute-conditioned exploration of stylegan-generated images using conditional continuous normalizing flows

venue: ACM Trans. Graph

authors:
J. Aneja
A. Schwing
J. Kautz
A. Vahdat

title: A contrastive learning approach for training variational autoencoder priors

venue: In Neural Information Processing Systems (NeurIPS)

authors:
A.F. Ansari
M.L. Ang
H. Soh

title: Refining deep generative models via discriminator gradient flow

venue: arXiv preprint arXiv:2012.00780

authors:
M. Arbel
L. Zhou
A. Gretton

title: Generalized energy based models

venue: In International Conference on Learning Representations

title: The tools of generative art, from flash to neural networks. Art in America

authors:
T. Che
R. Zhang
J. Sohl-Dickstein
H. Larochelle
L. Paull
Y. Cao
Y. Bengio

title: Your gan is secretly an energy-based model and you should use discriminator driven latent sampling. arXiv preprint arXiv:2003.06060

authors:
R.T.Q. Chen
Y. Rubanova
J. Bettencourt
D. Duvenaud

title: Neural ordinary differential equations

venue: Advances in Neural Information Processing Systems

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. PMLR

authors:
X. Chen
Y. Duan
R. Houthooft
J. Schulman
I. Sutskever
P. Abbeel

title: Infogan: Interpretable representation learning by information maximizing generative adversarial nets

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

authors:
Y. Choi
M. Choi
M. Kim
Ha
J.-W
S. Kim
J. Choo

title: Stargan: Unified generative adversarial networks for multi-domain image-to-image translation

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition

authors:
Y. Choi
Y. Uh
J. Yoo
Ha
J.-W

title: Stargan v2: Diverse image synthesis for multiple domains

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

authors:
Y. Du
S. Li
I. Mordatch

title: Compositional visual generation with energy based models

venue: In Advances in Neural Information Processing Systems

authors:
Y. Du
S. Li
J. Tenenbaum
I. Mordatch

title: Improved contrastive divergence training of energy based models

venue: arXiv preprint arXiv:2012.01316

authors:
Y. Du
I. Mordatch

title: Implicit generation and generalization in energy-based models. arXiv preprint arXiv:1903.08689

authors:
A. Gabbay
N. Cohen
Y. Hoshen

title: An image is worth more than a thousand words: Towards disentanglement in the wild. arXiv preprint arXiv:2106.15610

authors:
L. Goetschalckx
A. Andonian
A. Oliva
P. Isola

title: Ganalyze: Toward visual definitions of cognitive image properties

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

authors:
I. Goodfellow
J. Pouget-Abadie
M. Mirza
B. Xu
D. Warde-Farley
S. Ozair
A. Courville
Y. Bengio

title: Generative adversarial nets. In Advances in neural information processing systems (NeurIPS)

authors:
W. Grathwohl
Wang
K.-C
Jacobsen
J.-H
D. Duvenaud
M. Norouzi
K. Swersky

title: Your classifier is secretly an energy based model and you should treat it like one

venue: In Internal Conference of Learning Representations

authors:
T. Han
Y. Lu
Zhu
S.-C
Y.N. Wu

title: Alternating back-propagation for generator network

venue: In Proceedings of the AAAI Conference on Artificial Intelligence

authors:
E. Härkönen
A. Hertzmann
J. Lehtinen
S. Paris

title: Ganspace: Discovering interpretable gan controls

authors:
M. Heusel
H. Ramsauer
T. Unterthiner
B. Nessler
S. Hochreiter

title: Gans trained by a two time-scale update rule converge to a local nash equilibrium

venue: Advances in neural information processing systems

title: Products of experts

venue: In Proceedings of the Ninth International Conference on Artificial Neural Networks

authors:
J. Ho
A. Jain
P. Abbeel

title: Denoising diffusion probabilistic models. In Advances in neural information processing systems (NeurIPS)

authors:
A.G. Howard
M. Zhu
B. Chen
D. Kalenichenko
W. Wang
T. Weyand
M. Andreetto
H. Adam

title: Mobilenets: Efficient convolutional neural networks for mobile vision applications

venue: arXiv preprint arXiv:1704.04861

authors:
G. Huang
Z. Liu
L. Van Der Maaten
K.Q. Weinberger

title: Densely connected convolutional networks

authors:
A. Jahanian
L. Chai
P. Isola

title: On the" steerability" of generative adversarial networks

authors:
T. Karras
M. Aittala
J. Hellsten
S. Laine
J. Lehtinen
T. Aila

title: Training generative adversarial networks with limited data

authors:
T. Karras
S. Laine
T. Aila

title: A style-based generator architecture for generative adversarial networks

authors:
T. Karras
S. Laine
M. Aittala
J. Hellsten
J. Lehtinen
T. Aila

title: Analyzing and improving the image quality of StyleGAN

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

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980

authors:
D.P. Kingma
S. Mohamed
D.J. Rezende
M. Welling

title: Semi-supervised learning with deep generative models. In Advances in neural information processing systems

authors:
D.P. Kingma
M. Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114

authors:
A. Krizhevsky
G. Hinton

title: Learning multiple layers of features from tiny images. (Technical Report

authors:
Y. LeCun
S. Chopra
R. Hadsell
M. Ranzato
F. Huang

title: A tutorial on energy-based learning

venue: Predicting structured data,

authors:
Z. Liu
P. Luo
X. Wang
X. Tang

title: Deep learning face attributes in the wild

venue: In Proceedings of International Conference on Computer Vision (ICCV)

authors:
M. Mirza
S. Osindero

title: Conditional generative adversarial nets. arXiv preprint arXiv:1411.1784

authors:
A. Nguyen
A. Dosovitskiy
J. Yosinski
T. Brox
J. Clune

title: Synthesizing the preferred inputs for neurons in neural networks via deep generator networks. In Advances in neural information processing systems

authors:
W. Nie
T. Karras
A. Garg
S. Debnath
A. Patney
A. Patel
A. Anandkumar

title: Semisupervised stylegan for disentanglement learning

venue: In International Conference on Machine Learning. PMLR

title: Stochastic differential equations. In Stochastic differential equations, pages 65–84

authors:
B. Pang
T. Han
E. Nijkamp
Zhu
S.-C
Y.N. Wu

title: Learning latent space energy-based prior model

title: Semantic image synthesis with spatially-adaptive normalization

authors:
O. Patashnik
Z. Wu
E. Shechtman
D. Cohen-Or
D. Lischinski

title: Styleclip: Text-driven manipulation of stylegan imagery

authors:
A. Radford
J.W. Kim
C. Hallacy
A. Ramesh
G. Goh
S. Agarwal
G. Sastry
A. Askell
P. Mishkin
J Clark

title: Learning transferable visual models from natural language supervision

venue: arXiv preprint arXiv:2103.00020

authors:
S. Reed
Z. Akata
X. Yan
L. Logeswaran
B. Schiele
H. Lee

title: Generative adversarial text to image synthesis

title: Euclidean, metric, and Wasserstein} gradient flows: an overview

venue: Bulletin of Mathematical Sciences,

authors:
Y. Shen
B. Zhou

title: Closed-form factorization of latent semantics in gans

venue: In CVPR

authors:
K. Sohn
D. Berthelot
Li
C.-L
Z. Zhang
N. Carlini
E.D. Cubuk
A. Kurakin
H. Zhang
C. Raffel

title: Fixmatch: Simplifying semi-supervised learning with consistency and confidence

venue: arXiv preprint arXiv:2001.07685

authors:
Y. Song
S. Ermon

title: Generative modeling by estimating gradients of the data distribution

title: Improved techniques for training score-based generative models. In Advances in neural information processing systems

authors:
Y. Song
D.P. Kingma

title: How to train your energy-based models. arXiv preprint arXiv:2101.03288

authors:
Y. Song
J. Sohl-Dickstein
D.P. Kingma
A. Kumar
S. Ermon
B. Poole

title: Score-based generative modeling through stochastic differential equations

authors:
C. Vaccari
A. Chadwick

title: Deepfakes and disinformation: Exploring the impact of synthetic political video on deception, uncertainty, and trust in news

venue: Social Media+ Society

authors:
A. Vahdat
E. Andriyash
W.G. Macready

title: DVAE#: Discrete variational autoencoders with relaxed Boltzmann priors

authors:
A. Vahdat
J. Kautz

title: NVAE: A deep hierarchical variational autoencoder

authors:
A. Vahdat
K. Kreis
J. Kautz

title: Score-based generative modeling in latent space. In Neural Information Processing Systems (NeurIPS)

authors:
A. Vahdat
W.G. Macready
Z. Bian
A. Khoshaman
E. Andriyash

title: DVAE++: Discrete variational autoencoders with overlapping transformations

venue: In International Conference on Machine Learning (ICML)

authors:
M. Welling
Y.W. Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: In Proceedings of the 28th international conference on machine learning (ICML)

authors:
Z. Xiao
K. Kreis
J. Kautz
A. Vahdat

title: Vaebm: A symbiosis between variational autoencoders and energy-based models

authors:
J. Xie
Y. Lu
Zhu
S.-C
Y. Wu

title: A theory of generative convnet

authors:
J. Xie
Z. Zheng
X. Fang
Zhu
S.-C
Y.N. Wu

title: Cooperative training of fast thinking initializer and slow thinking solver for conditional learning

venue: IEEE Transactions on Pattern Analysis and Machine Intelligence

heading: 1 Introduction

text:
Deep generative learning has made tremendous progress in image synthesis. High-quality and photorealistic images can now be generated with generative adversarial networks (GANs) [17, 28, 29], score-based models [23, 51], and variational autoencoders (VAEs) [32, 54]. However, a key requirement for the success of these generative models in many real-world applications is the controllability of the generation process. Controllable generation, in particular the compositional ability to generate novel concept combinations, still remains a great challenge to current methods.
A common approach for achieving controllable generation is to train a conditional model [36, 31, 10, 41, 38], where the conditional information, such as semantic attributes, is specified during training to guide the generation. There are two major drawbacks with this approach: (i) Since the set of attributes is fixed and pre-defined, it is difficult to introduce new attributes to an existing model. This, in turn, limits the model’s compositional ability for unseen attributes and their novel combinations [12]. (ii) Training conditional generative models from scratch for new attributes is computationally expensive, and it can be challenging to reach to the same image quality of the uncontrollable counterparts [10, 41].
Recent works have attempted to overcome these issues by first training an unconditional generator, and then converting it to a conditional model with a small cost [26, 20, 1, 42]. This is often achieved by discovering semantically meaningful directions in the latent space of the unconditional model. This way, one would pay the most computational cost only once for training the unconditional model. However, these approaches often struggle with compositional generation, in particular with rare
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
combination of the attributes and the introduction of new attributes. An appealing solution to the compositionality problem is to use energy-based models (EBMs) for controllable generation [12, 18, 13]. This is due to the fact that energy functions representing different semantics can be combined together to form compositional image generators. However, existing approaches train EBMs directly in the pixel space, making them slow to sample and difficult to generate high-resolution images [13].
In this paper, we leverage the compositionality of EBMs and the generative power of state-of-the-art pre-trained models such as StyleGANs [28, 29] to achieve high-quality controllable and compositional generation. Particularly, we propose an EBM that addresses the problem of controlling an existing generative model instead of generating images directly with EBMs or improving the sampling quality.
Our main contributions are summarized as follows:
• We propose a novel formulation of a joint EBM in the latent space of pre-trained generative models for controllable generation. In particular, the EBM training reduces to training a classifier only, and it can be applied to any latent-variable generative model.
• Based on our EBM formulation, we propose a new sampling method that relies on ordinary different equation (ODE) solvers. The sample quality is robust to hyperparameters.
• We show that our method is fast to train and efficient to sample, and it outperforms previous methods in both conditional sampling and sequential editing.
• We show strong compositionality with zero-shot generation on unseen attribute combinaions, and with logical compositions of energy functions in generating high-quality images.
Specifically, we build an EBM in the joint space of data and attributes where the marginal data distribution is denoted by an implicit distribution (e.g., a pre-trained GAN generator), and the conditional distribution of attributes, given data, is represented by an attribute classifier. Using the reparameterization trick, we show that our EBM formulation becomes an energy function in the latent space where the latent distribution is known (i.e., a standard Gaussian). This way, we only need to train the classifier in the data space and do the sampling in the latent space. Because our method only requires training a classifier to add controllability, it is conceptually simple and fast-to-train.
For sampling, most existing EBMs rely on the Langevin dynamics (LD) sampling which is often computationally expensive and sensitive to the choice of sampling parameters. Instead, we build on the recent score-based model [51], and define the sampling process in the latent space using the corresponding probability flow ODEs, induced by the reverse diffusion process. Thus, we rely on an ODE solver for sampling from our EBM formulation in the latent space. With its adaptive step size, the ODE sampling in the latent space is efficient and robust to the sampling parameters.
In experiments, we show our method achieves higher image quality and better controllability in both the conditional sampling and sequential editing, compared to various baselines, including StyleFlow [1] and JEM [18]. For the training time, we show our method is 25⇥ faster than StyleFlow on the FFHQ data. For the inference time, our sampling is at least 49⇥ and 876⇥ faster than EBMs [18, 12] and score-based models [51] in the pixel space, respectively, on CIFAR-10. More importantly, our method excels on zero-shot generation with unseen attribute combinations where StyleFlow almost completely fails. By composing energy functions with logical operators [12], our method is the first to show such compositionality in generating photo-realistic images of resolution 1024⇥1024.

heading: 2 Method

text: In this section, we will first give an overview of EBMs and then propose our new EBM formulation in the latent space, based on which, we introduce a new sampling method through the ODE solver.

heading: 2.1 Energy-based models

text:
EBMs [34] represent data x 2 Rd by learning an unnormalized probability distribution pθ(x) / e Eθ(x), where Eθ(x) : R
d ! R is the energy function, parameterized by a neural network. To train an EBM on a data distribution pdata with maximum likelihood estimation (MLE), we can estimate the gradient of the data log-likelihood L(✓) = Ex⇠pdata [log pθ(x)] as [22]:
rL(✓) = Ex⇠pdata [rθEθ(x)] Ex⇠pθ [rθEθ(x)] (1)
To sample x from pθ for training and inference, Langevin dynamics (LD) [57] is applied as follows,
x0 ⇠ p0(x), xt+1 = xt ⌘
2 rxEθ(xt) + ✏t, ✏t ⇠ N(0, ⌘I) (2)
where p0 is the initial distribution, ⌘ is the step size. When ⌘ ! 0 and t ! 1, xt is guaranteed to follow the probability pθ(x) under some regularity conditions [50].

heading: 2.2 Modelling joint EBM in the latent space

text:
To formulate our EBM for controllable generation, we consider the following setting: We have the data x 2 X ✓ Rd and its attribute code c = {c1, · · · , cn} 2 C ✓ R
n, which is a n-dimensional vector with each entry ci representing either a discrete mi-class attribute where ci 2 {0, · · · ,mi 1}, or a continuous attribute where ci 2 R. The goal is to learn a generative model and semantically control its generated samples by manipulating its conditioning attribute code c. To begin with, we define the joint generative model on both the data and attribute as:
pθ(x, c) := pg(x)pθ(c|x) / pg(x)e Eθ(c|x) (3)
where pg(x) is an implicit distribution defined by a pre-trained generator g (such as GANs) in the form x = g(z) with the latent variable z 2 Z . And, pθ(c|x) is conditional distribution on attributes given x, modeled by a conditional energy function Eθ(c|x).
In this paper, we assume that the generator g is fixed and we only train the energy function Eθ(c|x). We also assume the attributes are conditionally independent, i.e., Eθ(c|x) = Pn
i=1 Eθ(ci|x). Since, our goal is to preserve the image generation quality, we design the joint in Eq. (3) such that the marginal data distribution pθ(x) from our joint satisfies pθ(x) = pg(x). This is obtained with an energy function that is normalized up to a constant. Thus, we define:
Eθ(ci|x) =
⇢
fi(x; ✓)[ci] + log P
ci exp (fi(x; ✓)[ci]) if ci is discrete,
1 2σ2 (ci fi(x; ✓)) 2 if ci is continuous (4)
where fi(x; ✓) is the output of a multi-class classifier mapping from X to R mi if the i-th attribute is discrete or a regression network mapping from X to R if it is continuous. Without loss of generality, we will always call fi(x; ✓) a classifier. Here,
2 is a hyperparameter to adjust the continuous energy. In practice, we set 2 = 0.01 after normalizing all the continuous attributes to [0, 1].
We cannot use Eq. (3) for sampling, since pg(x) is defined implicitly in GAN generators. However, in Appendix A.1, we show that using the reparameterization trick [6, 58], sampling x from pθ(x, c) is reparameterized to sampling z from pθ(z, c)/ e
Eθ(c|g(z))+log p(z), and then transferring the z samples to the data space using generator x=g(z). In most generative models, the prior distribution p(z) is a standard Gaussian. Thus, the joint becomes pθ(z, c)/e Eθ(z,c) with the energy function:
Eθ(z, c) = X
i
Eθ(ci|g(z)) + 1
2 kzk22 (5)
where each conditional energy function Eθ(ci|g(z)) is given by Eq. (4). For sampling, we can run LD of Eq. (2) in the z-space using the energy function above. At the end of a chain, we pass the final z samples to the generator g to get the x samples. Since our latent-space EBM is formulated for controllable generation, we term it as LAtent Controllable EBM (LACE) throughout this work.
Inspecting the joint energy Eθ(z, c) in Eq. (5) shows that when the unconditional generator g is fixed, the only trainable component is the classifier fi(x; ✓) that represents Eθ(ci|x). Therefore, unlike the previous joint EBMs [18] that train both pθ(x) and pθ(c|x), our method only needs to train a classifier for the conditional pθ(c|x), and does not require sampling with LD in the x-space during training. That is, for controllable generation, we can train the classifier only in the x-space and do the LD sampling in the z-space. This fact makes our method conceptually simple and easy to train.
Note that the considered x-space for training the classifier does not have to be the pixel space. Instead, we can choose any intermediate layer of the pre-trained generator as the x-space. Take StyleGAN2 [29] as an example, if we train the classifier in the w-space, then g(z) in Eq. (5) actually corresponds to the mapping network of StyleGAN2, and once we obtain the w samples via the LD sampler, we can pass them to the synthesis network of StyleGAN2 to get the final images. In other words, since data x is a deterministic function of z, the distribution of attributes is uniquely defined by z. We have the freedom to define pθ(ci|z) by applying the classifier on top of z or any representation extracted from z including the generator g(z) or its intermediate representations such as w-space in StyleGAN2.

heading: 2.3 Sampling through an ODE solver

text:
Sampling with LD requires hand-crafted tricks to speed up and stabilize its convergence [14, 18]. Even in the latent space, our experiments show that LD tends to be sensitive to its hyperparameters. In
this section, we introduce another sampling method, called ODE sampler, that relies on an ODE solver. Our method with the LD and ODE sampler is termed as LACE-LD and LACE-ODE, respectively.
Our idea is inspired by the prior work [51], which shows that controllable generation can be achieved by solving a conditional reverse-time SDE. Specifically, if we consider a Variance Preserving (VP) SDE [51], the forward SDE is defined as dx = 12 (t)xdt + p
(t)dw, where w is a standard Wiener process and x0 ⇠ pdata is a data sample. The scalar time-variant diffusion coefficient (t) has a linear form (t) = min + ( max min)t, where t 2 [0, T ]. Then, the conditional sampling from p0(x|c) is equivalent to solving the following reverse VP-SDE:
dx = 1
2 (t)[x+ 2rx log pt(x, c)]dt+
p
(t)dw̄ (6)
where w̄ is a standard reverse Wiener process when time flows backwards from T to 0, and pt(x, c) is the join data and attribute distribution at time t. Song et al. [51] show that there exist a corresponding ODE for the SDE in Eq. (6) which can be used for sampling. In Appendix A.2, using the reparameterization trick, we show that the ODE in the latent space for our model takes a simple form:
dz = 1
2 (t)
X
i
rzE(ci|g(z))dt (7)
with negative time increments from T to 0. To generate a sample x conditioned on c, all we need is to draw z(T )⇠N(0, I) and solve the ODE in Eq. (7) using a black-box solver. The final z(0) samples are transferred to the data space through the generator x= g(z(0)). Our main insight in Appendix A.2 that leads to the simple ODE is that the latent z in most generative models follows a standard Gaussian distribution, and diffusing it with VP-SDE will not change its distribution, i.e., pt(z(t))=N (0, I).
Remarks There are some key observations from our ODE formulation in Eq. (7):
(i) Connections to gradient flows. Similar to [45, 3] that build connections between SDE/ODE and gradient flows, our ODE sampler can be seen as a gradient flow that refines a latent z from a random noise to a z vector conditioned on an attribute vector. While [3] relied on the Euler-Maruyama solution of the SDE, we convert our generative SDE to an ODE. Adaptive discretization with a higher order method (e.g., Runge-Kutta [7]) is often preferred when solving ODEs. But our ODE formulation also works well with the first-order Euler method, and its performance lies in-between LACE-LD and LACE-ODE (See Appendix A.11.3).
(ii) Advantages of ODE in the latent space. If the SDE/ODE is built in the pixel space as in [51], it requires 1) estimating the score function rx log pt(x(t)), and 2) training a time-variant classifier pt(c|x(t)), both of which make its training and inference challenging. Instead, our ODE sampler is much simpler: we do not need to train any score function. We only train a time-invariant classifier. Compared with the LD sampler that uses a fixed step size and is sensitive to many hyperparameters (e.g., step size, noise scale and number of steps), our ODE sampler is adaptive in step sizes and only needs to tune the tolerances, making it more efficient and robust to hyperparameters.

heading: 3 Experiments

text:
In this section, we show the effectiveness of our method in conditional sampling, sequential editing and compositional generation, and we also perform an ablation study on the sampling method.
Experimental setting We use StyleGAN-ADA [27] as the pre-trained model for experiments on CIFAR-10 [33], and StyleGAN2 [29] for experiments on FFHQ [28]. We train the classifier in the w-space, where our method works best (see ablation studies in Appendix A.5). To train the latent classifier in the w-space, we first generate (image, w) pairs from StyleGAN, and then label each w latent by annotating its paired image with an image classifier (see data preparation in Appendix A.3).
For LACE-ODE, we use the ‘dopri5’ solver [7] with the tolerances of (1e-3, 1e-3), and we set T = 1, min = 0.1 and max = 20. For LACE-LD, the step size ⌘ and the standard deviation of ✏t are chosen separately for faster training [18], where the number of steps N = 100, step size ⌘ = 0.01 and standard deviation = 0.01. For metrics, we use (i) conditional accuracy (ACC) to measure the controllability, where we generate images using randomly sampled attribute codes, and pass them to a pre-trained image classifier to predict the attribute codes, and (ii) FID to measure the generation quality and diversity [21]. See Appendix A.3 for more details.

heading: 3.1 Conditional sampling

text:
Baselines For comparison, we consider a set of baselines: StyleFlow [1], JEM [18], Conditional EBM (Cond-EBM) [12], and SDEs [51]. We also propose Latent-JEM, a variant of JEM modelled in the latent space, and LACE-PC, which replaces the ODE sampler with the Predictor-Corrector (PC) sampler that solves the reverse SDE (Eq. 6) [51]. See the Appendix A.4 for more details.
CIFAR-10 The results of our method against baselines on CIFAR-10 are shown in Table 1 and Figure 1. Our method requires only four minutes to train, and it takes less than one second to sample a batch of 64 images on a single NVIDIA V100 GPU, which significantly outperforms previous EBMs (at least 49⇥ faster) and score-based models (at least 876⇥ faster) in the pixel space. Latent-JEM that works in the latent space is also slower in training as it performs the LD for each parameter update.
For the conditional sampling performance, LACE-ODE and LACE-LD largely outperform all the baselines in precisely controlling the generation while maintaining the relatively high image quality. In particular, LACE-PC performs similarly to VE-SDE: they can achieve better image quality but have a problem with precisely controlling the generation (ACC0.75). Besides, we observe that EBMs in the latent space achieve much better overall performance than EBMs in the pixel space.
FFHQ To test on FFHQ, we use 10k (w, c) pairs created by [1] for training, where w is sampled from the w-space of StyleGAN2, and c is the attribute code. Unless otherwise specified, we use truncation = 0.7 for our method. Following the evaluation protocol from [1], we use 1k generated samples from StyleGAN2 to compute the FID. Note that given the small sample size, FID values tend to be high. For the reference, the original unconditional StyleGAN2 with 1k samples has FID 20.87.
The results of comparing our method with baselines are shown in Table 2 and Figure 2, where we condition on glasses and gender_smile_age, respectively. For the training time, our method only takes 2 minutes, which is around 25⇥ faster than StyleFlow, and 5⇥ faster than Latent-JEM. Our inference time increases with the number of attributes to control. For instance, LACE-ODE needs similar inference time with StyleFlow (0.68s vs. 0.61s) on glasses, but more inference time than StyleFlow (4.81s vs. 0.61s) on gender_smile_age. In practice, we could adjust the ODE tolerances, allowing for trade-offs between inference time and overall performance. Besides, we can also optimize the network to further reduce the inference time (see results in Appendix A.7).
For the controllability, both LACE-ODE and LACE-LD outperform the baselines by a large margin. Also, the generation quality of our method is on par with our proposed baselines, and much better than the prior work StyleFlow. For instance, LACE-ODE has much lower FID than StyleFlow (24.52 vs. 43.55). Latent-JEM achieves better FID but always has worse controllability. These quantitative results could be verified by the visual samples in Figure 2, where our method achieves high-quality controllable generation. However, StyleFlow has difficulty with (i) the full controllability in different cases, and (ii) the lack of image diversity specifically when conditioning on more attributes.

heading: 3.2 Sequential editing

text:
In sequential editing, we semantically edit the images by changing an attribute each time without affecting other attributes and face identity. Given a sequence of attributes {c1, · · · , cn}, we adapt our method for sequential editing by relying on the compositionality of energy functions. We define the
joint energy function of the i-th edit as
E seq θ (z, {cj} i j=1) := Eθ(z, {cj} i j=1) + µd̃(z, zi 1) +
X
j>i d (fj(g(z); ✓), fj(g(zi 1); ✓))
where the joint energy function Eθ(z, {cj} i j=1) is from Eq. (5), and fj(·; ✓) is the classifier output for the j-th attribute, d̃(z, zi 1) = kg(z) g(zi 1)k 2 2 + kz zi 1)k 2 2 prevents z from moving too far from the previous zi 1, and the last term penalizes z for changing other attributes, with d(·, ·) defined as the squared L2 norm. Note that in the i-th edit, we use zi 1 as the new initial point of the sampling for a faster convergence. By default, we set µ = 0.04 and = 0.01.
We compare our method against StyleFlow [1], the state-of-the-art in sequential editing, with two additional metrics that quantify the disentanglement of the edits: (i) the face identity loss (ID) [1, 42], which calculates the distance between the image embeddings before and after editing to measure the identity preservation. (ii) the disentangled edit strength (DES), defined as DES = 1
n
Pn
i=1 DESi and DESi = Epθ(x)[∆i maxj 6=i |∆j |], where ∆i = ACCi ACC0i 1 ACC0i denotes the normalized ACC improvement, with ACCi and ACC0i being the ACC score of i-th attribute after and before the i-th edit, respectively. Intuitively, the maximum DES is achieved when each edit precisely control the considered attribute only but leave other attributes unchanged, leading to good disentanglement.
Table 3 shows the quantitative results, where the truncation coefficient = 0.5, and we apply the subset selection strategy as proposed in StyleFlow [1] to alleviate the background change and the reweighting of energy functions (see Appendix A.3 for details) for better disentanglement. Our method outperforms StyleFlow in terms of disentanglement (DES), identity preservation (ID), image quality (FID) and controllability (ACC). This is confirmed by the qualitative results in Figure 3, where
StyleFlow usually has the following issues: (i) changing unedited attributes, such as accidentally modifying or losing glasses and changing face identities, and (ii) having incomplete edits, for example, the smiling face still appears after setting smile=0 (last row in Figure 3). On the contrary, our method suffers less from the above problems. Moreover, as the sequential editing in our method is defined by simply composing energy functions, it can incorporate novel attributes in a plug-and-play way.

heading: 3.3 Compositional generation

text:
Zero-shot generation The goal is to generate novel images conditioned on unseen combinations of attributes that are not present in the training data, which is used for evaluating a model’s compositionality in controllable generation. We compare our method against StyleFlow on zero-shot generation in Figure 4, where we condition on (a) {beard=1, smile=0, glasses=1, age=15} and (b) {gender=female, smile=0, glasses=1, age=10}, respectively. We can see our method still performs well in zero-shot generation while StyleFlow suffers from a severe deterioration of image quality and diversity, and almost completely fails in the controllability. Quantitatively, the ACC scores of our method are much larger than StyleFlow, as we have 0.935 vs. 0.679 (smile), and 0.982 vs. 0.694 (glasses) in the setting (a), and 0.906 vs. 0.408 (smile), and 0.939 vs. 0.536 (glasses) in the setting (b). These results clearly demonstrate the strong compositionality of our method.
Compositions of energy functions EBMs have shown great potential in concept compositionality with various ways of composing energy functions [12]. Our method can inherit their compositionality, while achieving high image quality. Inspired by [12], we also consider composing energy functions with three logical operators: conjunction (AND), disjunction (OR) and negation (NOT). In particular,
given two attributes {c1, c2}, we can define the joint energy function of each logical operator as
E(z, {c1 AND c2}) := E(c1|g(z)) + E(c2|g(z)) + 1/2kzk 2 2
E(z, {c1 OR c2}) := log ⇣ eβ−E(c1|g(z)) + e−E(c2|g(z)) ⌘ + 1/2kzk22
E(z, {c1 AND (NOT c2)}) := E(c1|g(z)) αE(c2|g(z)) + 1/2kzk 2 2
where the AND operator actually boils down to the conditional sampling with multiple attributes, and ↵, > 0 are tunable hyperparameters to balance the importance of different energy functions. For more complex compositionality, we can recursively apply these logical operators.
Figure 6 demonstrates the concept compositionality of our method for glasses and yaw where we set = ln 20, and ↵ = min( 0.1|E(z,c2)| , 1). We can see the generated images not only precisely follow the rule of the given logical operators, but are also sufficiently diverse to cover all possible logical cases. For instance, given {glasses=1 OR yaw=front}, some images have glasses regardless of the yaw while other images satisfy yaw=front regardless of the glasses. Our method also works well for the recursive combinations of logical operators, as shown in the bottom-right of Figure 6. To the best of our knowledge, our method is the first to show such strong compositionality when controllably generating photo-realistic images of resolution 1024⇥ 1024.

heading: 3.4 Ablation study on the ODE and LD sampler

text:
Here, we carefully examine the ODE and LD samplers with a grid search in a large range of hyperparameter settings (see Appendix A.11 for details). Figure 5 shows the ACC and FID trade-off of the two samplers on CIFAR-10, where each dot in the figure refers to the result of a particular hyperparameter setting. After grid research, there are 81 and 104 hyperparameter settings for the ODE and LD sampler, respectively.
We can see that (i) the best ACC-FID scores (on the top-left of Figure 5) of the two samplers heavily overlap, implying that they perform equally well in their best tuned hyperparameter settings. (ii) With the concentrated ACC-FID scores for the ODE sampler on the top-left, we can see that the ODE sampler is more stable and less sensitive to the choice of hyperparameters than the LD sampler. (iii) When focusing
on the top-left with ACC 95% and FID  10, the ODE sampler needs much smaller (less than
21%) number of minimum average steps than the LD sampler (41.8 vs. 200). It shows that given the similar performance, the ODE sampler tends to be more efficient due to its adaptivity in step sizes.

heading: 3.5 Results on other high-resolution images

text:
MetFaces [27] By applying the image classifier pre-trained on the FFHQ data to label the painting faces data (i.e., MetFaces), we train a latent classifier in the w-space of StyleGAN-ADA on MetFaces. Accordingly, our method can effectively control the generation of painting faces with one or multiple attributes, as shown in Figure 7(a). It demonstrates the robustness of our method to the classification noise and its generalization ability regarding the domain gap between photos and paintings.
AFHQ-Cats [11] Since the original AFHQ-Cats dataset does not contain the ground-truth attributes (i.e., breed, haircolor, and age, etc.), we have to rely on extra annotators for efficient labeling. Inspired by [15], we apply the CLIP [43] to annotate the AFHQ-Cats data by designing proper prompts that contain the controlling attributes. Similarly, we then train a latent classifier in the w-space of StyleGAN-ADA on AFHQ-Cats to apply our method for controllable generation. As shown in Figure 7(b), we can effectively control the generation of cat images based on the CLIP annotations.

heading: 4 Related work

text:
Conditional generative models Significant efforts for conditional sampling and image editing have been invested in conditional generative models. In these models, various conditional information can be used to guide the generation process, such as class labels [36, 32], attribute codes [9, 38], source images [10], text descriptions [44], and semantic maps [41]. Training these models is costly, in particular for generating high-resolution images. Additionally, adding or modifying the conditional information requires re-training the whole network. On the contrary, our method is based on pretrained unconditional models, and it requires little effort to introduce new conditional information.
Latent code manipulation for image editing An alternative approach to conditional GANs is to manipulate the latent code of pre-trained unconditional models, in particular GANs. Some works explore linear manipulations of latent code [26, 20, 46], while others consider more complicated nonlinear manipulations [37, 16, 1, 42]. Similar to our work, these methods can discover the directions that often correspond to meaningful semantic edits, and need not train conditional generative models from scratch. However, they tend to have issue with compositionality. For instance, StyleFlow [1] specifies a fixed length of attributes during training, limiting its ability to generalize to new attributes and novel combinations (Section 3.3). Notably, [37, 16] share the similar idea of optimizing latent variables of generators with a classifier. But with the intuition of using generative models as a powerful prior, their classifiers are defined in the pixel space only, making their sampling more challenging. Instead, our method is formulated in a principled way from the EBM perspective (that unifies these methods) and can be defined in various spaces (pixel / w / z) of a generator.
EBMs and score-based models EBMs [34] and score-based models [48] have been widely used for image generation [59, 12, 49, 51], due to their flexibility in probabilistic generative modeling. Recently, many works have applied them for controllable and compositional generation. [18] proposes the joint EBM and shows its ability in class-conditional generation. [12] explores the compositionality of conditional EBMs in controllable generation. [51] proposes a unified framework of score-based models from the SDE perspective, and performs controllable generation by using both a conditional reverse-time SDE and its corresponding ODE. Unlike our method that works in the latent space, they all formulate energy functions or score functions in the pixel space, making them both more challenging to train and orders of magnitude slower to sample for high-resolution images.
Also, many other works have built EBMs or score-based models in the latent space of generative models, most of which, however, focus on improving generation quality with EBMs as a structural prior [2, 4, 19, 56, 55, 58, 53]. More similar to our work, [60, 40] considers using the conditional EBMs as a latent prior (conditioning on labels or attributes) of a generator for controllable generation. A key difference is that they require inferring latent variables from data or sampling from the EBMs using LD during training while we do not need to. This advantage distinguishes our method from these prior methods regarding training speed and inference efficiency.

heading: 5 Discussions and limitations

text:
We proposed a novel formulation of a joint EBM in the latent space of pre-trained generative models for controllable generation. Based on our formulation, all we need for controllability is to train an attribute classifier, and sampling is done in the latent space. Moreover, we proposed a more stable and adaptive sampling method by formulating it as solving an ODE. Our experimental results showed that our method is fast to train and efficient to sample, and it outperforms state-of-the-art techniques in both conditional sampling and sequential editing. With our strong performance in the zero-shot generation with unseen attribute combinations and compositions of energy functions with logical operators, our method also demonstrated compositionality in generating high-quality images.
One limitation of this work is that the generation quality is limited by the generative power of the underlying pre-trained generator. Since our EBM can be applied to any latent-variable model including GANs and VAEs, it is interesting to apply our method to different generative models based on downstream tasks. Another limitation is that training the attribute classifier requires the availability of attribute labels. We believe that advanced semi-supervised and self-supervised techniques [47, 8] can improve the classifier training while reducing the dependency on labels.

heading: 6 Broader impact

text:
The method presented in this work enables high-quality controllable image synthesis, built upon pre-trained generative models (e.g., StyleGAN2 [29]). Technically, it inherits the compositionality of EBMs and overcomes their difficulty in generating high-resolution images. Since our method does not train any conditional generator from scratch, it significantly reduces the computational cost of training generators for new conditioning attributes.
On the application side, it shares with other image synthesis tools similar potential benefits and risks, which have been discussed extensively in [5, 52]. Our method does not produce new images but only guides the generation process of existing generators. Thus, it inherits potential biases from the pre-trained generators, e.g. StyleGAN2 trained on the FFHQ dataset [28]. On the other hand, by providing a semantic control of image generation with strong compositionality, our method can be used to discover and proactively reduce unknown biases in existing generators.

title: Controllable and Compositional Generation with Latent-Space Energy-Based Models


summaries: The paper studies a federated setup where each client faces a stochastic contextual bandit.  The parameters of the bandits are coupled across the clientsm. The clients share their local estimates with the server which aggregate them and share back with the clients. The client are hetogenous and the server leverages the geomereic structure of the linear rewars to facilitates the clients learn the optimal arms.  Authors propose a collaborative algorithm called Fed-PE that do not require the nodes to exchange the local feature vectors thus preserving the privacy of the nodes. The authors consider the cases where the paramers are disjoint across the arms or coupled. The performance of the algorithms are valued on both synthetic and real datasets.

summaries: This paper considers a linear contextual bandits problem in a federated setting where different users have different contexts but the rewards are generated through common global parameters. The goal is to minimise the regret for each user while sharing the exploration across users. The paper proposes an algorithm termed Fed-PE with regret bounded as O(\sqrt{dMT}), where d is the dimension of contexts, M is the number of users and T is the time-horizon. A matching lower bound is also presented showing that this regret is tight. The main challenge in this setting is to \emph{balance} the exploration across users as contexts might be very diverse. In order to address this challenge the paper proposes a new G-optimal design which optimises for balanced exploration across users.

summaries: This paper discusses the federated multi-armed bandit problem with linear reward functions. The formulation utilizes a linear structure with agent-dependent contexts to model differences between agent reward functions. The authors propose an algorithm based on a multi-agent G-optimal design with competitive theoretical performance under heterogeneity. The authors additionally establish minimax rates under a reasonable collinearity assumption. The algorithms FedPE and Enhanced FedPE additionally provide competitive experimental performance on both simulated and real-world benchmarks. In summary, the paper is a good contribution, however, there are some issues that I would like to see addressed as well.

summaries: I think the paper sufficiently characterizes the federated linear contextual bandits problem (i) it proposes an algorithm for both the shared parameter and disjoint parameter setting (ii) proves regret guarantees and (iii) characterizes a reasonable class of policies for which a lower bound holds; this lower bound shows that the achieved regret guarantees are near optimal. The communication cost of the algorithm is also reasonable w.r.t to prior work. I think the writing of the paper can be improved in some areas and some technical points need to be explained more, but over all I think this is a good paper.

abstractText: This paper presents a novel federated linear contextual bandits model, where individual clients face different K-armed stochastic bandits coupled through common global parameters. By leveraging the geometric structure of the linear rewards, a collaborative algorithm called Fed-PE is proposed to cope with the heterogeneity across clients without exchanging local feature vectors or raw data. Fed-PE relies on a novel multi-client G-optimal design, and achieves near-optimal regrets for both disjoint and shared parameter cases with logarithmic communication costs. In addition, a new concept called collinearly-dependent policies is introduced, based on which a tight minimax regret lower bound for the disjoint parameter case is derived. Experiments demonstrate the effectiveness of the proposed algorithms on both synthetic and real-world datasets.

authors:
A. Agarwal
J. Langford
Wei
C.-Y.

title: Federated residual learning

authors:
S. Agrawal
N. Goyal

title: Analysis of Thompson sampling for the multi-armed bandit problem

venue: Conference on Learning Theory, pages 39–1.

title: Further optimal regret bounds for Thompson sampling

venue: Artificial Intelligence and Statistics, pages 99–107.

title: Thompson sampling for contextual bandits with linear payoffs

venue: Proceedings of the 30th International Conference on International Conference on Machine Learning, pages 1220–1228.

title: Using confidence bounds for exploitation-exploration trade-offs

venue: J. Mach. Learn. Res., 3:397–422.

authors:
P. Auer
N. Cesa-Bianchi
P. Fischer

title: Finite-time analysis of the multiarmed bandit problem

venue: Machine learning, 47(2-3):235–256.

authors:
P. Auer
R. Ortner

title: UCB revisited: Improved regret bounds for the stochastic multi-armed bandit problem

venue: Periodica Mathematica Hungarica,

authors:
I. Bogunovic
A. Losalka
A. Krause
J. Scarlett

title: Stochastic linear bandits robust to adversarial attacks

venue: International Conference on Artificial Intelligence and Statistics, pages 991–999. PMLR.

authors:
S. Bubeck
N. Cesa-Bianchi

title: Regret analysis of stochastic and nonstochastic multi-armed bandit problems

venue: Foundations and Trends® in Machine Learning, 5(1):1–122.

authors:
W. Chu
L. Li
L. Reyzin
R. Schapire

title: Contextual bandits with linear payoff functions

venue: Gordon, G., Dunson, D., and Dudík, M., editors, Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 208–214, Fort Lauderdale, FL, USA. PMLR.

authors:
M. Dimakopoulou
Z. Zhou
S. Athey
G. Imbens

title: Estimation considerations in contextual bandits

venue: arXiv preprint arXiv:1711.07077.

authors:
A. Dubey
A. Pentland

title: Differentially-private federated linear bandits

venue: Advances in Neural Information Processing Systems, 33.

authors:
F.M. Harper
J.A. Konstan

title: The MovieLens Datasets: History and context

venue: ACM Trans. Interact. Intell. Syst., 5(4).

authors:
R. Nock
A. Özgür
R. Pagh
M. Raykova
H. Qi
D. Ramage
R. Raskar
D. Song
W. Song
S.U. Stich
Z. Sun
A.T. Suresh
F. Tramèr
P. Vepakomma
J. Wang
L. Xiong
Z. Xu
Q. Yang
F.X. Yu
H. Yu
S. Zhao

title: Advances and open problems in federated learning

venue: Foundations and Trends® in Machine Learning, 14(1).

authors:
N. Korda
B. Szorenyi
S. Li

title: Distributed clustering of linear bandits in peer to peer networks

venue: Balcan, M. F. and Weinberger, K. Q., editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1301–1309, New York, New York, USA. PMLR.

authors:
T.L. Lai
H. Robbins

title: Asymptotically efficient adaptive allocation rules

venue: Advances in applied mathematics, 6(1):4–22.

authors:
P. Landgren
V. Srivastava
N.E. Leonard

title: On distributed cooperative decision-making in multiarmed bandits

venue: 2016 European Control Conference (ECC), pages 243–248. IEEE.

title: Social imitation in cooperative multiarmed bandits: partition-based algorithms with strictly local information

venue: 2018 IEEE Conference on Decision and Control (CDC), pages 5239–5244. IEEE.

authors:
J. Langford
T. Zhang

title: The epoch-greedy algorithm for multi-armed bandits with side information

venue: Advances in Neural Information Processing Systems, pages 817–824.

authors:
T. Lattimore
C. Szepesvári

title: Bandit Algorithms

venue: Cambridge University Press.

authors:
L. Li
W. Chu
J. Langford
R.E. Schapire

title: A contextual-bandit approach to personalized news article recommendation

venue: Proceedings of the 19th International Conference on World Wide Web, pages 661–670.

authors:
T. Li
L. Song
C. Fragouli

title: Federated recommendation system via differential privacy

venue: IEEE International Symposium on Information Theory (ISIT), pages 2592–2597.

authors:
D. Martínez-Rubio
V. Kanade
P. Rebeschini

title: Decentralized cooperative stochastic bandits

venue: Advances in Neural Information Processing Systems, pages 4529–4540.

authors:
B. McMahan
E. Moore
D. Ramage
S. Hampson
B.A. y Arcas

title: Communicationefficient learning of deep networks from decentralized data

venue: In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS),

title: Collaborative-filtering

authors:
C. Shi
C. Shen

title: Federated multi-armed bandits

venue: Proceedings of the 35th AAAI Conference on Artificial Intelligence (AAAI).

authors:
C. Shi
C. Shen
J. Yang

title: Federated multi-armed bandits with personalization

venue: Proceedings of the 24rd International Conference on Artificial Intelligence and Statistics (AISTATS).

authors:
B. Szorenyi
R. Busa-Fekete
I. Hegedus
R. Ormandi
M. Jelasity
B. Kegl

title: Gossipbased distributed stochastic bandit algorithms

venue: Proceedings of the 30th International Conference on Machine Learning, pages 19–27.

title: Minimum-Volume Ellipsoids

venue: Society for Industrial and Applied Mathematics, Philadelphia, PA.

authors:
Y. Wang
J. Hu
X. Chen
L. Wang

title: Distributed bandit learning: Near-optimal regret with efficient communication

venue: International Conference on Learning Representations.

authors:
Z. Zhu
J. Zhu
J. Liu
Y. Liu

title: Federated bandit: A gossiping approach

venue: Proceedings of the ACM on Measurement and Analysis of Computing Systems, 5(1):1–29.

heading: 1 Introduction

text:
Federated learning (FL) (McMahan et al., 2017) is an emerging distributed machine learning (ML) paradigm where massive number of clients collaboratively learn a shared prediction model while keeping all the training data on local devices. Compared with standard centralized machine learning, FL has the following characteristics (Kairouz et al., 2021):
• Heterogeneous local datasets. The local datasets, which are often generated at edge devices, are likely drawn from non-independent and identically distributed (non-IID) distributions.
• Communication efficiency. The communication cost scales with the number of clients, which is one of the primary bottlenecks of FL. It is critical to minimize the communication cost while maintaining the learning accuracy.
• Privacy. FL protects local data privacy by only sharing model updates instead of the raw data.
While the main focus of the state-of-the-art FL is on the supervised learning setting, recently, a few researchers begin to extend FL to the multi-armed bandits (MAB) framework (Lai and Robbins, 1985; Auer et al., 2002; Bubeck and Cesa-Bianchi, 2012; Agrawal and Goyal, 2012, 2013a). In the canonical setting of MAB, a player chooses to play one arm from a set of arms at each time slot. An arm, if played, will offer a reward that is drawn from its distribution which is unknown to the player. With all previous observations, the player needs to decide which arm to pull each time in order to maximize the cumulative reward. MAB thus represents an online learning model that naturally captures the intrinsic exploration-exploitation tradeoff in many sequential decision-making problems.
Extending FL to the MAB framework is naturally motivated by a corpus of applications, such as recommender systems, clinical trials, and cognitive radio. In those applications, the sequential decision making involves multiple clients and is distributed by nature. While classical MAB models assume immediate access to the sequentially generated data at the learning agent, under the new
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
realm of FL, local datasets can be stored and analyzed at the clients, thus reducing the communication load and potentially protecting the data privacy.
Despite the potential benefits of FL, the sequential decision making and bandit feedback bring new challenges to the design of FL algorithms in the MAB setting. Different from the supervised learning setting where static datasets are collected beforehand, under the MAB setting, data is generated sequentially as decisions are made, actions are taken, and observations are collected. In order to maximize the cumulative reward and minimize the corresponding learning regret, it thus requires sophisticated coordination of the actions of the clients. The heterogeneous reward distributions across clients make the coordination process even more convoluted and challenging. Besides, the data privacy and communication efficiency requirements result in significant challenges for efficient information exchange and aggregation between local clients and the central server.
In this work, we attempt to address those challenges in a federated linear contextual bandits framework. This particular problem is motivated by the following exemplary applications.
• Personalized content recommendation. For content (arm) recommendation in web-services, user engagement (reward) depends on the profile of a user (context). The central server may deploy a recommender system on each user’ local device (client) in order to personalize recommendations without knowing the personal profile or behavior of the user.
• Personalized online education. In order to maximize students performances (reward) in online learning, the education platform (central server) needs to personalize teaching methods (arms) based on the characteristics of individual students (context). With the online learning software installed at local devices (client), it is desirable to personalize the learning experiences without allowing the platform to access students’ characteristics or scores.
In those examples, the reward of pulling the same arm at different clients follows different distributions dependent on the context as in contextual bandits (Auer, 2003; Langford and Zhang, 2008). We note that conventional contextual bandits is defined with respect to a single player, where the time-varying context can be interpreted as different incoming user profiles. In contrast, we consider a multi-client model, where each client is associated with a fixed user profile. The variation of contexts is captured over clients as opposed to over time. Although the set of clients remains fixed through the learning process, the reward of pulling the same arm still varies across clients. Such a model naturally takes data heterogeneity into consideration. Besides, we adopt a linear reward model, which has been widely studied in contextual bandits (Li et al., 2010; Agrawal and Goyal, 2013b).
Main contributions. Our main contributions are summarized as follows.
First, we propose a new federated linear contextual bandits model that takes the diverse user preferences and data heterogeneity into consideration. Such a model naturally bridges local stochastic bandits with linear contextual bandits, and is well poised to capture the tradeoffs between communication efficiency and learning performances in the federated bandits setting.
Second, we design a novel algorithm named Fed-PE and further develop its variants to solve the federated linear contextual bandits problem. Under Fed-PE, clients only upload their local estimates of the global parameters without sharing their local feature vectors or raw observations. It not only keeps the personal information private, but also reduces the upload cost. We explicitly show that Fed-PE and its variants achieve near-optimal regret performances for both disjoint and shared parameter cases with logarithmic communication costs.
Third, we generalize the G-optimal design from the single-player setting (Lattimore and Szepesvári, 2020) to the multi-client setting. We develop a block coordinate ascent algorithm to solve the generalized G-optimal design efficiently with convergence guarantees. Such a multi-client G-optimal design plays a vital role in Fed-PE, and may find broad applications in related multi-agent setups.
Finally, we introduce a novel concept called collinearly-dependent policy and show that the celebrated LinUCB type of policies (Li et al., 2010), Thompson sampling based policies with Gaussian priors (Agrawal and Goyal, 2013b), and least squared estimation based policies, such as Fed-PE, are all in this category. By utilizing the property of collinearly-dependent policies, we are able to characterize a tight minimax regret lower bound in the disjoint parameter setting. We believe that this concept may be of independent interest for the study of bandits with linear rewards.
√
Notations. Throughout this paper, we use ‖x‖V to denote √ xᵀV x. The range of a matrixA, denoted by range(A), is the subspace spanned by the column vectors of A. We use A† and Det(A) to denote the pseudo-inverse and pseudo-determinant of square matrix A, respectively. The specific definitions can be found in Appendix B of the supplementary material.

heading: 2 Related Works

text:
Collaborative and distributed bandits. Our model is closely related to the collaborative and distributed bandits when action collision is not considered. Landgren et al. (2016, 2018) and Martínez-Rubio et al. (2019) study distributed bandits in which multiple agents face the same MAB instance, and the agents collaboratively share their estimates over a fixed communication graph in order to design consensus-based distributed estimation algorithms to estimate the mean of rewards at each arm. Szorenyi et al. (2013) considers a similar setup where in each round an agent is able to communicate with a few random peers. Korda et al. (2016) considers the case where clients in different unknown clusters face independent bandit problems, and every agent can communicate with only one other agent per round. The communication and coordination among the clients in those works are fundamentally different from our work.
Wang et al. (2020) investigates communication-efficient distributed linear bandits, where the agents can communicate with a server by sending and receiving packets. It proposes two algorithms, namely, DELB and DisLinUCB, for fixed and time-varying action sets, respectively. The fixed action set setting is similar to our setup, except that it assumes that all agents face the same bandits model, which does not take data heterogeneity into consideration.
Federated bandits. A few recent works have touched upon the concept of federated bandits. With heterogeneous reward distributions at local clients, Shi and Shen (2021) and Shi et al. (2021) investigate efficient client-server communication and coordination protocols for federated MAB without and with personalization, respectively. Agarwal et al. (2020) studies regression-based contextual bandits as an example of the federated residual learning framework, where the reward of a client depends on both a global model and a local model. Li et al. (2020) and Zhu et al. (2021) focus on differential privacy based local data privacy protection in federated bandits. While the linear contextual bandit model considered in Dubey and Pentland (2020) is similar to this work, it focuses on federated differential privacy and proposes a LinUCB-based FedUCB algorithm, which incurs a higher regret compared with our result for the shared parameter case. A regret and communication cost comparison between Fed-PE and other baseline algorithms is provided in Table 1.

heading: 3 Problem Formulation

text:
Clients and local bandits model. We consider a federated linear contextual bandits setting where there are M clients pulling the same set of K items (arms) denoted as [K] := {1, 2, . . . ,K}. At each time t, each client i ∈ [M ] pulls an arm ai,t ∈ [K] based on locally available information. The incurred reward yi,t is given by yi,t = ri,ai,t + ηi,t, where ηi,t is a random noise, and ri,ai,t is
1Results adapted from Dubey and Pentland (2020) by letting the privacy budget 1/ go to 0. 2Results adapted from the single-player linear contextual bandits studied in Dimakopoulou et al. (2017) by
assuming instantaneous information exchange and sequential decision-making at the central server.
the unknown expected reward by pulling arm ai,t. We note that without additional assumptions or interaction among the clients, each local model is a standard single-player stochastic MAB, where classic algorithms such as UCB (Auer and Ortner, 2010) and Thompson sampling (Agrawal and Goyal, 2012) are known to achieve order-optimal regret.
Linear reward structure with global parameters. In order to capture the inherent correlation between rewards of pulling the same arm by different clients, we assume ri,a has a linear structure, i.e., ri,a = x ᵀ i,aθa, where xi,a ∈ Rd is the feature vector associated with client i and arm a, and θa ∈ Rd is a fixed but unknown parameter vector for each a ∈ [K]. Here we use xᵀ to denote the transpose of vector x. The same arm a may have different reward distributions for different clients, due to potentially varying xi,a across clients. Such a linear model naturally captures the heterogeneous data distributions at the clients, yet admits possible collaborations among clients due to the common parameters {θa}a∈[K]. When θa varies for different arm a, it is called the disjoint parameter case; when θa is known to be a constant across the arms, it is the shared parameter case. We investigate both cases in Sections 4 and 5, respectively.
Communication model. We assume there exists a central server in the system, and similar to FL, the clients can communicate with the server periodically with zero latency. Specifically, the clients can send “local model updates” to the central server, which then aggregates and broadcasts the updated “global model” to the clients. (We will specify these components later.) Note that just as in FL, communication is one of the major bottlenecks and the algorithm has to be conscious about its usage. Similar to Wang et al. (2020), we define the communication cost of an algorithm as the number of scalars (integers or real numbers) communicated between server and clients. We also make the assumption that clients and server are fully synchronized (McMahan et al., 2017).
Data privacy concerns. Similar to Dubey and Pentland (2020), our contextual bandit problem involves two sets of information that are desirable to be kept private to client i: the feature vectors {xi,a}a∈[K] and the observed rewards {yi,t}t∈[T ]. Different from the differential privacy mechanism adopted in Dubey and Pentland (2020), in this work, we aim to communicate estimated global model parameters {θa}a between the clients and the server. This is consistent with the FL framework, where only model updates are communicated instead of the raw data.
Assumption 1 We make the following assumptions throughout the paper:
1) Bounded parameters: For any i ∈ [M ], a ∈ [K], we have ‖θa‖2 ≤ s, 0 < ` ≤ ‖xi,a‖2 ≤ L. 2) Independent 1-subgaussian noise: ηi,t is a 1-subgaussian noise parameter sampled indepen-
dently at each time for each client with E[ηi,t] = 0, E[exp(ληi,t)] ≤ exp(λ 2
2 ) for any λ > 0.
Assumption 1.1 is a standard assumption in the bandit literature, which ensures that the maximum regret at any step is bounded. We emphasize that our work does not make any assumption on the knowledge of suboptimality gaps, nor do we assume the existence of a unique optimal arm at each client.
Our objective is to minimize the expected cumulative regret among all clients, defined as:
E[R(T )] = E [ M∑ i=1 T∑ t=1 ( xᵀi,a∗i θa∗i − x ᵀ i,ai,t θai,t )] , (1)
where a∗i ∈ [K] is an optimal arm for client i: ∀b 6= a∗i , x ᵀ i,a∗i θa∗i − x ᵀ i,bθb ≥ 0.

heading: 4 Federated Linear Contextual Bandits: Disjoint Parameter Case

heading: 4.1 Challenges

text:
Solving the federated linear contextual bandits model faces several new challenges. The first challenge is due to the constraint that only locally estimated parameters {θa}a∈[K] are uploaded to the central server. While this is not an issue for stochastic MAB where the {θa}a∈[K] are scalars (Shi and Shen, 2021; Shi et al., 2021), this brings significant challenges for the aggregation of the local estimates into a “global model” in our setup. This is because under the linear reward structure, the locally received rewards {yi,t}t∈[T ] only contain the projection of θa along the direction of xi,a, while the
portion of information lying outside range(xi,a) is not captured in {yi,t}t∈[T ]. Thus, by utilizing {yi,t}t∈[T ], the locally estimated θa, denoted as θ̂i,a, cannot provide any information of θa beyond range(xi,a). Since {xi,a}i∈[M ] are different for the same arm a, the locally estimated {θ̂i,a}i∈[M ] are essentially lying in different subspaces. The central server thus needs to take such geometric structure into account when aggregating {θ̂i,a} to construct the global estimate of θa. The geometric structure of the local rewards also brings another challenge for the coordination of actions of local clients. Intuitively, in order to help client i accurately estimate the expected reward by pulling arm a, it suffices to obtain an accurate projection of θa on range(xi,a); any part of θa lying outside this subspace is irrelevant. Therefore, if two clients i and j have xi,a and xj,a orthogonal to each other, exchanging the local estimates θ̂i,a and θ̂j,a does not help the other client improve her own local estimation. On the other hand, if xi,a and xj,a are completely aligned with each other, θ̂i,a and θ̂j,a can be aggregated directly to improve the local estimation accuracy of both. With M possible subspaces spanned by {xi,a}i, it is highly likely that different clients receive different amounts of relevant information (i.e., information lying in range(xi,a)) through information exchange facilitated by the central server. Therefore, in order to reduce the overall regret, it is necessary to coordinate the actions of clients in a sophisticated fashion.
Third, since the exact knowledge of the local feature vectors {xi,a} are kept from the central server, the server may not have an accurate estimation of the uncertainty level of the local estimates at each client, or how much the coordination would help individual clients. This would make efficient and effective coordination even more challenging.
4.2 Federated Phased Elimination (Fed-PE) Algorithm
To address the aforementioned challenges, we propose the Federated Phased Elimination (Fed-PE) algorithm. The Fed-PE algorithm works in phases, where the length of phase p is fp+K. It contains a client side subroutine (Algorithm 1) and a server side subroutine (Algorithm 2). Throughout the paper we use superscript p to indicate phase p barring explicit explanation. We use Api ⊂ [K] to denote the subset of active arms at client i in phase p, Ap := ∪Mi=1A p i , Rpa := {i : a ∈ A p i }, and define T pi,a as the time indices at which client i pulls arm a during the collaborative exploration step in phase p. Then, the algorithm works as follows.
Algorithm 1 Fed-PE : client i Input: T , M , K, α, fp
1: Initialization: Pull each arm a ∈ [K] and receive reward yi,a; θ̂0i,a ← yi,axi,a ‖xi,a‖2 ; Send {θ̂0i,a}a to the server; A0i ← [K]; p← 1. 2: while not reaching the time horizon T do 3: Receive {(θ̂pa, V pa )}a∈Ap−1 from the server. . Arm elimination 4: for a ∈ Ap−1i do
r̂pi,a ← x ᵀ i,aθ̂ p a, u p i,a ← α ‖xi,a‖V pa /`. (2)
5: end for 6: âpi ← arg maxa∈Ap−1i r̂ p i,a, A p i ← { a ∈ Ap−1i | r̂ p i,a + u p i,a ≥ r̂ p i,â p i − up i,â p i } . 7: Send Api to the central server. . Active arm set updating 8: Receive fpi,a for all a ∈ A p i .
9: for a ∈ Api do . Collaborative exploration 10: Pull arm a for fpi,a times and receive rewards {yi,t}t∈T pi,a .
θ̂pi,a ← ( 1
fpi,a ∑ t∈T pi,a yi,t ) xi,a ‖xi,a‖2 . (3)
11: end for 12: Send {θ̂pi,a}a∈Api to the server; Pull â p i until phase length equals f p +K.

heading: 13: p← p+ 1.

text:
14: end while
At the initialization phase, each client i pulls every arm a ∈ [K] once and receives a reward yi,a, based on which it obtains an estimate of the projection of θa. These estimates are sent to the server to construct a preliminary global estimate of θa for each a. The global estimates {θ̂1a}a and the potential matrices {V 1a }a are then broadcast to all clients, after which phase p = 1 begins. Note that after receiving {θ̂0i,a}i,a, the central server will keep a unit vector ēi,a = θ̂ 0 i,a/‖θ̂0i,a‖ for all i ∈ [M ], a ∈ [K]. Since θ̂0i,a is a scaled version of xi,a, ēi,a lies in range(xi,a), and will be utilized to coordinate the arm pulling process (coined as collaborative exploration), as elaborated in Section 4.3.
At the beginning of phase p, after receiving the broadcast {θ̂pa}a and {V pa }a from the server, each client i will utilize the (θ̂pa, V p a ) pair to estimate the expected rewards ri,a and obtain the confidence level according to Eqn. (2) for each a ∈ Ap−1i . Based on the constructed confidence interval, client i then eliminates some arms in Ap−1i and obtains A p i . Next, each client i sends the newly constructed active arm set Api to the server. The server then decides fpi,a, the number of times client i pulling arm a during the collaborative exploration step in phase p for each i ∈ [M ] and a ∈ Api . The specific mechanism to decide f p i,a is elaborated in Section 4.3.
After the collaborative exploration step, client i performs least-square estimation (LSE) for each arm a ∈ Api based on local observations collected in the current phase according to Eqn. (3) and then sends it to the server for global aggregation. Note that although θ̂pi,a lies in range(xi,a), the exact value of xi,a is not revealed to the server, thus preserving the privacy to certain extent.
Algorithm 2 Fed-PE : Central server Input: T , M , K, α, fp
1: Initialization: Receive {θ̂0i,a}i,a; ēi,a ← θ̂0i,a
‖θ̂0i,a‖ for all i ∈ [M ], a ∈ [K]; V 1a ←
(∑ i∈[M ] θ̂0i,a(θ̂ 0 i,a) ᵀ
‖θ̂0i,a‖
)† ,
θ̂1a ← V 1a (∑ i∈[M ] θ̂ 0 i,a ) for all a ∈ [K]; Broadcast {θ̂1a, V 1a }a∈[K]; p← 1.
2: while not reaching the time horizon T do 3: Receive {Api }i∈[M ]; Set A p ← ∪Mi=1Api ; SetR p a ← {i : a ∈ Api }. 4: Solve the multi-client G-optimal design in (6), and obtain solution πp = {πpi,a}i∈[M ],a∈Api . 5: For every client i, send {fpi,a := dπ p i,af
pe}a∈Api . 6: Receive {(a, θ̂pi,a)}a∈Api from each client i. 7: for a ∈ Ap do . Global aggregation
V p+1a ← ∑ i∈Rpa fpi,a θ̂pi,a(θ̂ p i,a) ᵀ ‖θ̂pi,a‖2 † , θ̂p+1a ← V p+1a ∑ i∈Rpa fpi,aθ̂ p i,a  . (4) 8: end for 9: Broadcast {(θ̂p+1a , V p+1a )}a∈Ap to all clients.
10: p← p+ 1. 11: end while

heading: 4.3 Multi-client G-optimal Design

text:
In this subsection, we elaborate the core design of Fed-PE, the collaborative exploration step. There are three main design objectives we aim to achieve: 1) As explained in Section 4.1, one of the main challenges in our federated linear contextual bandits setting is that, each client may benefit differently from the information exchange through the central server. To minimize the overall regret, for each arm a ∈ Ap, it is desirable to ensure that after the global aggregation following the collaboration exploration in phase p, the uncertainty in r̂p+1i,a across the clients is balanced. 2) For each client i, in order to eliminate the sub-optimal arms efficiently, it is also important to guarantee that the uncertainty in r̂p+1i,a across the arms in A p i is balanced. 3) Finally, in order to ensure synchronized model updating, we aim to have each client perform the same number of arm pulling in each phase.
Motivated by those objectives, we propose a multi-client G-optimal design to coordinate the exploration of all clients. Specifically, we define πpi : A p i → [0, 1] as a distribution on the active arm set
Api for each i, and denote πp := (π p 1 , . . . , π p M ) as a vector in R
∑ i∈[M] |A
p i |. Let ei,a := xi,a/‖xi,a‖. We note that ei,a equals to either ēi,a or −ēi,a. Then, we define a feasible set Cp ⊂ R ∑ i∈[M] |A p i | as follows:
Cp = πp ∣∣∣∣∣∣ πpi,a ≥ 0,∀i ∈ [M ], a ∈ A p i ,∑ a∈Api πpi,a = 1,∀i ∈ [M ],
rank({πpi,aei,a}i∈Rpa) = rank({ei,a}i∈Rpa),∀a ∈ Ap  . (5) We can verify that Cp is a convex set. The first two conditions ensure that {πpi,a}a form a valid distribution for each client i. We name the last condition as the “rank-preserving” condition. We note that the subspace spanned by the LHS of the rank-preserving condition is always a subset of that spanned by the RHS. Thus, once the rank is preserved, the subspaces spanned by the LHS and the RHS are the same. Thus, this condition ensures that every dimension of θa lying in range({xi,a}i∈Rpa) will be explored under the collaborative exploration. Any violation of the “rank-preserving” condition will lead to information missing along the unexplored dimensions, which shall be prevented in order to reduce the uncertainty level regarding arm a at every client i ∈ Rpa. Then, we formulate the so called multi-client G-optimal design problem as follows:
minimize G(π) = M∑ i=1 max a∈Api eᵀi,a ( ∑ j∈Rpa πpj,aej,ae ᵀ j,a )† ei,a s.t. πp ∈ Cp. (6)
We note that eᵀi,a (∑ j∈Rpa π p j,aej,ae ᵀ j,a )† ei,a can be interpreted as an approximate measure of the uncertainty level along dimension xi,a if the arms are explored locally according to distributions {πpi }i. Thus, the objective function is an approximate measure of the total uncertainties in the least explored arms at each of the clients. By solving (6), the aforementioned three design objectives can be met. We point out that although the server does not known ei,a, the objective function remains the same when ei,a is replaced by ēi,a. Thus, the server can simply use ēi,a to solve (6).
After solving (6) and obtaining {πpi,a}, the server would set {f p i,a := dπ p i,af pe}a∈Api and send it to client i. Note that after taking the ceiling function, ∑ a∈Api fpi,a may be greater than f p. To ensure synchronized updating, each client i would keep pulling the estimated best arm âpi until the phase length equals fp +K.
We note that the multi-client G-optimal design formulated in (6) is related to the G-optimal design for the single-player linear bandits problem discussed in Lattimore and Szepesvári (2020), and the DELB algorithm for the distributed linear bandits in Wang et al. (2020). However, for such cases, the player(s) faces a single bandit problem, thus the objective is to simply obtain a distribution over a so-called core set of arms in order to minimize the maximum uncertainty across the arms. In contrast, due to the multiple clients involved in the federated bandits setting and the heterogeneous reward distributions, we are essentially solving M coupled G-design problems, one associated with each client. Such coupling effect fundamentally changes the nature of the problem, leading to very different characterization of the problem and numerical approaches.
In Appendix C.1 of the supplementary material, we analyze an equivalent problem of the multi-client G-optimal design. We note that such equivalence essentially generalizes the equivalence between the original G-optimal design and D-optimal design in Lattimore and Szepesvári (2020) to the coupled design case. While the original G-design problem can be approximately solved through the FrankWolfe algorithm under an appropriate initialization (Todd, 2016), solving the multi-client G-optimal design problem is numerically non-trivial. In Appendix C.2, we propose a block coordinate ascent algorithm to solve the equivalent problem of (6) efficiently with guaranteed convergence.
4.4 Theoretical Analysis of Fed-PE
We now characterize the performance of the Fed-PE algorithm.
Theorem 1 Under Assumption 1, with probability at least 1−δ, the cumulative regret under Fed-PE scales in O (√ dKMT (log(K(log T )/δ) + min{d, logM}) ) and the communication cost scales
in O(Md2K log T ).
The complete version of Theorem 1 and its proof can be found in Appendix D in the supplementary material. For the communication cost, at each phase p, client i uploads at most K local estimates with dimension d and downloads at most K global estimates and potential matrices with dimension d and d2, respectively. Thus, the upload cost is O(MdK log T ) and the download cost is O(Md2K log T ).
Remark 1 When we set δ = O( √ dK/MT), the overall regret scales in O( √ dKMT log(MKT )),
and the per-client regret scales in O( √ dKT log(MKT )/M). While the minimax lower bound for standard stochastic MAB scales in Ω( √ KT ), the collaborative learning induced by Fed-PE
leads to √ d/M -fold reduction of the per-client regret. We also note that for single-player linear contextual bandits with disjoint parameters, the best known upper bound scales in Õ( √ dKMT ) (Dimakopoulou et al., 2017) over MT arm pulls, which indicates that the regret of Fed-PE is close to the state-of-the-art centralized algorithms at a communication cost in O(log T ).

heading: 4.5 Enhanced Fed-PE

text: The original Fed-PE algorithm requires exponentially increasing fp in order to achieve the regret upper bound in Theorem 1. This is because in each phase p, we only utilize the rewards collected in phase p− 1 to estimate θa. While this simplifies the analysis, the measurements collected in earlier phases cannot be utilized. In order to overcome this limitation, we propose an Enhanced Fed-PE algorithm by leveraging all historical information. Enhanced Fed-PE achieves different tradeoffs between communication cost and regret performance by adjusting fp. The detailed description and analysis of Enhanced Fed-PE for different selection of fp can be found in Appendix E in the supplementary material.

heading: 4.6 Lower Bound

text:
To derive a tight lower bound, we focus on a set of representative policies defined as follows.
Definition 1 (Collinearly-dependent policy) Two clients i and j are called collinear if there exist an arm a ∈ [K] and a subset S ⊂ [M ] such that the following conditions are satisfied: 1) xi,a /∈ span({xm,a|m ∈ S}); and 2) xi,a ∈ span({xm,a|m ∈ S} ∪ {xj,a}). For any two clients i and j that are not collinear, if the action of client i is independent of the action of j under a policy π, then, the policy is called a collinearly-dependent policy.
We note that the definition of collinearly-dependent policies is actually quite natural. Intuitively, for two clients that are not collinear, their local observations on any arm a cannot be utilized to improve each other’s knowledge of their own local models. As a result, they should not affect each other’s decision-making process. As shown in the supplementary material, we can verify that the most celebrated ridge regression based LinUCB type of policies (Li et al., 2010), Thompson sampling based polices with Gaussian priors (Agrawal and Goyal, 2013b), and least-square estimation based policies, including Fed-PE, all fall in this category.
Theorem 2 For any collinearly-dependent policy, there exists an instance of the federated linear contextual bandits such that the regret is lower bounded as R(T ) = Ω( √ dKMT ).
Remark 2 Theorem 2 essentially shows that, even if raw data transmission and instantaneous communication are allowed and other collinearly-dependent policies are adopted, we cannot improve the order of the regret summarized in Theorem 1 much, i.e., Fed-PE is order-optimal up to √ log(KMT ).
The proof of Theorem 2 relies on the construction of a special instance of the federated linear contextual bandits where the clients can be divided into d groups. Clients in each group face the same K-armed stochastic bandits model locally, while clients from two distinct groups are not collinear. Analyzing the regret bound in each individual group, we can show that it is lower bounded by Ω( √ KMT/d). Then, by utilizing the property of collinearly-dependent policies, we can show that the overall regret is lower bounded by Ω( √ dKMT ) for this scenario. More discussions on the collinearly-dependent policies and the complete proof of Theorem 2 can be found in Appendix F in the supplementary material.

heading: 5 Federated Linear Contextual Bandits: Shared Parameter Case

text:
The Fed-PE algorithm can be slightly modified for the shared parameter case where θa = θ,∀a ∈ [K]. While the client side operation stays the same, the global aggregation step at the server side in (4) can be changed by letting the “potential matrix” V p be ( ∑ a∈[K](V p a ) †)†, and the global estimator θ̂p
be V p (∑ i∈[M ] ∑ a∈Ap−1i fp−1i,a θ̂ p−1 i,a ) . Below, we present the main result for this case and leave the detailed algorithm description and regret analysis in the supplementary material.
Theorem 3 Under Assumption 1, with probability at least 1−δ, the regret of the adapted Fed-PE for the shared parameter case is upper bounded by O (√ dMT (log((log T )/δ) + min{d, logMK}) ) ,
and the communication cost scales in O((KdM + d2M) log T ). Remark 3 By setting δ = O( √
1/MT), we can show that the overall regret scales in O( √ dMT log(MTK)). We note that this bound improves the regret bound for the no differential privacy guarantee case in Dubey and Pentland (2020) by a factor of √
log T , although our settings are slightly different. By assuming all clients face the same linear bandits in this shared parameter setting, the regret in the federated setting over horizon T must be worse than the linear bandits over horizon MT . Since the latter is lowered bounded by Ω( √ dMT ) (Chu et al., 2011), the minimax regret for the shared parameter setting is bounded by Ω( √ dMT ) as well. Thus, the modified Fed-PE is near-optimal for this case.
In terms of the communication cost, the uploading cost stays the same as in the disjoint parameter case, while the broadcast cost is reduced by a factor of K, since only one potential matrix needs to be broadcast for the shared parameter case.

heading: 6 Experiments

text:
Experiment results using both synthetic and real-world datasets are reported in this section to evaluate Fed-PE and the proposed enhancement. Additional experimental details and more experimental results can be found in the supplementary material. We consider four different algorithms, namely, Fed-PE, Enhanced Fed-PE, local UCB without communication, and a modified Fed-PE algorithm with full information exchange after collaborative exploration in each phase (coined as ‘Collaborative’ in Figure 1). For all experiments, we set T = 217, fp = 2p, p ∈ {1, 2, . . . , 16}, and run 10 trials. For Fed-PE and its variants, we choose δ = 0.1. Note that other values of δ may further improve the regret. We evaluate the algorithms on both synthetic and MovieLens-100K datasets.
Synthetic Dataset: We first set M = 100,K = 10, and d = 3. We set {θa} as the canonical basis of R3. The feature vectors xi,a are generated randomly ensuring that the suboptimality reward gaps lie in [0.2, 0.4] and ` = 0.5, L = 1. The per-client cumulative regret as a function of T is plotted in Figure 1(a). We see that Enhanced Fed-PE outperforms Fed-PE while being slightly worse than ‘Collaborative’. This indicates that keeping feature vectors xi,a private to clients does not impact the learning performance significantly. All Fed-PE related algorithms outperform local UCB when T is sufficiently large, demonstrating the effectiveness of communication in improving learning locally. We also set K = 10, d = 4, and vary the number of clients M . The performance of Enhanced Fed-PE is plotted in Figure 1(b). We note that the per-client regret monotonically decreases as M increases, corroborating the theoretical results.
Movielens Dataset: We then use the MovieLens-100K dataset (Harper and Konstan, 2015) to evaluate the performances. Motivated by Bogunovic et al. (2021), we first complete the rating matrix R = [ri,a] ∈ R943×1682 through collaborative filtering (Morabia, 2019), and then use non-negative matrix factorization with 3 latent factors to get R = WH , where W ∈ R943×3, H ∈ R3×1682. Let xi,a be the ith row vector of W . We apply the k-means algorithm to the row vectors of H to produce K = 30 groups (arms), and let θa be the center of the a-th group. Finally, we randomly choose M = 100 users’ feature vectors. We observe that 0.4 ≤ ‖xi,a‖2 ≤ 0.8, and the suboptimality gaps lie in [0.01, 0.8]. The regret performances of the algorithms are plotted in Figure 1(c). The curves show similar characteristics as in Figure 1(a). These results demonstrate the effectiveness of collaborative learning in the federated bandits setting.

heading: 7 Discussion and Conclusion

text:
In this work, we have considered a novel federated linear contextual bandits model, which naturally connects local stochastic MAB models with linear contextual bandits through common global parameters. While each client can only observe a projection of each global parameter in its own subspace, Fed-PE utilizes the geometric structure of the local estimates to reconstruct the global parameters and guide efficient collaborative exploration. Theoretical analysis indicates that Fed-PE achieves near-optimal regret for both disjoint and shared parameter cases with a communication cost in the order of O(log T ).
An interesting open question is whether we can further reduce the communication cost without downgrading the regret performance. In particular, we note the the original single-player G-optimal design allows for a sparse solution whose support is of size d(d+1)/2. Our numerical results indicate that such sparse solutions exist for the multi-client G-optimal design as well. Utilizing the sparsity of the solution may reduce the communication cost significantly. Theoretical characterization of the existence of such sparse solutions is our next step.
Another possible direction to explore is to incorporate the differential privacy mechanism to the Fed-PE framework. Although local feature vectors {xi,a} are kept private under Fed-PE, local estimate θ̂i,a lies in range(xi,a), thus revealing the direction of xi,a to the central server. We aim to add certain perturbation on θ̂i,a in order to obfuscate the direction information without significantly affecting the regret performance.
Acknowledgments and Disclosure of Funding
The work of RH and JY was supported by the US National Science Foundation under Grants CNS-1956276, CNS-2003131, CNS-2114542, and ECCS-2030026. CS acknowledges the funding support by the US National Science Foundation under Grants ECCS-2029978, ECCS-2033671, and CNS-2002902. WW’s work was done before he joined Facebook.

title: Federated Linear Contextual Bandits