aainstitutetext: School of Physics & Astronomy, University of Glasgow, Glasgow
G12 8QQ, United Kingdombbinstitutetext: Theoretical Physics Division, Physical
Research Laboratory, Shree Pannalal Patel Marg, Ahmedabad, 380009, Gujarat,
Indiaccinstitutetext: Discipline of Physics, Indian Institute of Technology,
Palaj, Gandhinagar - 382424, Gujarat, Indiaddinstitutetext: Institute for
Particle Physics Phenomenology, Durham University, Durham DH1 3LE, United
Kingdom

# IRC-safe Graph Autoencoder for unsupervised anomaly detection

Oliver Atkinson, a Akanksha Bhardwaj a Christoph Englert b Partha Konar b,c
Vishal S. Ngairangbam d Michael Spannowsky<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Anomaly detection through employing machine learning techniques has emerged as
a novel powerful tool in the search for new physics beyond the Standard Model.
Historically similar to the development of jet observables, theoretical
consistency has not always assumed a central role in the fast development of
algorithms and neural network architectures. In this work, we construct an
infrared and collinear safe autoencoder based on graph neural networks by
employing energy-weighted message passing. We demonstrate that whilst this
approach has theoretically favourable properties, it also exhibits formidable
sensitivity to non-QCD structures.

##  1 Introduction

New physics searches at the high-energy frontier of the Large Hadron Collider
(LHC) have so far not resulted in any significant deviation of experimental
results from the Standard Model (SM) expectation. However, with a growing
dataset of these high energy measurements, the pressure mounts for
theoretically motivated scenarios of beyond the SM (BSM) physics which have
been devised to tackle known shortcomings of the SM. So what are the ways out
of this juxtaposition of experimental agreement with the SM and its failure to
describe established physics at small and large distances?

On the one hand, there is an increasing emphasis on theoretically as-model-
independent-as-possible approaches based on effective field theory (EFT)
Weinberg:1978kz . EFT navigates QFT correlations away from the SM prediction
in any possible direction given the SM symmetry and particle content, thus
avoiding UV model biases. Alas, such an approach poses its own challenges:
looking for deviations from the SM expectation along these lines involves many
ad-hoc interactions. Concrete models will typically only source a subset of
relevant interactions, e.g. Englert:2019xhk ; DasBakshi:2020ejz ;
Bakshi:2021ofj . There has been great progress to facilitate matching
calculations Carmona:2021xtq , however, depending on the new physics scenario,
this can create a significant overhead that must be included in the parameter
fitting procedure itself Freitas:2016iwx ; Englert:2019rga .

On the other hand, we can look for phenomenological deviations from specific
SM signatures directly in collider results without any new physics bias. Under
the assumption that collider data can be modelled sufficiently adequately, we
can employ the SM expectation to identify regions where measurements do not
follow the SM expectation. This anomaly detection has emerged as a powerful
tool to look for any hidden signature of new physics in the data. Recently, a
range of state-of-the-art methods for anomaly detection Atkinson:2021nlt ;
Blance:2020ktp ; Collins:2021nxn ; Aaboud:2018ufy ; Collins:2018epr ;
Blance:2019ibf ; Hajer:2018kqm ; DeSimone:2018efk ; Araz:2021wqm ;
Nachman:2020lpy ; Hallin:2021wme ; Nachman:2020ccu ; Cheng:2020dal ;
Canelli:2021aps  using deep learning have been designed.

Theoretical consistency when confronting collider data with theoretical
expectations is pivotal. The formulation of infrared and collinear (IRC)-safe
observables is necessary to guarantee the comparability of experimental
measurements and theoretical predictions to all orders in perturbation theory
employing the Kinoshita-Lee-Nauenberg (KLN) theorem Kinoshita:1962ur ;
Lee:1964is  and collinear factorisation Collins:1989gx  of parton
distributions. Any sensitivity enhancement observed by algorithms that
inadvertently employ IRC-unsafe information will be critically assessed in
subsequent studies, with a potentially meaningless theoretical outcome. This
is a tedious task all too familiar from the use of IRC-unsafe jet clustering
algorithms (e.g. iterative cone algorithms) by the CDF and D0 experiments
Kilgore:1996sq , which was later only partially addressed with the midpoint
algorithm during Tevatron Run II, until fully IRC-safe algorithms
Catani:1993hr ; Cacciari:2008gp  were established as the only theoretically
meaningful community consensus.

It seems prudent to avoid mistakes of the past: in this paper, we devise an
IRC-safe Graph Neural Network (GNN) autoencoder algorithm, employing an
Energy-Weighted Message-Passing Network (EMPN) Konar:2021zdg  for unsupervised
anomaly detection. While the IRC-safe loss function is the primary observable
of our autoencoder, we also study the latent space (graph) representation
structure as a motivating tool for new physics discrimination Dillon:2021nxw ;
Atkinson:2021nlt  and highlight the relations to known and more “traditional”
IRC-safe observables. This paper is organised as follows: in Sec. 2, we
outline our EMPN approach and detail our IRC-safe graph construction before we
introduce the IRC-safe autoencoder architecture and simulation framework in
Sec. 3. Sec. 4 is devoted to the discussion of the sensitivity performance of
the autoencoder; we also highlight the correlation of sensitivity with more
traditional jet-based observables. We conclude in Sec. 5.

##  2 A brief outline of Energy-Weighted Message Passing algorithm

This section presents a brief overview of the IRC safe Energy-weighted Message
passing algorithm Konar:2021zdg . It generalises Energy Flow Networks
Komiske:2018cqr ; Dolan:2020qkr , an IRC safe feature extraction on point
clouds, by learning relational information between two elements (nodes) by
constructing a graph out of the point cloud. This procedure is similar to
message-passing networks like the Dynamic Graph Convolutional Neural Network
(DGCNN) wang2019dynamic  that extract local features beyond the global feature
extraction via point-cloud-based architectures such as deep-sets
zaheer2017deep  and PointNet charles2017pointnet ; qi2017pointnet++ . The
algorithm consists of two necessary ingredients:

  * •

an IRC safe prescription for constructing graphs which guarantees that the
graph is invariant under soft and collinear splittings;

  * •

an energy-weighted summed aggregation of messages (and node features after the
final message-passing layer) taking the directional (unit vectors or angles)
inputs $\hat{p}_{i}$ and $\hat{p}_{j}$ of the nodes connected by an edge
$(j,i)$ at the initial layer.

In the following, we discuss these two elements separately.

###  2.1 IRC safe graph construction

Figure 1: Representation of radius graph with $R_{0}$ in the $(\eta,\phi)$
plane undergoing a QCD splitting. The black arrows correspond to the
connections of a radius graph, while the red arrows highlight the 3-nearest
neighbours connections. One can see that the radius neighbourhoods have the
same total energy, which is not the case for those obtained by the nearest
neighbours method, leading to an IRC-unsafe construction.

The inductive biases that a message-passing algorithm imposes on its input
data are highly dependent on the graph structure. For instance, the
neighbourhood sets (the set of particles a node is connected with) determine
the local connectivity of the nodes. Thus, the graph construction algorithm
from a point cloud gives a strong indication that a graph neural network is
the best avenue to pursue. The IRC safety of a message-passing algorithm also
depends on the graph construction, and we highlight such an IRC safe graph
construction algorithm in this section.

Let $\mathcal{S}=\\{p_{1},p_{2},p_{3},.....,p_{N}\\}$ be the set of four-
vectors of the particles within a jet, while
$\mathcal{S}^{\prime}=\\{p_{1},p_{2},....,p_{N+1}\\}$ is the same set in the
presence of an additional splitting. The collinear limit is when the emitted
particles $r$ and $s$ with the angular separation $\Delta_{rs}$ tending to
zero, while the soft limit refers to the case when one of the particle’s
energy tends to zero. These four-vectors can be written as

$p_{i}=(z_{i},\hat{p}_{i})\,,~{}\hbox{with}~{}z_{i}={p^{i}_{T}\over\sum_{j\in\mathcal{S}}p^{j}_{T}}~{}\hbox{and}~{}\hat{p}=(\eta,\phi)\,,$
(1)

for hadron colliders, with the separation in the $\eta-\phi$ plane between two
particles $i$ and $j$, defined as $\Delta
R_{ij}=\sqrt{\Delta\eta_{ij}^{2}+\Delta\phi_{ij}^{2}}$ denoting the quantity
analogous to $\Delta_{ij}$. Since we will be taking directed edges, the
neighbourhood set of a node $i$ will be the set of all nodes with incoming
connections to $i$. For all particles $i$ in $\mathcal{S}$ or
$\mathcal{S}^{\prime}$, a graph construction algorithm will construct
neighbourhood sets $\mathcal{N}[i]$ and $\mathcal{N}^{\prime}[i]$,
respectively. We will use a “closed” neighbourhood with $i\in\mathcal{N}[i]$
instead of an “open” neighbourhood $i\notin\mathcal{N}(i)$, since the second
choice will always be IRC unsafe when the node $i$ splits. To illustrate this,
we show the radius graph with $R_{0}$ in the $(\eta,\phi)$ plane in Fig. 1,
where the node $q$ undergoes a splitting. The black arrows highlight the
connections of the radius graph. Fig. 1 also demonstrates a nearest
neighbourhood connection as an example of an IRC unsafe graph construction.

To formalise the graph construction algorithm in terms of the four-vectors of
the particles, we define a decision function $\mathbf{D}(p_{i},p_{j})$ and a
threshold function $\mathbf{T}(p_{i},p_{j})$, such that any particle $j$ with
four-vector $p_{j}$ will be assigned to the neighbourhood of particle $i$ with
four-vector $p_{i}$ if $\mathbf{D}(p_{i},p_{j})$ is less-than or equal-to
$\mathbf{T}(p_{i},p_{j})$. This can be summarised as

$\mathbf{D}(p_{i},p_{j})\leq\mathbf{T}(p_{i},p_{j})\implies
j\in\mathcal{N}[i]\,.$ (2)

Since we are interested in the soft and collinear limits, constructing an IRC
safe graph requires putting conditions on these functions in the respective
kinematical configurations.

The required condition on these functions for a “parent” splitting $q\to r+s$
when the “daughters” $r,s$ become collinear is

$\begin{split}\mathbf{D}(p_{i},p_{r}+p_{s})\leq\mathbf{T}(p_{i},p_{r}+p_{s})\Leftrightarrow\mathbf{D}(p_{i},p_{r})\leq\mathbf{T}(p_{i},p_{r})\;\land&\;\mathbf{D}(p_{i},p_{s})\leq\mathbf{T}(p_{i},p_{s})\,,\\\
\mathbf{D}(p_{r}+p_{s},p_{i})\leq\mathbf{T}(p_{r}+p_{s},p_{i})\Leftrightarrow\mathbf{D}(p_{r},p_{i})\leq\mathbf{T}(p_{r},p_{i})\;\land&\;\mathbf{D}(p_{s},p_{i})\leq\mathbf{T}(p_{s},p_{i})\,,\end{split}$
(3)

where the second condition arises since the nodes $q,\;r$ or $s$ can also be
the node whose neighbourhood is being determined. The only requirement in the
IR limit for a daughter particle is that all the particles in the set
$\mathcal{N}[i]$ are also present in $\mathcal{N}^{\prime}[i]$, with the only
potential addition of a soft particle. This is guaranteed by the form of Eq,
(2), since both functions depend only on the four-vector of the two nodes of
interest.111This is not the case for popular graph construction algorithms
like $k$-nearest neighbours, for which the decision and threshold has a
complicated dependence on the distance of the primary node $i$ with every
other particle in the graph, and on the number of elements in the
neighbourhood set. The conditions (c.f. Eq. 3) are satisfied in the collinear
limit $\Delta_{rs}\to 0$ if

$\mathbf{D}=\mathbf{D}(\hat{p}_{i},\hat{p}_{j})\,,\qquad\mathbf{T}=\mathbf{T}(\hat{p}_{i},\hat{p}_{j})\,,$
(4)

employing the definitions Eq. (1). Therefore, graphs formed by connecting
particles within a constant radius $R_{0}$ in the $\eta-\phi$ plane are IRC
safe when the decision and threshold functions take the form

$\mathbf{D}=\Delta R_{ij}\,,\qquad\mathbf{T}=R_{0}\,.$ (5)

Note that these choices of $\mathbf{D},\mathbf{T}$ yield closed neighbourhoods
without additional requirements. We will use these graphs in the remainder of
this paper; the neighbourhood of a particle of such a radius graph is shown in
Fig. 1.

###  2.2 Energy-weighted message passing

We detail the IRC safe message passing operation in this section. Before doing
so, we summarise the general definition of message passing operation in the
following steps. The first step, the message-passing stage, involves
calculating the messages for all edges present in the graph. The message
function, parametrised as a multilayer perceptron shared for all edges, takes
the node features of the two nodes connected by an edge and evaluates the
message. Since the message function does not need to be symmetric for the two
node features, a direction convention is necessary for the second phase. In
our convention, the message originates from all nodes in the neighbourhood
$\mathcal{N}[i]$ and flows towards the particle $i$. The second step, the
node-readout stage, updates the node features of each node in the graph as a
permutation-invariant function of all incoming messages.

IRC safety of the updated node features after a message-passing operation is
crucially dependent on the nature of the node readout. A readout based on the
maximum or minimum value of the node features depends on a single node feature
in the neighbourhood, and a soft or collinear splitting of this particular
node would render the updated node feature IRC-unsafe. This is ultimately
related to identifying a specific node in the neighbourhood as special,222This
is also the reason for using closed neighbourhoods $\mathcal{N}[i]$, as an
open neighbourhood $\mathcal{N}(i)$, would give a special status to the node
$i$. which impedes KLN cancellations. A mean readout, on the other hand,
explicitly depends on the cardinality of the neighbourhood sets
$\mathcal{N}[i]$ which is not a well-defined QCD quantity either since there
can be an arbitrary but finite amount of resolvable emissions in the enhanced
collinear or soft regions of phase space. Thus we use a summed readout, which
will inclusively take all the particles in the neighbourhood into account and
will not explicitly depend on their size.

An IRC safe graph construction algorithm ensures two things: the equality of
the sum of energy (transverse energy in the case of hadron colliders) of all
particles in either neighbourhood sets and the presence of both collinear
daughters in $\mathcal{N}^{\prime}[i]$ if the parent is present in
$\mathcal{N}[i]$. Defining a scope-dependent energy weight-factors analogous
to $z_{i}$ as

$\omega_{j}^{(\mathcal{K})}=\frac{p_{T}^{j}}{\sum_{k\in\mathcal{K}}\;p_{T}^{k}}\quad,$

with $\mathcal{K}$ denoting the set of particles in the particular readout
operation, any message passing of the form

$\mathbf{h}^{(l+1)}_{i}=\sum_{i\in\mathcal{N}[i]}\;\omega^{(\mathcal{N}[i])}_{j}\;\;\hat{\Phi}^{(l)}(\mathbf{h}^{(l)}_{i},\mathbf{h}^{(l)}_{j})\,,$
(6)

with $\mathbf{h}^{(0)}_{i}=\hat{p}_{i}$ and $\mathbf{h}^{(l)}_{i}$ denoting
the updated node-features after $l$ message-passing operations satisfies IRC
safety; in the infrared limit, it is straightforward to see that any soft
particle with $z_{r}\to 0\implies\omega_{r}^{(\mathcal{N}[i])}\to 0$ for any
node $i$. The splitting $q\to r+s$ for IRC-safe graphs therefore yields

$\omega_{q}^{(\mathcal{N}[i])}=\omega_{r}^{(\mathcal{N}[i])}+\omega_{s}^{(\mathcal{N}[i])}\,.$
(7)

In the collinear limit with $\hat{p}_{q}=\hat{p}_{r}=\hat{p}_{s}$ we have
$\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{q})=\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{r})=\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{s})$.
Combining this with Eq. (7), we obtain (for $l=0$)

$\omega_{q}^{(\mathcal{N}[i])}\;\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{q})=\omega_{r}^{(\mathcal{N}[i])}\;\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{r})+\omega_{s}^{(\mathcal{N}[i])}\;\hat{\Phi}^{(0)}(\hat{p}_{i},\hat{p}_{s})\,.$

When evaluating Eq. (6) for the neighbourhood of a node $i$, the terms on the
RHS and LHS of this expression are the only ones which will not be common
between $\mathcal{N}[i]$ and $\mathcal{N}^{\prime}[i]$, due to the IRC safe
graph construction. The same expression is followed when $i=q$ on the left,
and $i=r$ or $i=s$ on the right, since
$\\{r,s\\}\subset\mathcal{N}^{\prime}[s]$ and
$\\{r,s\\}\subset\mathcal{N}^{\prime}[r]$, with all three neighbourhoods
(including $\mathcal{N}[q]$) containing the same particles except for $q$,
$r$, and $s$. Therefore, from Eq. (6) we have
$\mathbf{h}^{(1)}_{q}=\mathbf{h}^{(1)}_{r}=\mathbf{h}^{(1)}_{s}$ for collinear
splittings. On the other hand, for a soft daughter, say $r$, we have
$\mathbf{h}^{(1)}_{q}=\mathbf{h}^{(1)}_{s}$, but
$\mathbf{h}^{(1)}_{r}\neq\mathbf{h}^{(1)}_{q}$, with $\mathbf{h}^{(1)}_{r}$
not necessarily zero. The presence of the node features of the daughter
particles, even in the soft or collinear limit, impedes an IRC safe
examination of the full jet graph unless observables are specifically designed
to be insensitive to their presence in the IRC limit. The procedures to take
care of these additional nodes are explained in the following sections, which
are different for supervised and unsupervised methods. Since the above
derivation used the collinearity of $q$, $r$, and $s$, for IRC safe
neighbourhoods, for the same neighbourhoods and any successive application of
an energy-weighted message passing of the form Eq. (6), we have
$\mathbf{h}^{(l)}_{q}=\mathbf{h}^{(l)}_{r}=\mathbf{h}^{(l)}_{s}$ for any $l$.

##  3 IRC-safe graph autoencoder

In a supervised machine learning scenario, the IRC-safe graph readout acting
on the node features of the final message-passing operation gives an IRC-safe
graph representation, and one loses the graph’s structure. The graph
representation, a fixed-length vector obtained after applying a permutation
invariant function on the node features for any variable-length graph, feeds
into the downstream network. Therefore, training a classifier on the loss
function defined with the downstream network’s output proceeds without any
complications from the presence of additional soft or collinear nodes. On the
other hand, a graph autoencoder similar to the one proposed in Ref.
Atkinson:2021jnj  preserves the graph structure until the output. Therefore,
the autoencoder’s output graph will have additional nodes in the soft and
collinear limits in the case of extra emissions. Since the observable of
interest for anomaly detection with an autoencoder is the loss function, we
need to ensure its IRC safety. In this section, we first devise an IRC safe
loss function and give details of the network architecture and training.

###  3.1 An IRC-safe loss function

The definition of the loss function involves input which changes with a soft
or collinear splitting. Therefore, the loss which is normally used as an
observable in anomaly detection, needs to be IRC-safe. A simple IRC-safe loss
function for a jet with constituent set $\mathcal{G}$ is of the form

$\mathcal{L}_{\mathcal{G}}=\sum_{i\in\mathcal{G}}z_{i}~{}d(\hat{p}_{i},\hat{\bar{p}}_{i})\,.$
(8)

The barred quantities are the output of the network, while the unbarred
quantities are the inputs to the network. The function
$d(\hat{p}_{i},\hat{\bar{p}}_{i})\geq d_{0}$ denotes a well-behaved metric
(one-to-one) between the input and the output space, with
$d(\hat{p}_{i},\hat{p}_{i})=d_{0}$. We now show that this is indeed an IRC
safe choice:

Any soft particle $s$, will not contribute to the sum since $z_{s}\to 0$, and
hence it is IR safe. For the splitting $q\to r+s$ we have

$\begin{split}\mathcal{L}_{\mathcal{S}}&=...+z_{q}\;d(\hat{p}_{q},\hat{\bar{p}}_{q})+...\\\
\mathcal{L}_{\mathcal{S^{\prime}}}&=...+z_{r}\;d(\hat{p}_{r},\hat{\bar{p}}_{r})+z_{s}\;d(\hat{p}_{s},\hat{\bar{p}}_{s})+...\quad.\end{split}$

Since, by construction, a GNN’s node output after $L$ total message-passing
operations $\mathbf{h}^{(L)}_{i}=\hat{\bar{p}}_{i}$, is a function of the
input four-vectors $\\{p_{1},p_{2},p_{3},....p_{N}\\}$, in general, they can
have a very complicated dependence on all the input node features. However,
due to the IRC safety of the EMPN, we have

$\mathcal{L}_{\mathcal{S^{\prime}}}-\mathcal{L}_{\mathcal{S}}=z_{r}\;d(\hat{p}_{r},\hat{\bar{p}}_{r})+z_{s}\;d(\hat{p}_{s},\hat{\bar{p}}_{s})-z_{q}\;d(\hat{p}_{q},\hat{\bar{p}}_{q})\,.$
(9)

In the collinear limit with
$\hat{p}_{q}=\hat{p}_{r}=\hat{p}_{s}\implies\hat{\bar{p}}_{q}=\hat{\bar{p}}_{r}=\hat{\bar{p}}_{s}$,
we therefore have (since $z_{q}=z_{r}+z_{s}$),

$z_{q}\;d(\hat{p}_{q},\hat{\bar{p}}_{q})=z_{r}\;d(\hat{p}_{r},\hat{\bar{p}}_{r})+z_{s}\;d(\hat{p}_{s},\hat{\bar{p}}_{s})\implies\mathcal{L}_{\mathcal{S^{\prime}}}-\mathcal{L}_{\mathcal{S}}=0\,,$
(10)

i.e. collinear safety. In the following analysis of the EMPN autoencoder we
will use mean-squared error between the input and output node features for
$d(\hat{p}_{i},\hat{\bar{p}}_{i})$.

###  3.2 Jet graph definition

To demonstrate the performance of the described algorithm, we use the publicly
available top-tagging dataset of Refs. kasieczka_gregor_2019_2603256 ;
Butter:2017cot . The dataset contains a training, validation and testing set
of 600k, 200k, and 200k QCD jets, respectively. The training and validation
are done only with the background QCD samples since the total cross-section of
their production would be orders of magnitude higher than most probable
signals. Although the dataset has the same number of top jets for each of the
three analysis stages, we use the 200k top jets of the test dataset as a
benchmark signal scenario. These jets are simulated using Pythia8
Sjostrand:2014zea ; Sjostrand:2007gs  and passed through Delphes3
deFavereau:2013fsa  for the detector simulations using the default ATLAS
parameter card. Jets are clustered from particle flow (Eflow) constituents
with a distance parameter $\Delta R=0.8$ using the anti-$k_{t}$ algorithm
Cacciari:2008gp . The transverse momentum of the jets is in the range
$p_{T}\in[550,650]$ GeV.

Using the constituents of these jets, we construct the radius graphs which
serve as the input to the IRC safe graph network. To construct the jet radius
graph, we first calculate the inter-particle distance $\Delta R_{ij}$ in the
$(\eta,\phi)$ plane. Next, we define a set of all the particles $i$ as the
neighbourhood set $\mathcal{N}[i]$ such that $\Delta R_{ij}\leq R_{0}$, where
$R_{0}$ is an external tunable parameter. Each node is associated with three
node features

$\mathbf{h}^{0}_{i}=({\Delta\eta_{i},\Delta\phi_{i},\Delta R_{i}})\,,$ (11)

where $\Delta\eta_{i}$, $\Delta\phi_{i}$, $\Delta R_{i}$ are calculated with
respect to the jet axis. For the network analysis, we choose $R_{0}=0.3$.
Since the dependence of the classification power on $R_{0}$ for the supervised
case was found to be mild Konar:2021zdg , with the AUC values changing in the
third decimal value for different values of $R_{0}$ between $0.1$ and $0.5$,
we restrict ourselves to a single value in the intermediate range. The final
node vectors contain information about the $L$-hop neighbourhood with an
effective radius of $R_{0}\times L$. On the other hand, the primary region of
activity for the one-prong QCD jets used to train the network lies in a
relatively small central region of the total jet of radius $\Delta R=0.8$.
Therefore, the features learnt by the autoencoder would be weakly dependent on
$R_{0}$, once the effective radius covers a significant portion of the central
region.

###  3.3 Network architecture and training

Figure 2: A schematic diagram of an IRC safe graph-autoencoder.

Now that we have described the construction of the jet graphs, we discuss the
details of the network architecture and training in this section. Follow from
Fig. 2 where we sketch a schematic diagram of an IRC safe graph-autoencoder.
The encoder consists of three edge convolution operations with output
dimensions of 128, 64 and 2, which is the dimension of the latent
representation. Since we take three-dimensional node features, we restrict
ourselves to a 2-dimensional latent space $(g_{1},g_{2})$ to induce an
information bottleneck.333The effective number of inputs to a message function
could be twice the number of input features–one each for the two nodes
connected by an edge. However, a concrete understanding of the universal
approximation properties of graph neural networks NEURIPS2020_e4acb4c8  is yet
to be achieved, making it difficult to precisely determine the actual input
dimensions when looking at the complete graph neural network. The decoder also
has three edge convolution operations, with the first two dimensions mirroring
the encoder network dimensions (excluding the latent dimension). Finally, the
last edge convolution operation maps the 128-dimensional node vectors at the
penultimate message passing the layer to a three-dimensional space to
reconstruct the input node features.

We take $\hat{\Phi}^{(l)}$ at each message-passing layer to be a multilayer
perceptron (MLP). For an edge convolution operation, we have for two node
features $\mathbf{h}^{(l)}_{i}$ and $\mathbf{h}^{(l)}_{i}$ connected by an
edge in Eq. (6),

$\hat{\Phi}^{(l)}(\mathbf{h}^{(l)}_{i},\mathbf{h}^{(l)}_{j})=\hat{\Phi}^{(l)}\left(\mathbf{h}^{(l)}_{i}\oplus(\mathbf{h}^{(l)}_{j}-\mathbf{h}^{(l)}_{i})\right)\,.$

Therefore the input vector to the MLP has twice the node-feature’s dimensions,
since the direct sum
$\mathbf{h}^{(l)}_{i}\oplus(\mathbf{h}^{(l)}_{j}-\mathbf{h}^{(l)}_{i})$, is a
concatenation of the two vector quantities of equal dimensions. The dimension
of the MLP’s output is the same as the output dimension of the message passing
operations and has a linear activation. We fix the MLP to have two hidden
layers with ReLU activation and the same number of nodes as the output
dimension. The network is implemented using the Pytorch-Geometric
Fey/Lenssen/2019  package. Note that we have not performed any hyperparameter
scan as part of this present, proof-of-concept study. We train the network for
fifty epochs with a learning rate of 0.001 using the Adam kingma2014adam
optimiser. The training and validation losses are compared after each epoch to
ensure that there is no overfitting or a premature termination of training.
The epoch with minimum validation loss is used to infer the anomaly detection
on the test dataset.

##  4 Anomaly detection performance and results

Figure 3: The distribution of the loss function of an IRC safe graph
autoencoder trained only with QCD jets with graph radius $R_{0}$=0.3

We now discuss the performance of the designed IRC safe loss function in
detecting anomalous jets when the network is trained only on the QCD
background. We choose boosted top jets from the aforementioned public dataset
as our benchmark. In Fig. 3 (left), we show the distribution of the loss
function for the QCD and top jets (our inputs are the node features given in
Eq. (11)). As can be seen, the distributions of the loss function values for
the QCD and top jets are significantly different, highlighting the capability
of the architecture to detect anomalous jets in an IRC-safe way. The Receiver-
Operator-Characteristic (ROC) curve and the Area Under the Curve (AUC) of
0.902 shown in Fig. 3 (right) confirm the good separation shown in the loss
distribution, rivalling convolutional autoencoders Farina:2018fyg ;
Heimel:2018mkt ; Roy:2019jae ; Finke:2021sdf  which also have AUCs close to
such values (up to 0.93 Heimel:2018mkt  and 0.91 Finke:2021sdf ) on the same
dataset. Although we did not perform a hyperparameter scan for this study, we
observed a decrease in performance for a one-dimensional latent space.

Top jets possess a different and hard kinematical structure that is typically
not present in QCD jets. The ability to look into the soft and collinear
splittings from the QCD shower evolution in an IRC safe way enables the
network to access such information and the hard radiation pattern in a
theoretically meaningful way. Modifications of the soft and collinear
radiation patterns that are seeded by novel hard scales (see e.g. Refs.
Soper:2011cr ; Englert:2011cg ; Gerwick:2011tm ; Gerwick:2012hq ;
Soper:2014rya ; Prestel:2019neg  for a more traditional jet-based approach to
this) are therefore consistently included in the anomaly detection
performance. Therefore, when such non-QCD structures are present, the anomaly
detection performance is considerably improved.

In light of these results, it is worthwhile to compare our autoencoder results
to phenomenological aspects of QCD in jet substructure analyses. From the
point of view of soft and collinear features, Energy Correlation Functions
(ECF) Larkoski:2013eya  are particularly relevant for such a comparison as we
will motivate below. Furthermore, given that our autoencoder condenses the QCD
information into the latent space in an IRC-safe way, it is interesting to see
how it correlates with ECF observables. To this end, we define

$\mathbf{g}=\sum_{a\in\mathcal{G}}\;z_{a}\;\mathbf{h}_{a}\,,$ (12)

where $\mathbf{h}_{a}$ are the latent node features. Similar to the graph
readout in a classification scenario Konar:2021zdg , this is an IRC safe
representation of the jet. The distribution of the individual components of
the two-dimensional graph representation are shown in Fig. 4. The good
performance of the autoencoder is reflected in the good separation in the
latent space. The two latent space directions are, however, completely anti-
correlated; see Fig. 5 (they are also highly correlated with the loss). Thus,
restricting ourselves to any of these three variables would be sufficient for
the anomaly detection problem studied in this work. The loss would most likely
be a better choice when one focuses on anomaly detection capabilities since it
condenses the information of the two-dimensional latent space into a single
quantity. On the other hand, any latent feature would be more suitable for
applications demanding lower execution times, since in this case only the
encoder needs to be evaluated during inference.

Figure 4: The distribution of each dimension of the two-dimensional latent
spaces obtained after an IRC safe graph readout given in Eq. (12).

Figure 5: The correlation of IRC safe loss (cf. Eq. (8)) and latent dimension
(obtained with Eq. (12)) is shown with the Energy Correlation Functions (13).
One can see a very high correlation of the ECFs with the variables obtained
from the network, hinting at a close connection between them.

Moving on to the relation of the learned information with ECFs, we first
define the ECFs as

$\text{ECF}(N,\beta)=\sum_{i_{1}<i_{2}<....i_{N-1}<i_{N}\in
J}\;\left(\prod_{a=1}^{N}z_{i_{a}}\right)\;\left(\prod_{b=1}^{N-1}\prod_{c=b+1}^{N}\Delta
R_{i_{b}i_{c}}^{\beta}\right)\,.$ (13)

Focussing specifically on the case $N=2$, we obtain

$\text{ECF}(2,\beta)=\sum_{j<i\in I_{J}}\;z_{i}\;z_{j}\;\Delta
R_{ij}^{\beta}=\sum_{i=1}^{|J|}z_{i}\;\sum_{j=i+1}^{|J|}z_{j}\;\Delta
R_{ij}^{\beta}\,,$ (14)

where $I_{J}$ is the index set of the constituent set $J$, and
$z_{i}=p_{T}^{i}/(\sum_{k\in I_{J}}p_{T}^{k})$. We can rewrite the expression
as

$\text{ECF}(2,\beta)=\sum_{i=1}^{|J|}z_{i}\;H_{i}\quad,\quad
H_{i}=\sum_{j=i+1}^{|J|}z_{j}\;\Delta R_{ij}^{\beta}\,.$ (15)

Therefore, the quantity $H_{i}$ can be regarded as a scalar node feature
obtained from the message function $\Delta R_{ij}^{\beta}=\Delta
R_{ij}^{\beta}(\hat{p}_{i},\hat{p}_{j})$, with a weighted (by $z_{j}$) summed
readout, while the sum over $i$ to get the ECF is similar to a graph readout
operation on all the nodes (or constituents) of the jet. Although the graph
structure in the current case is the 2-combinatorial graph, such an analogy
suggests that the features extracted by the EMPN are closely connected to
ECFs.

This expectation is analysed in more detail in Fig. 5, where we show the
correlation of different order ECFs with each dimension $g_{i}$ of the latent
graph readout and the loss function . There is a strong correlation between
the 2-point ECFs and the network outputs, which decreases when considering the
4-point ECFs. This difference illustrates the close relation of the message
passing architecture to the 2-point ECFs. The latent dimensions show a higher
correlation for $\beta=1$ than $\beta=2$, while the opposite holds for the
loss function. This may be due to the ReLU activation, which is essentially a
linear function for all positive arguments, while the loss function’s higher
correlation to the quadratic ECFs may be due to the usage of the mean-squared
error as $d(\hat{p}_{i},\hat{\bar{p}}_{i})$.

##  5 Conclusions

Infrared and collinear safety is not a luxury but an essential requirement to
guarantee the theoretical consistency of particle physics collider data
interpretations. The emerging and fast-developing area of anomaly detection
should therefore incorporate IRC safety when analysing data at the LHC where
QCD activity plays a dominant role. New heavy physics significantly deviates
from QCD phenomenology, predominantly characterised by soft and collinear
emissions. Reflecting the QCD expectation adequately helps isolate anomalies
further; the ability to meaningfully interpolate into the soft and collinear
regime is crucial for extending the reach of such techniques to lower scales.
Despite this, IRC safety has not played an essential role in the
implementation of anomaly detection. In this paper, we have placed IRC safety
at the heart of anomaly detection for the first time by constructing a graph
neural network autoencoder that employs Energy-Weighted Message-Passing, which
gives rise to an IRC-safe architecture Konar:2021zdg .

Graph neural networks are well-adapted approaches for isolating tell-tale
correlations of final states Dreyer:2020brq ; Atkinson:2021jnj  and we find
that our algorithm shows a high anomaly detection capability whilst having
theoretically appealing properties. We have demonstrated this by injecting top
jets as an anomaly and finding excellent discriminating sensitivity. While
this partly results from the direct presence of a novel hard scale in the
jet’s substructure, additional sensitivity is accessed from a different soft
and collinear shower pattern that accompanies the hard scale. To highlight
this relation to well-studied observables in QCD phenomenology, we have shown
a strong relation of the information encoded in our autoencoder’s latent space
with energy correlation functions. This motivates extending anomaly detection
analyses using our framework to new physics scenarios of lighter BSM degrees
of freedom, which we leave for future work.

## Acknowledgements

O.A. is supported by the UK Science and Technology Facilities Council (STFC)
under grant ST/V506692/1. A.B. and C.E. are supported by the STFC under grant
ST/T000945/1. C.E. is also supported by the Leverhulme Trust under Research
Project Grant RPG-2021-031 and the IPPP Associate Scheme. P.K. and V.S.N. are
supported by the Physical Research Laboratory (PRL), Department of Space,
Government of India. M.S. is supported by the STFC under grant ST/P001246/1.
Part of the computational work detailed in this paper was performed using the
HPC resources (Vikram-100 HPC) and TDP project at PRL.

## References

  * (1) S. Weinberg, _Phenomenological Lagrangians_ , _Physica A_ 96 (1979) 327–340.
  * (2) C. Englert, P. Galler, A. Pilkington and M. Spannowsky, _Approaching robust EFT limits for CP-violation in the Higgs sector_ , _Phys. Rev. D_ 99 (2019) 095007, [1901.05982].
  * (3) S. Das Bakshi, J. Chakrabortty, C. Englert, M. Spannowsky and P. Stylianou, _$CP$ violation at ATLAS in effective field theory_, _Phys. Rev. D_ 103 (2021) 055008, [2009.13394].
  * (4) S. D. Bakshi, J. Chakrabortty, C. Englert, M. Spannowsky and P. Stylianou, _Landscaping CP-violating BSM scenarios_ , _Nucl. Phys. B_ 975 (2022) 115676, [2103.15861].
  * (5) A. Carmona, A. Lazopoulos, P. Olgoso and J. Santiago, _Matchmakereft: automated tree-level and one-loop matching_ , 2112.10787.
  * (6) A. Freitas, D. López-Val and T. Plehn, _When matching matters: Loop effects in Higgs effective theory_ , _Phys. Rev. D_ 94 (2016) 095007, [1607.08251].
  * (7) C. Englert, P. Galler and C. D. White, _Effective field theory and scalar extensions of the top quark sector_ , _Phys. Rev. D_ 101 (2020) 035035, [1908.05588].
  * (8) O. Atkinson, A. Bhardwaj, C. Englert, V. S. Ngairangbam and M. Spannowsky, _Anomaly detection with convolutional Graph Neural Networks_ , _JHEP_ 08 (2021) 080, [2105.07988].
  * (9) A. Blance and M. Spannowsky, _Unsupervised event classification with graphs on classical and photonic quantum computers_ , _JHEP_ 21 (2020) 170, [2103.03897].
  * (10) J. H. Collins, P. Martín-Ramiro, B. Nachman and D. Shih, _Comparing Weak- and Unsupervised Methods for Resonant Anomaly Detection_ , 2104.02092.
  * (11) ATLAS collaboration, M. Aaboud et al., _A strategy for a general search for new phenomena using data-derived signal regions and its application within the ATLAS experiment_ , _Eur. Phys. J. C_ 79 (2019) 120, [1807.07447].
  * (12) J. H. Collins, K. Howe and B. Nachman, _Anomaly Detection for Resonant New Physics with Machine Learning_ , _Phys. Rev. Lett._ 121 (2018) 241803, [1805.02664].
  * (13) A. Blance, M. Spannowsky and P. Waite, _Adversarially-trained autoencoders for robust unsupervised new physics searches_ , _JHEP_ 10 (2019) 047, [1905.10384].
  * (14) J. Hajer, Y.-Y. Li, T. Liu and H. Wang, _Novelty Detection Meets Collider Physics_ , _Phys. Rev. D_ 101 (2020) 076015, [1807.10261].
  * (15) A. De Simone and T. Jacques, _Guiding New Physics Searches with Unsupervised Learning_ , _Eur. Phys. J. C_ 79 (2019) 289, [1807.06038].
  * (16) J. Y. Araz and M. Spannowsky, _Combine and Conquer: Event Reconstruction with Bayesian Ensemble Neural Networks_ , _JHEP_ 04 (2021) 296, [2102.01078].
  * (17) B. Nachman and D. Shih, _Anomaly Detection with Density Estimation_ , _Phys. Rev. D_ 101 (2020) 075042, [2001.04990].
  * (18) A. Hallin, J. Isaacson, G. Kasieczka, C. Krause, B. Nachman, T. Quadfasel et al., _Classifying Anomalies THrough Outer Density Estimation (CATHODE)_ , 2109.00546.
  * (19) B. Nachman, _Anomaly Detection for Physics Analysis and Less than Supervised Learning_ , 2010.14554.
  * (20) T. Cheng, J.-F. Arguin, J. Leissner-Martin, J. Pilette and T. Golling, _Variational Autoencoders for Anomalous Jet Tagging_ , 2007.01850.
  * (21) F. Canelli, A. de Cosa, L. L. Pottier, J. Niedziela, K. Pedro and M. Pierini, _Autoencoders for semivisible jet detection_ , _JHEP_ 02 (2022) 074, [2112.02864].
  * (22) T. Kinoshita, _Mass singularities of Feynman amplitudes_ , _J. Math. Phys._ 3 (1962) 650–677.
  * (23) T. D. Lee and M. Nauenberg, _Degenerate Systems and Mass Singularities_ , _Phys. Rev._ 133 (1964) B1549–B1562.
  * (24) J. C. Collins, D. E. Soper and G. F. Sterman, _Factorization of Hard Processes in QCD_ , _Adv. Ser. Direct. High Energy Phys._ 5 (1989) 1–91, [hep-ph/0409313].
  * (25) W. B. Kilgore and W. T. Giele, _Next-to-leading order gluonic three jet production at hadron colliders_ , _Phys. Rev. D_ 55 (1997) 7183–7190, [hep-ph/9610433].
  * (26) S. Catani, Y. L. Dokshitzer, M. H. Seymour and B. R. Webber, _Longitudinally invariant $K_{t}$ clustering algorithms for hadron hadron collisions_, _Nucl. Phys. B_ 406 (1993) 187–224.
  * (27) M. Cacciari, G. P. Salam and G. Soyez, _The anti- $k_{t}$ jet clustering algorithm_, _JHEP_ 04 (2008) 063, [0802.1189].
  * (28) P. Konar, V. S. Ngairangbam and M. Spannowsky, _Energy-weighted message passing: an infra-red and collinear safe graph neural network algorithm_ , _JHEP_ 02 (2022) 060, [2109.14636].
  * (29) B. M. Dillon, T. Plehn, C. Sauer and P. Sorrenson, _Better Latent Spaces for Better Autoencoders_ , _SciPost Phys._ 11 (2021) 061, [2104.08291].
  * (30) P. T. Komiske, E. M. Metodiev and J. Thaler, _Energy Flow Networks: Deep Sets for Particle Jets_ , _JHEP_ 01 (2019) 121, [1810.05165].
  * (31) M. J. Dolan and A. Ore, _Equivariant Energy Flow Networks for Jet Tagging_ , _Phys. Rev. D_ 103 (2021) 074022, [2012.00964].
  * (32) Y. Wang, Y. Sun, Z. Liu, S. E. Sarma, M. M. Bronstein and J. M. Solomon, _Dynamic graph cnn for learning on point clouds_ , _Acm Transactions On Graphics (tog)_ 38 (2019) 1–12, [1801.07829].
  * (33) M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov and A. J. Smola, _Deep sets_ , _Advances in Neural Information Processing Systems_ 30 (2017) , [1703.06114].
  * (34) R. Q. Charles, H. Su, M. Kaichun and L. J. Guibas, _Pointnet: Deep learning on point sets for 3d classification and segmentation_ , in _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 77–85, IEEE Computer Society, 2017.  1612.00593.
  * (35) C. R. Qi, L. Yi, H. Su and L. J. Guibas, _Pointnet++: Deep hierarchical feature learning on point sets in a metric space_ , _Advances in Neural Information Processing Systems_ 30 (2017) , [1706.02413].
  * (36) O. Atkinson, A. Bhardwaj, S. Brown, C. Englert, D. J. Miller and P. Stylianou, _Improved Constraints on Effective Top Quark Interactions using Edge Convolution Networks_ , 2111.01838.
  * (37) G. Kasieczka, T. Plehn, J. Thompson and M. Russel, _Top quark tagging reference dataset_ , Mar., 2019.  10.5281/zenodo.2603256.
  * (38) A. Butter, G. Kasieczka, T. Plehn and M. Russell, _Deep-learned Top Tagging with a Lorentz Layer_ , _SciPost Phys._ 5 (2018) 028, [1707.08966].
  * (39) T. Sjöstrand, S. Ask, J. R. Christiansen, R. Corke, N. Desai, P. Ilten et al., _An introduction to PYTHIA 8.2_ , _Comput. Phys. Commun._ 191 (2015) 159–177, [1410.3012].
  * (40) T. Sjostrand, S. Mrenna and P. Z. Skands, _A Brief Introduction to PYTHIA 8.1_ , _Comput. Phys. Commun._ 178 (2008) 852–867, [0710.3820].
  * (41) DELPHES 3 collaboration, J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaître, A. Mertens et al., _DELPHES 3, A modular framework for fast simulation of a generic collider experiment_ , _JHEP_ 02 (2014) 057, [1307.6346].
  * (42) R. Brüel Gabrielsson, _Universal function approximation on graphs_ , in _Advances in Neural Information Processing Systems_ (H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan and H. Lin, eds.), vol. 33, pp. 19762–19772, Curran Associates, Inc., 2020.
  * (43) M. Fey and J. E. Lenssen, _Fast graph representation learning with PyTorch Geometric_ , in _ICLR Workshop on Representation Learning on Graphs and Manifolds_ , 2019.  1903.02428.
  * (44) D. P. Kingma and J. Ba, _Adam: A method for stochastic optimization_ , 1412.6980.
  * (45) M. Farina, Y. Nakai and D. Shih, _Searching for New Physics with Deep Autoencoders_ , _Phys. Rev. D_ 101 (2020) 075021, [1808.08992].
  * (46) T. Heimel, G. Kasieczka, T. Plehn and J. M. Thompson, _QCD or What?_ , _SciPost Phys._ 6 (2019) 030, [1808.08979].
  * (47) T. S. Roy and A. H. Vijay, _A robust anomaly finder based on autoencoders_ , 1903.02032.
  * (48) T. Finke, M. Krämer, A. Morandini, A. Mück and I. Oleksiyuk, _Autoencoders for unsupervised anomaly detection in high energy physics_ , _JHEP_ 06 (2021) 161, [2104.09051].
  * (49) D. E. Soper and M. Spannowsky, _Finding physics signals with shower deconstruction_ , _Phys. Rev. D_ 84 (2011) 074002, [1102.3480].
  * (50) C. Englert, T. Plehn, P. Schichtel and S. Schumann, _Jets plus Missing Energy with an Autofocus_ , _Phys. Rev. D_ 83 (2011) 095009, [1102.4615].
  * (51) E. Gerwick, T. Plehn and S. Schumann, _Understanding Jet Scaling and Jet Vetos in Higgs Searches_ , _Phys. Rev. Lett._ 108 (2012) 032003, [1108.3335].
  * (52) E. Gerwick, T. Plehn, S. Schumann and P. Schichtel, _Scaling Patterns for QCD Jets_ , _JHEP_ 10 (2012) 162, [1208.3676].
  * (53) D. E. Soper and M. Spannowsky, _Finding physics signals with event deconstruction_ , _Phys. Rev. D_ 89 (2014) 094005, [1402.1189].
  * (54) S. Prestel and M. Spannowsky, _HYTREES: Combining Matrix Elements and Parton Shower for Hypothesis Testing_ , _Eur. Phys. J. C_ 79 (2019) 546, [1901.11035].
  * (55) A. J. Larkoski, G. P. Salam and J. Thaler, _Energy Correlation Functions for Jet Substructure_ , _JHEP_ 06 (2013) 108, [1305.0007].
  * (56) F. A. Dreyer and H. Qu, _Jet tagging in the Lund plane with graph networks_ , _JHEP_ 03 (2021) 052, [2012.08526].


# FedSynth: Gradient Compression via Synthetic Data in Federated Learning

Shengyuan Hu111Work done as an intern at Meta
CMU
<EMAIL_ADDRESS>Jack Goetz
Meta
<EMAIL_ADDRESS>Kshitiz Malik
Meta
<EMAIL_ADDRESS>Hongyuan Zhan
Meta
<EMAIL_ADDRESS>Zhe Liu
Meta
<EMAIL_ADDRESS>Yue Liu
Meta
<EMAIL_ADDRESS>
###### Abstract

Model compression is important in federated learning (FL) with large models to
reduce communication cost. Prior works have been focusing on sparsification
based compression that could desparately affect the global model accuracy. In
this work, we propose a new scheme for upstream communication where instead of
transmitting the model update, each client learns and transmits a light-weight
synthetic dataset such that using it as the training data, the model performs
similarly well on the real training data. The server will recover the local
model update via the synthetic data and apply standard aggregation. We then
provide a new algorithm FedSynth to learn the synthetic data locally.
Empirically, we find our method is comparable/better than random masking
baselines in all three common federated learning benchmark datasets.

##  1 Introduction

Federated Learning(FL) has been widely studied recently to train machine
learning models without directly accessing user’s data. Despite being
successful in achieving high utility performance compared to centralized
training, huge communication costs induced by current FL algorithms like
FedAvg[8] prevents using federated data to train large scale models.
Specifically, communicating the entire model between each client and the
server could drastically slower the training process, and imposes
communication costs on the users. Many prior efforts have focused on
sparsifying the model to reduce communication cost, including masking the
model updates, low precision training with quantization, distillation, etc.
However, sparsification based methods usually suffers from communication cost-
utility trade-off: high compression rate could hurt model quality.

To overcome such limitation, we propose a different way to think about model
compression in FL. Instead of a model update, which is the same size as the
original model, we propose sending a batch of carefully optimized synthetic
data, which is significantly smaller in size. Each client crafts a set of
synthetic data such that the model updated by the synthetic data performs well
on the client’s original training data. In this way, we could use a dataset
that is significantly smaller than the client training set to obtain a similar
model update as if we use the original data. Each client then send the
synthetic data to the server. Upon receiving the synthetic data from each
client, the server will use it to recover the model updated by synthetic data.

Having this intuition in mind, we formally propose a new objective in
federated learning for local clients at each communication round. Using this
formulation, we propose an effective solver that could adapt to a wide family
of existing federated learning algorithm. Specifically, we develop an
algorithm that learns synthetic data for local clients under the FedAvg[8]
framework. We empirically evaluate and compare our method with prior works,
demonstrating advantage of transmitting synthetic data as an effective
compression technique.

##  2 Background and Related Work

#### Compression in FL

Model compression techniques has been widely studied in machine learning
community with a centralized dataset. Some widely used methods include
gradient quantization[1], gradient ternarization[12], using the sign of
gradient[2], pruning based methods like masking[3], etc. In federated
learning, compression could happen at two places: transmitting model updates
from local client to the server (upstream); transmitting updated global model
from the server to local client (downstream). Recent works have proposed using
sparsification techniques for both upstream and downstream communication to
reduce the cost of large scale federated learning [9, 7, 10]. Different from
our work, these works focused on communicating a sparse model between server
and client, where the model performance could significantly degrade given the
same number of training steps. Goetz and Tewari [6] proposed a similar scheme
to transmit synthetic data via upstream communication. However, they propose
optimizing synthetic data to minimize the distance between model updated by
synthetic data and model updated by training data. Our proposed objective
directly optimizes the synthetic data so that the resulting model achieves
good performance on the true training data.

#### Dataset Distillation

A motivation of our work is to learn a small set of synthetic data that could
perform equally well on a given model compared to real data used to train the
model. [11] proposed a dataset distillation algorithm that optimizes synthetic
data $w_{syn}$ such that the model learned using $w_{syn}$ as the training
data approximates the model learned using the true training data. Although
similar to our approach, they only considered distillation from a centralized
dataset at one time while in our case we learn the objective locally for each
client at every communication round. Our proposed method is also more general
in the method to update the model using synthetic data (See Section 3.2)
rather than restricted to SGD.

##  3 Communication via Synthetic Data

###  3.1 Formulation

Traditional Federated Learning(FL) aims at solving the following objective:

$\min_{w}\sum_{k=1}^{K}p_{k}F_{k}(w)$ (1)

where $F_{k}(w)$ is the local objective for client $k$, $p_{k}$ are pre-
defined weights such that $\sum_{k}p_{k}=1$. At each communication round, the
central server selects a subset of clients and send the current model to the
them. Each client then separately optimizes its local objective iteratively
using stochastic gradients. Then the server collects and aggregates the model
updates from every client to obtain the new global model. Note that model
updates have the same size as the actual global model, which means if a large-
scale model is used as the model, the client would need to send a model as
large as the global model. Under our proposed method, instead of sending the
model updates, each client now sends batches of synthetic data generated
locally to the server. The server will then utilize the synthetic data to
recover the local model updates. We formalize the optimization process of
synthetic data as the following.

For client $k$, let $D_{k}^{tr}=(X_{k},Y_{k})$ be the training data and
$w_{k}^{t}$ the local copy of global model at the $t$-th communication round.
In traditional FL, at every communication round, client $k$ tries to solve

$\min_{w}F_{k}(D^{tr}_{k};w)$ (2)

using $w_{k}^{t}$ is as the initialization of $w$. Note that a lot of existing
federated learning methods rely on using iterative gradient methods to solve
for Equation 2. Let’s define the update process for any client $k$ at
communication round $t$ to be $\texttt{ClientUpdate}^{k}(\cdot;w_{k}^{t})$.
Thus, in the traditional FL setting, local optimization process could be
written as $\texttt{ClientUpdate}^{k}(D_{k}^{tr};w_{k}^{t})$. The goal is to
find a set of synthetic data
$D_{k}^{syn}=\\{x_{k}^{i},y_{k}^{i}\\}_{i=1,\cdots,m}$ that is significantly
smaller in memory size than $w$, such that
$\texttt{ClientUpdate}^{k}(D_{k}^{syn};w_{k}^{t})$ is similar to optimizing
$\texttt{ClientUpdate}^{k}(D_{k}^{tr};w_{k}^{t})$. At the end of that
communication round, client $k$ will send $D^{syn}_{k}$ to the server and the
server could recover
$w^{syn}_{k}=\texttt{ClientUpdate}^{k}(D_{k}^{syn};w_{k}^{t})$ and utilizes
$w^{syn}_{k}$ as client $k$’s updated model for aggregation.

Now the problem becomes how can we find $D^{syn}_{k}$ that distills the
knowledge from $D^{tr}_{k}$. The most direct way to do so is to minimize a
distance metric between the model generated from synthetic data and the model
generated from true train data. However, note that our purpose is that using
$D^{syn}_{k}$, we could obtain an updated model $w^{syn}_{k}$ such that
$F_{k}(D^{tr}_{k};w^{syn}_{k})$ is as good as $F_{k}(D^{tr}_{k};w^{tr}_{k})$.
With certain purpose in mind, we propose the following objective:

$\min_{D^{syn}_{k}}F_{k}\left(D^{tr}_{k};\operatorname*{arg\,min}_{w}F_{k}\left(D^{syn}_{k};w\right)\right)$
(3)

Note that when Equation 2 is convex in $w$ and Equation 3 is convex in
$D^{syn}_{k}$, given the same $D^{tr}_{k}$, both equations are essentially
finding the same optimal local model. However, when there doesn’t exist a
closed form solution for
$\operatorname*{arg\,min}_{w}F_{k}\left(D^{syn}_{k};w\right)$, the inner
optimization problem for Equation 3 could not be solved exactly with finite
number of steps at every communication round. To find an approximate solution,
most existing FL methods utilize gradient based methods like SGD. Therefore,
we propose to optimize the following objective in practice instead of Equation
3:

$\min_{D^{syn}_{k}}F_{k}\left(D^{tr}_{k};\texttt{ClientUpdate}_{k}\left(D_{k}^{syn};w_{k}^{t}\right)\right)$
(4)

Without loss of generality, ClientUpdate could be any local optimization
methods including GD, SGD, etc.

###  3.2 Algorithm

We summarize our algorithm for federated learning via synthetic data in
Algorithm 1. Our algorithm is based off of FedAvg[8], a communication
efficient method widely used in federated learning. At each communication
round, instead of performing SGD on the local training data, each selected
client $k$ first initializes a synthetic dataset $D^{syn}_{k}$ (line 5). To
find the best $w^{syn}_{k}$, synthetic updated model generated by
$D^{syn}_{k}$ (line 7) that minimizes the its loss on the original training
data $D^{tr}_{k}$, we propose to apply gradient descent on $D^{syn}_{k}$ for
multiple iterations (line 8). After that, client $k$ would send $D^{syn}_{k}$,
an entity that requires significantly less storage compared to the model
weight, back to the server. To recover client $k$’s learned model, the server
updates $w_{k}^{t}$ with $D^{syn}_{k}$ using the same process client $k$
generated $w^{syn}_{k}$ (line 12). We also provide an example of the
ClientUpdate method: running SGD on $w_{k}^{t}$ using the $D^{syn}_{k}$ (line
16-19). This is consistent to the local update rule in FedAvg, where client
applies SGD to update $w_{k}^{t}$ using its local training data. In order to
fully utilize the advantage of using synthetic data to distill the information
from the original training data, we also propose the following techniques
while learning $D^{syn}_{k}$.

Multiple batches of synthetic data At every communication round, FedAvg allows
a selected client $k$ to split its local training data into multiple batches
and perform minibatch-SGD for multiple epochs. Motivated by this, we allow
client $k$ to create multiple batches of synthetic data. Instead of running
one step gradient descent on the entire synthetic data, client $k$ updates
$w_{k}^{t}$ sequentially using different batches of synthetic data, as
specified in Line 18 of Algorithm 1.

Trainable label In a traditional supervised classification task, the data
usually has fixed label $y$. When using cross entropy as the loss function,
fixed $y$ is encoded as an one-hot vector in $\mathbb{R}^{|C|}$ where $C$ is
the set of all labels. However, this is not necessary for synthetic data. The
purpose of using synthetic data is only to generate a model that performs well
on the real training data. Restricting any synthetic $x_{i}$ to have a fixed
label $y_{i}$ is too stringent and limit the search space for pairs of
$(x_{i},y_{i})$ to learn the information of the original training data. Hence,
we propose randomly initializing $y_{i}\sim Uniform(0,1)^{|C|}$. While
updating synthetic data $(x_{i},y_{i})$, we calculate

$\displaystyle x_{i}$ $\displaystyle\leftarrow
x_{i}-\eta_{x}\nabla_{x_{i}}F_{k}(D^{tr}_{k};w^{syn}_{k})$ (5) $\displaystyle
y_{i}$ $\displaystyle\leftarrow
y_{i}-\eta_{y}\nabla_{y_{i}}F_{k}(D^{tr}_{k};w^{syn}_{k})$ (6)

It is worth noting that under certain scenario, we do not limit $y_{i}$ to be
a vector representing the probability that $x_{i}$ belongs to a certain class.
Each entry for $y_{i}$ could be arbitrary real numbers so that we could search
in the entire $\mathbb{R}^{|C|}$ to find a good local minima for $y_{i}$.

Algorithm 1 FedSynth

1: Input: $T$, $E$, $\eta$, $\eta_{w}$, $w^{0}$,
$\\{D^{tr}_{k}\\}_{k=1,\cdots,K}$

2: for $t=0,\cdots,T-1$ do

3: Server selects a subset of clients $S_{t}$ and broadcasts $w^{t}$ to
$S_{t}$.

4: for all $k\in S_{t}$ in parallel do

5: Client $k$ initializes $w_{k}^{t}=w^{t}$ and $m$ batches of synthetic data
$D^{syn}_{k}=\\{x_{i},y_{i}\\}_{i=1,\cdots,m}$.

6: for $j=0,1,\cdots,E$ do

7: Client $k$ obtains the model updated by $D^{syn}_{k}$
$w^{syn}_{k}=\texttt{ClientUpdate}(D^{syn}_{k};w_{k}^{t})$

8: Client $k$ updates $D^{syn}_{k}$ by  $D^{syn}_{k}\leftarrow
D^{syn}_{k}-\eta\nabla_{D^{syn}_{k}}F_{k}(D^{tr}_{k};w^{syn}_{k})$

9: end for

10: Client $k$ sends $D^{syn}_{k}$ back to the server.

11: end for

12: Server recovers
$\widehat{w}^{syn}_{k}=\texttt{ClientUpdate}(D^{syn}_{k},w_{k}^{t})$ for every
$k$.

13: Server aggregates the weight  $w^{t+1}=w^{t}+\frac{1}{|S_{t}|}\sum_{k\in
S_{t}}(\widehat{w}^{syn}_{k}-w^{t})$

14: end for

15: return $w^{T}$

16: ClientUpdate($\\{x_{i},y_{i}\\}_{i=1,2,\cdots,m};w$)

17: for $j=1,\cdots,m$ do

18: Client performs minibatch-SGD locally  $w\leftarrow
w-\eta_{w}\nabla_{w}F_{k}((x_{i},y_{i});w)$

19: end for

##  4 Experiments

In this section we empirically evaluate our Algorithm 1 on common large scale
federated learning benchmarks. We first demonstrate that given the same
compression rate our method could achieve higher test accuracy then random
masking, a popular compression method used in federated learning. We also show
how number of batches of synthetic data($m$) and the size of each batch could
affect the resulting model’s performance.

###  4.1 Experimental setup details

For fair comparison, all methods are trained for the same amount of
communication rounds for each dataset. For baseline method(Random Masking), we
only apply the compression technique during the upstream communication (i.e.
only compress the model sent from client to the server) in order to be
consistent to our method. For a fixed compression rate, we apply grid search
to tune the hyperparameters($E$, $\eta$, $\eta_{w}$) for our method on the
validation data and report the test accuracy corresponding to the best
validation accuracy. We similarly finetune the hyperparameters for random
masking and baseline FedAvg as well. For all our experiments, we evaluate the
test accuracy and compression rate for a fixed number of communication rounds
$T$. All experiments are performed on common federated learning benchmark
datasets. Data for FEMNIST and Reddit are naturally partitioned among all the
users.

###  4.2 Comparison between FedSynth and Random Masking

We evaluated our method on three commonly used federated benchmark datasets:
FEMNIST, MNIST, and Reddit [4]. The results are shown in Table 1. Under all
three datasets, our method achieves comparable/better performance then
baseline random masking methods. Specifically, under a low compression rate,
there is an advantage of using our method over random masking in all three
datasets. Under Reddit next word prediction task where the utility performance
is extremely sensitive to masking, our method that utilizes synthetic data
without trainable $y$ achieves higher test accuracy than prior works under all
compression rates we experimented. It is also worth noting that our method
with trainable $y$ does not always outperform synthetic data with a fixed
label, as shown in the FEMNIST experiment.

Table 1: Comparing our method with previous compression baselines. The best
performance under each given compression rate is highlighted.

FEMNIST |  | FedAvg | Random Masking | FedSynth(Ours) | FedSynth w/ Trainable $y$ (Ours)
---|---|---|---|---|---
1x |  | 69.29 | 69.29 | 69.29 | 69.29
5.8x |  | - | 68.21 | 68.63 | 46.67
11.6x |  | - | 67.34 | 63.27 | 39.98
MNIST |  | FedAvg | Random Masking | FedSynth(Ours) | FedSynth w/ Trainable $y$ (Ours)
1x |  | 97.74 | 97.74 | 97.74 | 97.74
7.8x |  | - | 97.08 | 95.28 | 97.25
15.6x |  | - | 96.94 | 93.68 | 96.62
Reddit |  | FedAvg | Random Masking | FedSynth(Ours) | FedSynth w/ Trainable $y$ (Ours)
1x |  | 14.19 | 14.19 | 14.19 | 14.19
1.3x |  | - | 8.20 | 8.86 | -
2.6x |  | - | 4.87 | 4.89 | -

###  4.3 Number of synthetic batches vs. batch size

As we mentioned in Section 3, each client could split their synthetic data
into multiple batches. In Table 2, we demonstrate how different number of
synthetic batches and batch size could influence the model performance. Given
the same number of data points, having more small batches significantly
outperforms having few large batches. On one extreme where we treat the entire
synthetic dataset as a large batch, i.e. $w_{k}^{t}$ is only updated once to
get $w^{syn}_{k}$, model trained using our methods is barely useful. However,
on the other extreme where every single piece of data is treated as a separate
batch, our method is able to achieve significantly better performance. We
would also like to highlight that the more synthetic data we use, the better
model performance we could obtain, given the same number of synthetic batches
used for getting $w^{syn}_{k}$.

Table 2: The effect of number of synthetic batches and batch size on the test
accuracy of FedSynth.

FEMNIST |  | FedSynth | FedSynth w/ Trainable $y$
---|---|---|---
Synthetic Batches | Batch size |  |  |
1 | 10 |  | 11.87 | 10.93
5 | 5 |  | 64.21 | 29.43
10 | 5 |  | 67.92 | 46.67
10 | 2 |  | 60.19 | 39.72
20 | 2 |  | 68.63 | 42.32
10 | 1 |  | 56.47 | 26.29
20 | 1 |  | 62.37 | 39.98

##  5 Conclusion and Future works

In this work, we propose a new objective for communication efficient federated
learning along with a practical algorithm to solve it. We showed empirically
that our methods outperforms baseline methods at low compression level in all
three datasets we evaluated. In future works, we aim at making the algorithm
more scalable so that learning of synthetic data would require less
iterations. We also want to look at sending differentially private synthetic
data to protect the local data from potential privacy leakage.

## References

  * Alistarh et al. [2017] D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic.  Qsgd: Communication-efficient sgd via gradient quantization and encoding.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Bernstein et al. [2018] J. Bernstein, Y.-X. Wang, K. Azizzadenesheli, and A. Anandkumar.  signsgd: Compressed optimisation for non-convex problems.  In _International Conference on Machine Learning_. PMLR, 2018.
  * Beznosikov et al. [2020] A. Beznosikov, S. Horváth, P. Richtárik, and M. Safaryan.  On biased compression for distributed learning.  _arXiv preprint arXiv:2002.12410_ , 2020.
  * Caldas et al. [2018] S. Caldas, P. Wu, T. Li, J. Konečnỳ, H. B. McMahan, V. Smith, and A. Talwalkar.  Leaf: A benchmark for federated settings, https://leaf.cmu.edu/.  _arXiv preprint arXiv:1812.01097_ , 2018.
  * Cohen et al. [2017] G. Cohen, S. Afshar, J. Tapson, and A. Van Schaik.  Emnist: Extending mnist to handwritten letters.  In _2017 International Joint Conference on Neural Networks (IJCNN)_ , pages 2921–2926, 2017.
  * Goetz and Tewari [2020] J. Goetz and A. Tewari.  Federated learning via synthetic data.  _arXiv preprint arXiv:2008.04489_ , 2020.
  * Haddadpour et al. [2021] F. Haddadpour, M. M. Kamani, A. Mokhtari, and M. Mahdavi.  Federated learning with compression: Unified analysis and sharp guarantees.  In _International Conference on Artificial Intelligence and Statistics_. PMLR, 2021.
  * McMahan et al. [2017] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial Intelligence and Statistics_ , pages 1273–1282. PMLR, 2017.
  * Reisizadeh et al. [2020] A. Reisizadeh, A. Mokhtari, H. Hassani, A. Jadbabaie, and R. Pedarsani.  Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization.  In _International Conference on Artificial Intelligence and Statistics_. PMLR, 2020.
  * Sattler et al. [2020] F. Sattler, S. Wiedemann, K.-R. Müller, and W. Samek.  Robust and communication-efficient federated learning from non-i.i.d. data.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2020\.
  * Wang et al. [2018] T. Wang, J.-Y. Zhu, A. Torralba, and A. A. Efros.  Dataset distillation.  _arXiv preprint arXiv:1811.10959_ , 2018.
  * Wen et al. [2017] W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li.  Terngrad: Ternary gradients to reduce communication in distributed deep learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Xie et al. [2019] C. Xie, S. Koyejo, and I. Gupta.  Asynchronous federated optimization.  _arXiv preprint arXiv:1903.03934_ , 2019.

##  Appendix A Appendix

###  A.1 Datasets and Models

We summarize the details of the datasets and models we used in our empirical
study in Table 3. Our experiments include both text (Reddit) and image (MNIST
and FEMNIST) datasets with both classification task (MNIST and FEMNIST) and
next-word precition task (Reddit).

Table 3:

Dataset | Number of clients | Model | Task Type
---|---|---|---
FEMNIST [5, 4] | 1000 | 4-layer CNN[13] | 62-class image classification
MNIST | 60 | 4-layer CNN[13] | 10-class image classification
Reddit [4] | 100 | Stacked LSTM | Next word precition