# A New Spin on Particle Flow Networks

Chase Shimmin
Department of Physics
Yale University
New Haven, CT 06511
<EMAIL_ADDRESS>

###### Abstract

Point-based neural network architectures such as Deep Sets and point
convolutions have recently proved to be effective in jet tagging. These models
are adaptable to variable-sized inputs and are more efficient than recurrent
networks. Some models, such as the Energy Flow Network, even have desirable
physics properties built-in, such as IRC-safety. We generalize these pointwise
architectures to more naturally learn distributions on the $\eta-\phi$ plane,
and to incorporate the approximate rotational symmetry about the jet axis. We
show that in this generalized form, the EFN remains IRC-safe, while achieving
performance comparable to non-IRC safe models. In the more general case where
the network operates on constituents with arbitrary non-spatial features, we
also show an increase in performance of existing models, while requiring
significantly fewer parameters.

##  1 Introduction

Hadronic jets are the most abundant products of proton collisions at the LHC.
They appear as a result of hard QCD processes, as well as decay products in
many rare signatures. Moreover, many LHC analyses involve searching for highly
boosted hadronically-decaying objects, in which two or more collimated decay
products deposit energy into an single overlapping region of the detector
[Butterworth:2008iy, Plehn:2009rk, Abdesselam:2010pt, Shimmin:2016vlc]. The
extremely large QCD cross sections often lead to a scenario where the rate of
background jets overwhelm the signal. Therefore, classifying the particle
origin of jets is of special importance.

In general, discrimination of candidates in a particular kinematic phase space
is based on features of the jet substructure. That is, we consider the
detailed patterns of radiation observed via the tens to hundreds of
consitutent measurements that comprise each jet. For decades, there has been
much progress in developing analytical variables to encode distinctive
structural features, such as jet shape, n-subjettiness, energy correlation
functions, etc. [Gallicchio:2012ez, Komiske:2017aww, Thaler:2010tr,
Larkoski:2014gra] Many of these are built with specific physical motivations;
for example, certain variables are designed to distinguish between the number
of sub-jet “prongs”, since it is expected that boosted bosons decaying to a
pair of quarks would tend to have two subjets, while a top quark may have
three.

Some of these observables are designed to also satisfy certain theoretical
guarantees, such as Infrared- and Collinear (IRC) safety [Banfi:2006hf,
Dasgupta:2007wa, Moult:2016cvt]. Such quantities are insensitive to the
effects of soft radiation and colliniar splitting, allowing for accurate
measurements free of systematics arising from nonperturbative modeling. These
developments in jet substructure have also yielded practical benefits for LHC
experimentalists, with the development of improved pileup suppression and jet
grooming algorithms [Soyez:2012hv, Larkoski:2014wba].

In addition to its role in surpressing QCD backgrounds, jet substructure
allows experiments to probe the difference between jets initiated by quarks
vs. gluons. Although formally defining jet flavor at the hadron level is a
hazy concept [Komiske:2018vkc, Banfi:2006hf], it is certainly the case that
individual jets can be described as “quark-like” or “gluon-like” based on
flavor-specific distributions. Probing the differential cross sections of
these jet categories can improve models for parton showers, parton
distribution functions, as well as deriving constraints on QCD couplings.
[Komiske:2018vkc, Gallicchio:2011xq, Dasgupta:2007wa, Gallicchio:2012ez,
Larkoski:2014pca].

More recently, there has been a surge of development in the application of
modern deep learning techniques to the area of jet substructure and tagging
[Larkoski:2017jix, Guest:2018yhq]. This is a problem for which deep learning
(DL) is well suited, as the input space of observations is naturally very high
dimensional, and existing LHC data and simulations provide a very large volume
of data from which to learn.

Many different “off the shelf” models developed for and by the field of
industrial and academic ML research have already adapted to LHC specific data.
These models range from standard dense networks [Luo:2017ncs], to 2D image-
based convolutional networks (CNNs) [Cogan:2014oua, deOliveira:2015xxd], to
recursive networks (RNNs) [Egan:2017ojy, Cheng:2017rdo, Louppe:2017ipp].

One tricket facet in the application of DL to problems in jet substructure is
the variable size of the input space. Typical jets at the LHC are comprised of
tens to hundreds of pseudo-momentum constituents; yet, a network with a fixed
architecture must be able to handle differently-sized jets on equal footing in
order to model the data distribution. In the jet image approach, this
collection of consitituents is flattened into a fixed binning scheme,
obviating the issue at the cost of generating very sparse images with which
standard CNNs may struggle. RNNs, which naturally operate on variable-length
sequences; however, most such models require a “temopral ordering” which acts
as an inductive bias that does not map meaningfully onto the problem at hand.
The remaining approach is to consider the inputs as an unordered set. These
so-called “point cloud” methods involve pooling operations which are invariant
w.r.t. particle ordering. The simplest of these architectures is based on Deep
Sets [2017arXiv170306114Z], where each input point is encoded independently to
a latent representation. Another possibility that has been recently explored
[Shlomi2020GraphNN, Moreno:2019neq, Moreno:2019bmu] is to include pairwise
relational features between inputs, which can be modeled with Graph Neural
Networks.

In this paper, we focus on the former point cloud approach. In particular, we
build upon the work of [Komiske:2018cqr] by generalizing the EnergyFlow and
ParticleFlow network models to incorporate rotational equivariance. This leads
to performance that is competitive with graph network models, while retaining
most of the simplicity of the original formulation. We also demonstrate how
preserving this natural spatial symmetry can lead to more interpretable
results.

##  2 Deep Sets and Flow Nets

We will now analyze the connection between Deep Sets, Particle Flow and Energy
Flow networks.

###  2.1 Deep Sets as Projections

The Deep Sets theorem [2017arXiv170306114Z] establishes that certain classes
of permutation-invariant functions $u:2^{X}\rightarrow\mathcal{Y}$ acting on
sets $S$ of inputs may be universally approximated by a function of the form:

$u(S)\approx
f(\\{s_{1},\dots,s_{k}\\})=F\left(\sum_{i=1}^{|S|}\Phi(s_{i})\right)$ (1)

for suitable functions $\Phi:X\rightarrow\mathbb{R}^{L}$ and
$F:\mathbb{R}^{L}\rightarrow\mathcal{Y}$ are some functions of sufficient
complexity (usually represented by neural networks). The hyperparameter $L$ is
the latent space dimension, which must be of sufficient size but is not
generally known in advance.

Strictly speaking, the universal approximator property is only proved for
functions $u$ on sets of fixed-size, i.e. on $X^{k}$. Moreover, the Deep Sets
theorem relies on results for polynomial approximation, and applies only to
functions $u$ which are continuous with respect to the Euclidean metric on
$X^{k}\cong\mathbb{R}^{k\times d}$. Some practical limitations of this as well
as weaker conditions for universality in more general settings have been
developed in [NEURIPS2018_f5f8590c]. Nonetheless, the architecture of Eq. 1
has yielded competetive results, even when the underlying space $X$ encodes
features such as color, size, energy, etc, and consequently has meaningful
topology at all. In our present study, we will ignore the universal
approximator property.

To facilitate our study of symmetry, let us suppose the feature space $X$ has
measure $\mu$, and the functions $f,g\in H$ are elements of the Hilbert space
$\mathcal{H}=L^{2}(X,\mu)$ of square integrable functions We may write the
inner product:

$\left<f,g\right>=\int_{X}f^{*}gd\mu\,.$ (2)

We also define improper integral of the Dirac delta function on this space in
the usual way111This might require something stricter than a measure space for
$X$, possibly a locally compact group?,

$\int_{X}\delta(x-y)f(x)d\mu(x)=:f(y)\,.$ (3)

Now we propose that rather than considering $S$ to be a set of features $s\in
X$, we instead view it as a generalized function
$\mathcal{S}_{S}:X\rightarrow\mathbb{R}$. We make this association via the
functional defined as $\mathcal{S}_{(\cdot)}:2^{X}\times
X\rightarrow\mathbb{R}$ where

$\mathcal{S}_{S}(x)\mapsto\sum_{i=1}^{|S|}\delta(x-s_{i})\,.$ (4)

This object (also known as an empirical distribution) can be viewed as a
“function” representing the spatial density of the observed elements $s\in S$.

Given these definitions, we can evaluate the projection of $\mathcal{S}_{S}$
onto an arbitrary function $\Phi\in\mathcal{H}$:

$\left<\Phi,\mathcal{S}_{S}\right>=\int_{X}\Phi^{*}(x))\mathcal{S}_{S}(x)d\mu=\sum_{i=1}^{|S|}\Phi^{*}(s_{i})\,.$
(5)

We immediately recognize this result as the inner term of Ref. 1, i.e., we may
write:

$u(S)\approx F\left(\left<\Phi^{*},\mathcal{S}_{S}\right>\right)\,.$ (6)

In other words, the Deep Sets networks can be thought of as a projection of
the input onto a collection of learned feature maps. For later convenience, we
have generalized $\Phi$ to be an arbitrary complex-valued function.

###  2.2 Energy Flow Networks

In Ref. [Komiske:2018cqr], the Particle Flow Network (PFN) is defined to be
exactly as in Eq. 1, with the input space $x=(p^{\mu},q_{0},\dots,q_{n})\in
X$, where $p^{\mu}$ is a particle’s 4-momentum, and $q_{i}$ are arbitrary
additional per-particle features such as charge or species. The Energy Flow
Network (EFN) is defined as a special case, where $\Phi$ is restricted such
that

$\Phi(p^{\mu},q_{0},\dots,q_{n})=z\ \Phi_{efn}(\hat{p}^{\mu})\,,$ (7)

that is, $\Phi$ depends only on the momentum direction $\hat{p}^{\mu}$ and an
energy scale222Both the particle energy $E$ and transverse momentum
$p_{\mathrm{T}}$ are reasonable choices for the scalar energy; we shall assume
$z=p_{\mathrm{T}}$, as it is invariant under translation in $\phi$ and $\eta$.
$z$. $\Phi$ is also held to be linear in $z$; this is the basis for the EFN’s
IRC-safety.

Although arbitrary spatial directions have the topology of a sphere,
$\hat{p}^{\mu}\in S_{2}$, momenta represented in detector coordinates
$\vec{x}=(\eta,\phi)$ locally resemble $\mathbb{R}^{2}$ on scales
$|\delta\vec{x}|\ll\pi$. In order to better utilize this simple topology, we
redefine our feature space to be $W=X\times V$, where $X=\mathbb{R}^{2}$
represents local spatial direction, and $V=\mathbb{R^{+}}$ represents energy.

This suggests defining a new functional, $\mathcal{E}_{(\cdot)}:2^{W}\times
X\rightarrow V$, which maps a jet’s constituents $S=\\{(\vec{x}_{i},z_{i})\\}$
to a function defined by:

$\mathcal{E}_{S}(\vec{x})=\sum_{i=1}^{|S|}z_{i}\delta(\vec{x}-\vec{x}_{i})\,.$
(8)

Since $\vec{x}$ now has a meaningful topology, we can easily interpret this
quantity as the directional distribution of energy for the observed jet $S$,
i.e. this quantity is the jet’s empirical energy flow. We can now replace the
summation term in the EFN with the projection of this energy flow onto
arbitrary functions $\Phi$ in the well-understood space of functions that are
square-integrable on $\mathbb{R}^{2}$:

$u_{efn}(S)=F\left(\sum_{i=1}^{|S|}z_{i}\Phi^{*}(\vec{x}_{i})\right)=F\left(\left<\Phi,\mathcal{E}_{S}\right>\right)\,.$
(9)

We emphasize the point that while the appearance of Eq. 7 suggests that the
EFN is a restriction of the functional form of $\Phi$, an alternative
interpretation is that the EFN is a generalization of a deep sets-like
projection on the topological space $X$. It is a generalization in the sense
that we have promoted the empirical distribution of observed points
$\sum\delta(\vec{x}-\vec{x_{i}})$ to an energy-valued scalar field
$\mathcal{E}_{S}(\vec{x})$ over $X$. It is a scalar field in that the values
of the field at point $\vec{x}$ are invariant under group transformations on
$X$. However, the field may take values in a vector space such as
$\mathbb{R}^{n}$, as we shall see in the next section.

For example, the triplet of RGB color channels in a 2D image is a scalar under
rotations; rotating the coordinate system will reorient the image, but does
not affect the local color of any pixels. There are interesting scenarios in
which the particle cloud have non-scalar features, for example position and
momentum. In this case, a rotation of the position space will also locally
affect the orientation of momentum vectors. The representational
considerations for neural networks operating on such data has been explored
recently [2018arXiv180208219T, pmlr-v119-bogatskiy20a], but in the present
work we focus on the scalar-field case.

###  2.3 Particle Flow Networks

Adopting this viewpoint of the particle point-cloud as a scalar field, we
propose an alternative form for the PFN; this definition again involves
defining a functional $\mathcal{J}_{S}:2^{W}\times X\rightarrow V^{\prime}$

$\mathcal{J}_{S}(\vec{x})\mapsto\sum_{i=1}^{|S|}Z(\vec{q}_{i})\delta(\vec{x}-\vec{x}_{i})\,,$
(10)

where $V$ now represents the space of any non-directional features (energy,
charge, mass, etc.), and $Z:V\rightarrow V^{\prime}$ is an arbitrary
(nonlinear) function mapping from those features to a latent space
$V^{\prime}\subseteq\mathbb{C}^{L}$. Note that the learnable filter(s) $\Phi$
remain an arbitrary function on $X$, as this is the relevant space over which
the inner product is defined.

The map $Z$ allows for the transformation of arbitrary non-spatial features so
that $\mathcal{J}_{S}$ represents a scalar field taking values in the
$L$-dimensional latent space $V^{\prime}$. This map could be a statically
chosen function, or itself a trainable neural network. The case $Z(\vec{q})=z$
corresponds to the EFN architecture. With this representation, our network is
defined as

$u_{pfn}(S)=F\left(\sum_{i=1}^{|S|}\Phi_{pfn}(\vec{x}_{i},\vec{q}_{i})\right)=F\left(\left<\Phi,\mathcal{J}_{S}\right>\right)$
(11)

Ultimately, adopting the scalar field representation of Eq. 10 corresponds to
simply requiring the function
$\Phi_{pfn}(\hat{p}^{\mu},\vec{q})=\Phi(\hat{p}^{\mu})Z(\vec{q})$ be separable
in the direction features. Indeed, preventing the learnable function from
observing the joint distribution of spatial and non-spatial features seems a
suboptimal restriction. As we will see, however, there are a number of
advantages to this form.

##  3 Equivariant Projection Filters

EFNs and PFNs work by learning 2D pattern maps, and calculating latent
features by correlating inputs with these filters. However, due to continuous
factors of variation in the data, it could take infinitely many filters in
order to optimally match even a one-dimensional data manifold. As an example,
consider a jet containing two distinct subjets, arising from a quark pair
within the jet, illustrated in Fig. 1. Due to an (approximate) physical
symmetry about the jet axis, the orientation of these subjets relative to the
$+\eta$ axis is essentially randomly distributed from $[0,2\pi)$. That is, in
most cases we expect jets which differ only by an $SO(2)$ rotation about the
axis to be physically identical, and should have equal density in the data
distribution. Therefore, in order for a PFN to respond identically each one of
these orientations, it must either 1) learn infinitely many filters
corresponding to each rotation, 2) learn limit resolution in such a way that
filters are not sensitive to orientation, or 3) respond in unpredictable ways
when rotations are applied.

Figure 1: Illustration of the inner product between a collection of observed
jet constituents (red dots with brightness indicating $p_{\mathrm{T}}$) and
learned filter maps $\Phi_{1,2}(\eta,\phi)$. The observed jets $J_{1}$ and
$J_{2}$ are completely identical up to a rotation about the jet axis. Yet an
EFN would require two filters $\Phi_{1}(\eta,\phi)$ and $\Phi_{2}(\eta,\phi)$
in order to detect the same feature in both jets.

We can ameliorate this issue by imposing some structure on the space of
functions allowed for $\Phi$. One possibility would be to constrain the form
of $\Phi$ so that it is azimuthally symmetric about the jet axis. In this
case, a given jet would give the same filter response regardless of its
orientation; this is an example of invariance. Although in most applications
we do in fact want a network response which is invariant to such a symmetry,
this would be a severe restriction of the functions that could be learned by
the network. Instead of forcing the network to immediately produce an
invariant response, we can instead preserve the structure of the symmetry by
enforcing equivariance at each layer. A particular function $f:X\rightarrow Y$
is equivariant w.r.t. a group $G$ acting on $X$ if there is some
representation $\alpha$ of $G$ on $Y$ s.t. for all $g\in G$ and $x\in X$

$f(g\cdot x)=\alpha(g)f(x)\,.$ (12)

In other words, equivariance of $f$ implies that the group transformation
$x\rightarrow g\cdot x$ can be effected either before evaluating the function
(by acting $g$ on $\vec{x}$), or afterwards (by acting $\alpha(g)$ on the
result).

To obtain an equivariant filter, we could require $\Phi(r,\theta)$ to be
periodic in $\theta$, e.g. $\Phi(r,\theta)=\rho(r)e^{im\theta}$ for integer
$m$ and arbitrary radial profile $\rho$. It is easy to verify the equivariance
of this construction, as $\Phi(R_{\Delta}\vec{x})=e^{im\Delta}\Phi(\vec{x})$.
However, this is also too restrictive a form to be useful, as salient features
in jet substructure are not necessarily periodic.

Instead, we consider the more general case in which $\Phi$ is a steerable
function. A steerable function is one which may be written as a linear
combination of equivariant basis functions ${\phi_{i}}$, so that

$\Phi(g\cdot\vec{x})=\sum_{i=1}^{n}\alpha_{i}(g)\phi_{i}(\vec{x})=\boldsymbol{\alpha}^{\mathrm{T}}(g)\mathbf{\Phi}(\vec{x})\,.$
(13)

This expression does not exactly conform to the equivariance condition 12.
However we may associate the $n$-dimensional vector space
$\mathcal{F}=\mathrm{span}\\{\phi_{1},\dots,\phi_{n}\\}$ with a subspace of
the original Hilbert space, $\mathrm{F}\subseteq\mathcal{H}$. Then in terms of
the column vector $\mathbf{\Phi}$, we have

$\mathbf{\Phi}(g\cdot\vec{x})=\mathbf{A}(g)\mathbf{\Phi}(\vec{x})\,,$ (14)

where $\mathbf{A}(g)=\mathrm{diag}\\{\alpha_{1}(g),\dots,\alpha_{n}(g)\\}$,
which provides a manifestly equivariant form. This is known as the
interpolation equation, and depending on the group $G$ and the choice of basis
functions $\\{\phi_{i}\\}$, $\mathbf{A}$ need not be diagonal. Hence, this
type of equivariant representation generalizes Eq. 13; the function space,
rather than the function itself must be closed under transformations. Function
spaces whose elements satisfy Eq. 14 are called equivariant function spaces.

Functions in an equivariant space $\mathcal{F}$ may be much more generic and
expressive than any particular equivariant function. Returning to the
application of particle flow networks, we find that in order to learn
sufficiently interesting filters, while preserving equivariance, we must
represent intermediate results in the vector space $\mathcal{F}$. For
convenience, we define the projection of $\mathcal{J}_{S}$ onto
$\mathbf{\Phi}\in\mathcal{F}$ as

$\displaystyle\left<\mathbf{\Phi},\mathcal{J}_{S}\right>:=\begin{pmatrix}\left<\phi_{1},\mathcal{J}_{S}\right>\\\
\vdots\\\
\left<\phi_{n},\mathcal{J}_{S}\right>\end{pmatrix}\in\mathbb{C}^{n}\,.$ (15)

A rotation $R_{\Delta}$ of the jet $\mathcal{J}_{S}$ is equivalent to a change
of coordinates so that
$R_{\Delta}\mathcal{J}_{S}(\vec{x})=\mathcal{J}_{S}(R^{-1}_{\Delta}\vec{x})$.
From the definition333fix this argument of the inner product on $\mathcal{H}$,

$\displaystyle\left<\phi_{i},R_{\Delta}\mathcal{J}_{S}\right>$
$\displaystyle=\int_{X}\phi_{i}(\vec{x})\mathcal{J}_{S}(R_{\Delta}^{-1}\vec{x})$
(16)
$\displaystyle=\int_{X}\phi_{i}(R_{\Delta}\vec{x})\mathcal{J}_{S}(\vec{x})$
(17) $\displaystyle=\left<R^{-1}_{\Delta}\phi_{i},\mathcal{J}_{S}\right>$ (18)

From the interpolation Eq. 14,

$\displaystyle\left<R_{\Delta}^{-1}\phi_{j},\mathcal{J}_{S}\right>=A^{*}_{ij}(\Delta)\left<\phi_{i},\mathcal{J}_{S}\right>\,,$
(19)

and therefore the equivariance of the projection onto $\mathbf{\Phi}$ is
established in terms of the interpolation matrix $\mathbf{A}$:

$\displaystyle\left<\mathbf{\Phi},\mathcal{J}_{S}\right>=\mathbf{A}^{*}(\Delta)\left<\mathbf{\Phi},\mathcal{J}_{S}\right>$
(20)

##  4 SO(2) Equivariance

In the preceeding section, we established a general formulation for designing
projective filters by restricting the learnable filters to an equivariant
function space and utilizing the interpolation matrix. We now consider a
concrete example for the rotation group $SO(2)$. It can be shown
[cheng1998theory] that the most general $SO(2)$-equivariant space of functions
$\phi:\mathbb{R}^{2}\rightarrow\mathbb{C}$ are spanned by functions of the
form $\rho(r)e^{im\theta}$, where $r$ and $\theta$ are polar coordinates on
$\mathbb{R}^{2}$, in our case corresponding to $r=\sqrt{\eta^{2}+\phi^{2}}$
and $\theta=\tan^{-1}(\phi/\eta)$. Therefore the simplest possibility is to
define our filter as
$\mathbf{\Phi}=(\phi_{M},\phi_{M-1},\dots,\phi_{-M})^{\mathrm{T}}$ where

$\phi_{m}(r,\theta)=\rho_{m}(r)e^{im\theta}\,.$ (21)

Here $M$ is a hyperparameter that represents an angular frequency cutoff for
features. Since the functions $\rho_{m}:\mathbb{R}^{\geq
0}\rightarrow\mathbb{C}$ may be of arbitrary form, we can parameterize them as
trainable neural networks. Because of the orthogonality of the complex
exponentials, the interpolation matrix is diagonal and given simply by

$\mathbf{A}(\Delta)=\mathrm{diag}\\{e^{iM\Delta},\dots,e^{-iM\Delta}\\}\,.$
(22)

That is, under rotation $R_{\Delta}$, each of the components of
$\mathbf{\Phi}$ (and hence
$\left\mathbf{missing}{\Phi},\mathcal{J}_{S}\right>$) are simply modulated by
a phase with frequency $m$.

###  4.1 Convolution

Now that we have obtained a trainable equivariant intermediate representation,
$\vec{y}=\left<\mathbf{\Phi},\mathcal{J}_{S}\right>\in\mathbb{C}^{2M+1}$, we
will generally want to proceed to deeper representations by passing to
additional network layers. In order to preserve equivariance, the subsequent
layer $f:\mathbb{C}^{2M+1}\rightarrow X^{\prime}$ will need to satisfy the
equivariance condition

$f(\mathbf{A}(\Delta)\vec{y})=\alpha(\Delta)f(\vec{y})$ (23)

for some representation $\alpha$ of $SO(2)$ on $X^{\prime}$.

The most general group-equivariant operation is that of generalized
convolution (or correlation), which is defined for locally compact groups.
Using the interpolation equation, we may consider the projection of an
arbitrarily rotated jet to be a function $g:SO(2)\rightarrow\mathbb{C}^{m}$:

$g(\Delta)=\left<\mathbf{\Phi},R^{-1}_{\Delta}\mathcal{J}\right>=\mathbf{A}(\Delta)\left<\mathbf{\Phi},\mathcal{J}_{S}\right>\,.$
(24)

The generalized convolution for signals on $SO(2)\cong U(1)$ looks essentially
the same as usual 1D convolution, except that all functions on $SO(2)$ have
cyclic boundary conditions:

$\displaystyle[f\star g]_{U(1)}(\Delta)$
$\displaystyle=\int_{-\pi}^{\pi}d\theta f(\theta)g(\Delta-\theta)$ (25)
$\displaystyle=\left(\int_{-\pi}^{\pi}d\theta
f(\theta)\mathbf{A}(\Delta-\phi)\right)\left<\mathbf{\Phi},\mathcal{J}_{S}\right>\,.$
(26)

We have already obtained the analytic form of $\mathbf{A}$; however, depending
on the convolution kernel $f$, we may or may not be able to effect this
integral. In practice, convolution networks use a discretized approximation of
the convolution, so that we may simply sample $f$ and $\mathbf{A}$ at
intervals:

$[f\star g]_{U(1)}(\Delta_{j})=\sum_{i}f_{i}\mathbf{A}(\Delta_{j}-\theta_{i})$
(27)

where $f_{i}=f(\theta_{i})$ are the learnable parameters of the covolutional
filter kernel.

editorial? should i drop the convolutional story or should I add in the
results for the convolutional networks?

###  4.2 Invariance

General equivariant representations will be more powerful than invariant ones.
Once an invariant representation is formed, all subsequent representations
computed by the network must necessarily be invariant as well. But since
rotational invariance is desirable for most jet tagging objectives, an obvious
(and simple) choice is to compute an invariant at some point before the output
of the network.

In particular, the unitarity of $\mathbf{A}(\Delta)$ in Eq. 22 ensures that
the squared modulus
$|\left<\mathbf{\Phi},\mathcal{J}_{S}\right>|^{2}:=\sum_{m}|\left<\phi_{m}|\mathcal{J}_{S}\right>|^{2}$
is invariant under rotation. It also has the benefit of being real-valued,
making it easier to adapt to conventional neural network layers.

Using Parseval’s theorem, it is possible to rewrite this as:

$\sum_{m}|\left<\phi_{m},\mathcal{J}_{S}\right>|^{2}=\frac{1}{2\pi}\int_{-\pi}^{\pi}|\left<\Phi,R_{\Delta}^{-1}\mathcal{J}_{S}\right>|^{2}\,,$
(28)

Therefore we can interpret this invariant as the continuous average projection
of $\mathcal{J}_{S}$ onto the filter, averaged over all possible rotations. In
our experiments, we found that simply computing this norm was nearly as
effective as including additional convolutional layers.

$\left<\mathbf{\Phi},\mathcal{J}_{S}\right>\overset{\vec{x}\rightarrow
g\cdot\vec{x}}{\longrightarrow}\left<\mathbf{\Phi},g^{-1}\mathcal{J}_{S}\right>=\mathbf{A}(g)\left<\mathbf{\Phi},\mathcal{J}_{S}\right>\,.$
(29)

Bra-ket notation lends itself well to this dualistic interpretation;
$\left|\Phi\right>=\sum_{i}a_{i}\left|\phi_{i}\right>$ represents a column-
vector in any basis $\\{\left|\phi_{i}\right>\\}$ of $\mathcal{F}$, while we
understand
$\left<\vec{x}|\Phi\right>=\sum_{i}\left<\vec{x}|\phi_{i}\right>=\sum_{i}\phi_{i}(\vec{x})$.
Then in terms of elements of the function space $\mathcal{F}$, we recover the
equivariance condition

$\displaystyle\Phi(g\cdot\vec{x})$
$\displaystyle=\sum_{i}\alpha_{i}(g)\left<\vec{x}|\phi_{i}\right>=\left<\vec{x}|\mathbf{A}(g)|\Phi\right>$
(30) $\displaystyle\implies$
$\displaystyle\left|g^{-1}\Phi\right>=\mathbf{A}(g)\left|\Phi\right>$ (31)

where $\mathbf{A}(g)=\mathrm{diag}\\{\alpha_{1}(g),\dots,\alpha_{n}(g)\\}$.
This is known as the interpolation equation, and depending on the group $G$
and the choice of functions $\\{\phi_{i}\\}$, $\mathbf{A}$ does not
necessarily need to be diagonal. Hence, this type of equivariant
representation generalizes Eq. 13, and function spaces whose elements satisfy
Eq. 14 are called equivariant function spaces.

Returning to the application of flow projections, we suppose that $\Phi$ is
restricted to some equivariant function space
$\Phi\in\mathcal{F}=\mathrm{span}\\{\phi_{1},\dots,\phi_{n}\\}$. Using the
completeness of $\\{\phi_{i}\\}$ on $\mathcal{F}$, we may then write

$\displaystyle\left<\Phi|\mathcal{J}_{S}\right>$
$\displaystyle=\sum_{i}\left<\Phi|\phi_{i}\right>\left<\phi_{i}|\mathcal{J}_{S}\right>$
(32)
$\displaystyle=\sum_{i}\left<\Phi|\phi_{i}\right>\sum_{(\vec{q},\vec{x})\in
S}Z(\vec{q}_{j})\phi_{i}^{*}(\vec{x}_{j})$ (33)
$\displaystyle=\sum_{i}c_{S}\left<\Phi|\phi_{i}\right>\,,$ (34)

where $c_{S}\in\mathbb{C}$. Effecting a rotation of the jet $\mathcal{J}_{S}$
by an angle $\Delta$, we have

$\left<\Phi|R_{\Delta}\mathcal{J}_{S}\right>=\sum_{i}c_{S}\left<\Phi|A(\Delta)|\phi_{i}\right>Phi_{i}A_{i}kphi_{k}j$
(35)

we see that when $\mathbf{\Phi}$ is constructed as above, we should rewrite

$\displaystyle\left<\Phi|R_{\Delta}\mathcal{J}_{S}\right>$
$\displaystyle=\left<R_{\Delta}^{-1}\Phi|\mathcal{J}_{S}\right>$ (36)
$\displaystyle=\sum_{i}\left<R_{\Delta}^{-1}\Phi|\phi_{i}\right>\left<\phi_{i}|\mathcal{J}_{S}\right>$
(37)
$\displaystyle=\left<\Phi|\mathbf{A}(g^{-1})|\phi_{i}\right>\left<\phi_{i}|\mathcal{J}_{S}\right>$
(38)

To this end, we propose expressing $\Phi$ in terms of some basis of filters,
each of which is equivariant. $\Phi(\vec{x})=\sum_{m}a_{m}\phi_{m}(\vec{x})$

$\displaystyle\left<\Phi,R_{\Delta}\mathcal{J}_{S}\right>$
$\displaystyle=\int_{X}\Phi^{*}(\vec{x})\mathcal{J}_{S}(R^{-1}_{\Delta}\vec{x}_{i})$
(39)
$\displaystyle=\int_{X}\Phi^{*}(R_{\Delta}\vec{x})\mathcal{J}_{S}(\vec{x})$
(40)
$\displaystyle=\int_{X}\alpha(\Delta)\Phi^{*}(\vec{x})\mathcal{J}_{S}(\vec{x})$
(41) $\displaystyle=\alpha(\Delta)\left<\Phi,\mathcal{J}_{S}\right>$ (42)

In order to endow our network with the desired structure, we will require the
filters $\Phi$ to be an element of an equivariant function space
$\mathcal{F}=\mathrm{span}\\{\phi_{1},\phi_{n}\\}$. An $n$-dimensional
equivariant function space has the property that for every $g\inG$, there is
an $n\times n$ matrix $\mathbf{A}(g)$ such that

$\mathbf{\Phi}(g\cdot\vec{x})=\mathbf{A}(g)\mathbf{\Phi}(\vec{x})\,,$ (43)

where the bold-face notation $\mathbf{\Phi}$ denotes that we express the
function $\Phi(\vec{x})=\sum c_{i}\phi_{i}(\vec{x})$ as an $n$-dimensional
column vector $(c_{1},\dots,c_{n})^{\mathrm{T}}\in\mathcal{F}$. Eq. 14 is
known as the interpolation equation.

In the case of one-parameter transformation group, the interpolation matrix
may be expressed as $A(g)=e^{gB}$ for an arbitrary matrix $B$.

For the specific case of Particle Flow networks, we shall consider
equivariance w.r.t. the group of rotations about the jet axis, $SO(2)\cong
U(1)$. Since $SO(2)$ is a one-parameter group transformation, it is always
possible to change to coordinates in which the group action resembles
translation, with infinitesimal generator $L_{t}=\partial/\partial x$.

In this case the relevant coordinates are polar coordinates in the
$(\eta,\phi)$ plane, defined by $r=\sqrt{\eta^{2}+\phi^{2}}$,
$\theta=\tan^{-1}(\phi/\eta))$. In these coordinates, a rotation of $\Delta$
about the jet axis is given by $f(R_{\Delta}\vec{x})=f(r,\theta+\Delta)$.

For any Lie group, it is possible to construct a space of equivariant
functions satisfying Eq. 12 by requiring closure under the exponential map of
the group’s generators [cheng1998theory]. In the case of $SO(2)$, the
generator of rotation is $L=-y\partial/\partial x+x\partial/\partial y$, and
any equivariant function space is of the form $h(r)$

for any such group; for $SO(2)$ any such space is spanned by functions of the
form $f(r,\theta)=e^{im\theta}R(r)$ for $m\in\mathbb{Z}$. Therefore, we will
take $\Phi$ to be of the form:

$\Phi(\vec{x})=\Phi(r,\theta)=\sum_{m=-M}^{M}R_{m}(r)e^{im\theta}=:\sum_{m=-M}^{M}\phi_{m}(r,\theta)\,,$
(44)

where $M$ is a hyperparameter associated with some maximum azimuthal
frequency, and $R_{m}:\mathbb{R}^{\geq 0}\rightarrow\mathbb{C}$ are arbitrary
functions.

Clearly, for any rotation $R_{\Delta}$, we have
$\phi_{m}(R_{\Delta}\vec{x})=\phi_{m}(r,\theta+\Delta)=e^{im\Delta}\phi_{m}(r,\theta)$.
If we view $\Phi$ as a column vector with $2M+1$ components, i.e.:

$\vec{\Phi}(\vec{x})=\left(\phi_{-M}(\vec{x}),\phi_{-(M-1)}(\vec{x}),\dots,\phi_{M}(\vec{x}))\par$
(45)

##  5 Deep Sets as Projections

The Deep Sets architecture [2017arXiv170306114Z] provides a permutation-
invariant universal approximator for continuous functions $f\colon
2^{X}\rightarrow\mathcal{Y}$ which act on sets of features $S\in 2^{X}$ where
$S=\\{s_{i}\\}_{i\in\mathrm{Ind}_{S}};s_{i}\in X$. This is the case when $f$
is of the form

$f(S)=f(\\{s_{1},\dots,s_{k}\\})=F\left(\sum_{i=1}^{|S|}\Phi(s_{i})\right)\,,$
(46)

for suitable latent space dimensionality $L$ and functions $\Phi\colon
X\rightarrow\mathbb{R}^{L}$ and $F\colon\mathbb{R}^{L}\rightarrow\mathcal{Y}$.
In particular, when $F$ and $\Phi$ are modeled by neural networks, $f$ is a
neural network which can be trained with the usual algorithms, and is
manifestly invariant with respect to the input order of the features $x_{i}$.

Strictly speaking, the universal approximator property is only proved for
functions $f$ on sets of fixed-size, i.e. on $X^{k}$. The Deep Sets theorem
relies on results for polynomial approximation, and applies only to functions
$f$ which are continuous with respect to the Euclidean metric on
$X^{k}\cong\mathbb{R}^{k\times d}$. Weaker conditions have been obtained for
universality, e.g. [NEURIPS2018_f5f8590c]. Nevertheless, the architecture has
seen successful application even when the underlying feature space $X$
represents quantities such as color, words, etc., which does not have a
meaningful topology.

When the feature space is measurable444We suppress the specification of
$\sigma$-algebras throughout the discussion, assuming only that
$(X,\Sigma,\mu)$ forms a measure space for some $\Sigma$. with measure $\mu$,
we may naturally consider the inner product between functions $f,g$ in the
Hilbert space $H=L^{2}(X,\mu)$:

$\langle f,g\rangle=\int_{X}f^{*}gd\mu\,.$ (47)

We shall also define the improper integral of the Dirac delta function on this
space as error: this requires binary operation $+$ with inverses, so we
probably require $X$ to be a locally-compact group

$\int_{X}\delta(x-y)f(x)d\mu(x):=f(y)\,.$ (48)

Now we propose that rather than considering $S$ to be a set of features $s\in
X$, we instead view it as a generalized function $\mathcal{J}_{S}\colon
X\rightarrow\mathbb{C}$. We make this association via the functional defined
as ${\mathcal{J}_{(\cdot)}\colon 2^{X}\times X}\rightarrow\mathbb{C}$ where

$\mathcal{J}_{S}(x)\mapsto\sum_{i=1}^{|S|}\delta(x-s_{i})\,.$ (49)

This object can be viewed as a “function” representing the spatial density of
the observed elements $s\in S$.

Given these definitions, we can evaluate the projection of the object
$\mathcal{J}_{S}$ onto the function $\Phi\in H$:

$\langle\Phi,\mathcal{J}_{S}\rangle=\int_{X}\Phi^{*}(x)\mathcal{J}_{S}(x)d\mu(x)=\sum_{i=1}^{|S|}\Phi^{*}(s)\,.$
(50)

We immediately recognize this quantity as the inner term of Eq. 1, for the
case of real-valued $\Phi$. Therefore, when acting on spatial (or otherwise
locally compact) coordinates $X$, the Deep Sets architecture may be understood
as a function $F$ applied to the projection of the spatial distribution of
observations onto learned filters $\Phi\in H$:

$f(S)=F\left(\langle\Phi,\mathcal{J}_{S}\rangle\right)\,.$ (51)

##  6 Energy Flow Networks

Energy Flow Networks [Komiske:2018cqr] are a special case of the Deep Sets
architecture, in which coordinate space $X\subset\mathbb{R}^{3}$ represents
the momentum 4-vectors of massless particles measured within an energetic
hadronic jet. Each particle can represented by the coordinates
$(z,\eta,\phi)$, where $z$ is an energy scale (taken to be the transverse
momentum $p_{\mathrm{T}}$) and $(\eta,\phi)\in\mathbb{R}^{2}$ are the angular
coordinates of the particle relative to the jet axis. Specifically, the EFN
corresponds to the case where $\Phi(z,\eta,\phi)$ is linear in $z$, by
construction:

$f_{\mathrm{efn}}=F\left(\sum_{i=1}^{|S|}z_{i}\Phi(\eta_{i},\phi_{i})\right)\,.$
(52)

In the notation of Eq. 48, we may write

$f_{\mathrm{efn}}=F\left(\langle\Phi,\mathcal{E}_{S}\rangle\right)$ (53)

where the functional
$\mathcal{E}_{(\cdot)}\colon\mathbb{R}^{2}\rightarrow\mathbb{C}$ defined by

$\mathcal{E}_{S}(x)\mapsto\sum_{i=1}^{|S|}z_{i}\delta(\eta-\eta_{i})\delta(\phi-\phi_{i})$
(54)

has a natural physical interpretation as the spatial distribution of energy
(or momentum) within jet.

While a network of the form in Eq. 52 clearly has much less expressive power
than Eq. 1, this particular case is of special interest in high energy physics
applications. It is shown in [Komiske:2018cqr] that for suitable $F$ and
$\Phi$, Eq. 52 is a universal approximator for a specific class of functions
on the particles of a jet, known as IRC-safe observables.

Essentially, IRC-safety means that the observable is not measurably different
for any two events which differ only by the collinear splitting of an
energetic particle, or the infrared emission of additional low-energy
particles. IRC-safety is an important invariant which applies to many
perturbative calculations in QCD theory, allowing for these theoretical values
to be compared to experimental measurements of non-perturbative phenomena. As
experimentalists rely increasingly on deep learning methods in analyses, the
EFN is an important step towards learning powerful models with appropriately
regularized representations.

##  7 Dual Space Regression

In Section 6, we see an example of how linearity of the mapping $\Phi$ with
respect to a subspace of the features $X$ can have useful properties. We will
generalize this concept further by supposing that the input feature space is
$W=X\times V$, where $X$ is required to have locally-compact topology as
described in Section 5. For the present purposes, we assume only that
${V\subseteq\mathbb{C}^{K}}$ for some dimension $K$. We say that each element
of a set $S\in 2^{W}=\\{(x_{i},v_{i})\\}$ has coordinate features $x_{i}\in X$
and qualitative features $v_{i}\in V$.

Then we may define a new functional $\mathcal{G}_{(\cdot)}\colon 2^{W}\times
X\rightarrow V$ as

$\mathcal{G}_{S}(x)\mapsto\sum_{(x_{i},v_{i})\in S}v_{i}\delta(x-x_{i})\,,$
(55)

In particular, the EFN corresponds to the case where the $v_{i}\in\mathbb{R}$
are particle energies.

Now we may write the projection operation between a filter $\Phi$ and
functional $\mathcal{G}$:

$\langle\Phi,\mathcal{G}_{S}\rangle=\int_{X}\Phi^{*}\mathcal{G}_{S}d\mu=\sum_{(x_{i},v_{i})\in
S}\Phi^{*}(x_{i})v_{i}\,.$ (56)

Here we take the map $\Phi\colon X\rightarrow V$ to be an element of some
Hilbert space. Then by the Riesz representation theorem,
$\langle\Phi,\cdot\rangle$ uniquely represents an element of the dual space
$\mathcal{H}^{*}$.

Strictly speaking, neither $S\in 2^{W}$ nor the generalized function
$\mathcal{G}_{S}$ are members of any Hilbert space; however, they do have
well-defined inner product. Therefore, by analogy with the Riesz
representation theorem, we propose that the map
$\langle\Phi,\mathcal{G}_{(\cdot})\rangle\colon 2^{W}\rightarrow\mathcal{C}$
may be thought of as an element of the “dual space” to the set of inputs
$2^{W}$.

We now consider the case where the qualitative feature space $V$ has some kind
of tensor structure associated with the coordinates $X$. That is, suppose that
$G$ is a group that acts (isometrically?) on $X$ with group action
$\cdot:G\times X\rightarrow X$ and that $G$ has a representation $L\colon
G\rightarrow GL(V)$ on $V$. Then if each element of the input space $2^{W}$
transforms with group action $(x_{i},v_{i})\mapsto(g^{-1}x_{i},L_{g}v_{i})$,
we say that $V$ has a tensor structure (specified by the representation $L$)
with respect to the coordinates $X$. In the notation of the functional in Eq.
55, a transformation $g\in G$ acts on a tensor-valued $\mathcal{G}_{S}$ via:

$\mathcal{G}_{S}(x)\overset{g}{\rightarrow}L_{g}\mathcal{G}_{S}(g^{-1}x)$ (57)

The simplest case is when $V$ is a scalar: then $L_{g}$ is the identity on $V$
for all $g\in G$. So for example, if $x_{i}$ are 2D coordinates, and $V$
represents say, the colors, or masses of some particle at that location, we
expect that under a global coordinate rotation, these features will remain the
same even though their location changes, illustrated in [figure].

Another interesting case is when $V$ transforms as a vector under $G$. Taking
again the example of 2D rotations, we consider $v_{i}$ represents the 2D
(vector) velocity of a particle at point $x_{i}$. Then under a rotation, the

When the function $\Phi$ is parameterized, trainable neural network layers can
be constructed which act on inputs in the form of Eq. 56. Such networks can be
thought of as learning elements of $\mathcal{H}^{*}$ which optimally project
from potentially high-dimensional input spaces to a single number in
$\mathbb{C}$. These layers can therefore be used as a form of “pooling” which
is naturally invariant with respect to the ordering of the input set.
Moreover, they provide an interpretable artifact, $\Phi$, which can be thought
of as a “density” map of qualitative features $v$ as a function of coordinate
features $x$.

##  8 other

Given this projective interpretation, we may generalize a bit further by
considering a network of the form Eq. 48 to be a map
$F\colon\mathbb{C}^{L}\rightarrow\mathcal{Y}$ composed with a learnable
collection of dual space elements. That is, the left-bilinear action
$\langle\Phi,\cdot\rangle$ may be considered an element of the dual space of
$2^{X}$.

This interpretation can be generalized further to the case where the feature
space is a product of spaces $X=W\times Z$, where $W$ is locally compact and
$Z$ is invariant with respect to transformations in $W$. For example $W$ could
represent spatial coordinates of objects in $\mathbb{R}^{3}$, while $Z$
represents arbitrary features of those objects, such as color, size, etc. In
this case, we write $\mathbf{x}=\\{(w_{i},z_{i})\\}$, and define
$\mathcal{J}_{\mathbf{x}}\colon W\rightarrow V$:

$\mathcal{J}_{\mathbf{x}}(x):=\sum_{(x_{i},z_{i})\in\mathbf{x}}g(z_{i})\delta(x-x_{i})\,,$
(58)

where $g\colon Z\rightarrow V$ is an arbitrary function mapping the non-
spatial coordinates to some latent representation space $V$. Now
$\mathcal{J}_{\mathbf{x}}$ maps from the spatial coordinates of $W$ to a
latent representation in $V$. We also promote the functions $\Phi$ so that
$\Phi*\colon W\rightarrow V^{*}$ takes values in the dual space to $V$. In
this work we shall always consider $V\subseteq\mathbb{C}^{K}$, so that
conjugation $\Phi\rightarrow\Phi^{*}$ corresponds to the Hermitian transpose.
However, it may be interesting to consider other cases, such as tensor-valued
maps.

Taking the definition of Eq. 58, the projections now read

$\langle\Phi,\mathcal{J}_{\mathbf{x}}\rangle=\sum_{(w_{i},z_{i})\in\mathbf{x}}g(z_{i})\Phi^{*}(w_{i})\,,$
(59)

and we may define a new network architecture analogous to Eq. 1:

$f_{\mathrm{proj}}(\mathbf{x})=F\left(\langle\Phi,\mathcal{J}_{\mathbf{x}}\rangle\right)=F\left(\sum_{(w_{i},z_{i})\in\mathbf{x}}g(z_{i})\Phi^{*}(w_{i})\right)$
(60)

We note that Eq. 60 is strictly a special case of Eq. 1; essentially all that
has changed is that the function mapping from feature space to latent space
has been factorized. In the following sections, we shall demonstrate that
factorizing the covariant and invariant features has distinct benefits in
certain applications.

As a corollary to the EFN result, we note that there is a sort of generalized
“IRC-safety” that applies to networks of the form in Eq. 60. Specifically, it
is invariant under the exchange:

$\displaystyle\\{\dots,(w_{i},g(z_{i})),\dots\\}\leftrightarrow$ (61)
$\displaystyle\\{\dots,(w_{i}+\delta,ag(z_{i})),(w_{i}+\delta^{\prime},(1-a)g(z_{i})),\dots\\}\,,$

for sufficiently small $\delta,\delta^{\prime}\in W$ and $a\in[0,1]$. In other
words, elements may be added or removed from the set $\mathbf{x}$, as long as
the total quantity of $g(z_{i})$ in a localized area remains unchanged. This
follows in analogy with the EFN result from the fact that $\Phi$ is continuous
and $\langle\Phi,\mathcal{J}_{\mathbf{x}}\rangle$ is linear in $g(z_{i})$.

##  9 Plancherel Networks

Let $G$ be a  something? group which acts on $\mathcal{J}$. The Plancherel
network is of the form:

$f_{\mathrm{plan}}(\mathbf{x})=F\left(\int_{G}\left|\langle\Phi,g\cdot\mathcal{J}_{\mathbf{x}}\rangle\right|^{2}d\mu(g)\right)\,,$
(62)

where $G$ is  some group and $\mu$ is the Haar measure on $G$. In other words,
this network computes the square modulus of the projection
$\langle\Phi,\mathcal{J}\rangle$ averaged over all possible transformations
$g\in G$.

##  10 Harmonic Flow Networks

We start with the special case of the Energy Flow Network. Both EFNs and PFN
propose to solve the issue of jet representation in neural networks by
leveraging the universal approximation theorem known as the Deep Sets theorem
[2017arXiv170306114Z]. A network $f:X\rightarrow\mathcal{Y}$ applied to a jet
$\mathcal{J}\in X$ is represented by:

$f(\mathcal{J})=F\left(\sum_{i}z_{i}\Phi(\eta_{i},\phi_{i})\right)$ (63)

We note that if we choose to view a jet $J$ as the observed energy density as
a function of the spatial coordinates $\eta$ and $\phi$, we may write

$\mathcal{J}(\eta,\phi)=\sum_{i}z_{i}\delta(\eta-\eta_{i})\delta(\phi-\phi_{i})\,.$
(64)

Then for a filter $\Phi$ which is some arbitrary function in
$L^{2}(\mathbb{R}^{2})$, we may compute the projection of the jet onto the
filter using the ordinary inner product

$\left<\Phi,\mathcal{J}\right>=\int d\eta d\phi\
\Phi^{*}(\eta,\phi)\mathcal{J}(\eta,\phi)=\sum_{i}z_{i}\Phi^{*}(\eta_{i},\phi_{i})\,,$
(65)

which we may immediately recognize as the inner term of Eq. 52, for the case
of real-valued $\Phi$. In other words, we may understand energyflow networks
as a neural network $F$ applied to the projection of the jet’s spatial energy
distribution upon a learned filter $\Phi$.

We will first define the specific case of a Harmonic Energy Flow network in
order to demonstrate its advantages, while preserving the IRC-safe and
permutation invariant features of the EFN. We will then generalize the concept
to show that Harmonic Flow networks can be used on a variety problems.

We begin by noting that an arbitrary $L^{2}$ function
$g:\mathbb{R}^{2}\rightarrow\mathbb{C}$ may be expressed in polar coordinates
in the form

$\displaystyle g(r,\theta)$
$\displaystyle=\sum_{m=0}^{\infty}\sum_{n=m}^{\infty}A_{nm}\Psi_{nm}(r,\theta)\,;$
(66) $\displaystyle\Psi_{nm}(r,\theta)$
$\displaystyle=\sqrt{\frac{2n!}{(n+m)!}}e^{-r^{2}/2}r^{m}L^{m}_{n}(r^{2})e^{im\theta}\,,$
(67)

where $L^{m}_{n}$ are the generalized Laguerre polynomials. This is the so-
called Laguerre-Gaussian series, and it may come in handy later. For now, it
is sufficient to note in particular that we may therefore choose to write an
arbitrary filter $\Phi$ in the form

$\displaystyle\Phi(r,\theta)$ $\displaystyle=\sum_{m}\Phi_{m}(r,\theta)\,;$
(68) $\displaystyle\Phi_{m}:$ $\displaystyle=R_{m}(r)e^{im\theta}\,,$ (69)

where $R_{m}(r)$ depends only on the radius. This is essentially a multipole
expansion of $\Phi$.

Assuming $\Phi$ to take this form, we may now write

$\left<\Phi,\mathcal{J}\right>=\sum_{m}\left<\Phi_{m},\mathcal{J}\right>\,.$
(70)

While $\left<\Phi,\mathcal{J}\right>$ has no particular equivariance with
respect to rotations about the jet axis, the individual
$\left<\Phi_{m},\mathcal{J}\right>$ do. Under a rotation $R_{\Delta}$ about
the jet axis,

$\left<\Phi_{m},\mathcal{J}\right>\rightarrow\left<\Phi_{m},R_{\Delta}\mathcal{J}\right>=e^{im\Delta}\left<\Phi_{m},\mathcal{J}\right>\,.$
(71)

So the projections of the jet onto the individual angular frequency modes are
linearly equivariant with respect to rotations about the jet axis; in
particular, they just pick up a phase.

This immediately suggests replacing the argument to $F$ in Eq. 52 with a
vector of $M+1$ projections onto the definite-frequency modes of the filter
$\Phi_{m}$:

$f(\mathcal{J})=F\left(\left[\left<\Phi_{0},\mathcal{J}\right>,\dots,\left<\Phi_{M},\mathcal{J}\right>\right]^{\mathrm{T}}\right)\,.$
(72)

This forces the overall network to preserve equivariance w.r.t. rotations at
least until the first layer of the subnetwork $F$. We call this the
equivariant EFN, and note that this formulation is strictly more general than
the EFN of Eq. 52, since $F$ is always free to learn to simply sum the
individual components of the input vector; in the case of a neural network,
this would correspond to a neuron connected to the input with all weights
equal to one. Therefore, Eq.72 possesses the same universal-approximation
guarantees as Eq. 52, and since each of the $\Phi_{m}$ are continuous, this
form also preserves the IRC-safety.

While we shall show that preserving equivariance to a deeper level of
representation within the overall network $f$, it is unlikely that the outer
subnetwork $F$ will be able to continue to learn equivariant representations
without additional structure. It turns out that the natural way to ensure this
is by promoting $F$ to be a $U(1)$ group-convolutional architecture; we treat
this topic in detail in [OTHER PAPER].

Since the task at hand (quark vs. gluon tagging) is manifestly rotationally
invariant, we shall consider imposing rotational invariance, a special case of
equivariance. This can be easily accomplished by simply taking the magnitude
of the projection onto each $m$ mode, i.e. by replacing
$\left<\Phi_{m},\mathcal{J}\right>$ with
$\left|\left<\Phi_{m},\mathcal{J}\right>\right|^{2}$ in Eq. 72. However, in so
doing, we no longer can guarantee that every observable learned by Eq. 52 can
be learned in this form since. This is obvious, for instance, if
$\left<\Phi,\mathcal{J}\right>$ happens to be negative.

It turns out that simply summing the projections onto the various $m$ modes in
quadrature is a very powerful way to ensure equivariance while maintaining the
universally expressive and IRC-safe features of the EFN 52. This seems
counterintuitive, as it is strictly less expressive than Eq. 72. However,
using the equivariance property of Eq. 71, we note that the projection of a
rotated jet, $R_{\Delta}$ onto the filter $\Phi$ is

$\langle\Phi,R_{\Delta}\mathcal{J}\rangle=\sum_{m}\langle\Phi_{m},\mathcal{J}\rangle
e^{im\Delta}\,,$ (73)

where $\langle\Phi_{m},\mathcal{J}\rangle$ is simply a complex number.
Therefore, by Parseval’s theorem we have the fundamental result:

$\sum_{m}\left|\langle\Phi_{m},\mathcal{J}\rangle\right|^{2}=\frac{1}{2\pi}\int_{-\pi}^{\pi}d\Delta\left|\langle\Phi,R_{\Delta}\mathcal{J}\rangle\right|^{2}\,.$
(74)

We therefore propose the formulation:

$\displaystyle f(\mathcal{J})$
$\displaystyle=F\left(\sum_{m}\left|\left<\Phi_{m},\mathcal{J}\right>\right|^{2}\right)=F\left(\frac{1}{2\pi}\int_{-\pi}^{\pi}d\Delta\left|\langle\Phi,R_{\Delta}\mathcal{J}\rangle\right|^{2}\right)\,;$
(75) $\displaystyle\Phi_{m}(r,\theta)$ $\displaystyle=R_{m}(r)e^{im\theta}\,.$
(76)

which we refer to as the invariant form of the Harmonic Flow network. It is so
called because it is invariant with respect to rotations about the origin.

##  11 Experiments

| Model | AUC | $1/\epsilon_{b}|_{\epsilon_{s}=50\%}$
---|---|---|---
IRC-safe |  |  |
| EFN | 0.8824 | 28.6
| s-EFN (Ours) | 0.8895 | 30.5
w/o PID |  |  |
| ResNeXt-50 | 0.8960 | 30.9
| P-CNN | 0.8915 | 31.0
| PFN | 0.8911 | 30.8
| s-PFN (Ours) | 0.8958 | 32.4
| ParticleNet-Lite | 0.8993 | 32.8
| ParticleNet | 0.9014 | 33.7
w/ PID |  |  |
| P-CNN | 0.9002 | 34.7
| PFN-Ex | 0.9005 | 34.7
| s-PFN-Ex (Ours) | 0.9058 | 37.6
| ParticleNet-Lite | 0.9079 | 37.1
| ParticleNet | 0.9116 | 39.8
Table 1:  Comparison of similar network architectures on the quark-gluon
tagging dataset of Ref. [Komiske:2018cqr]. ResNeXt is a deep 2D CNN adapted
for jet images in [Qu:2019gqs], while P-CNN [CMS-DP-2017-049] is a 1D
convolution operating on ordered lists of particles, also implemented in
[Qu:2019gqs]. Our s-EFN model is slightly more sensitive than the EFN, while
maintaining IRC-safety. The s-PFN model achiaves comparable performance to
ParticleNet without adding more inputs, by no longer guaranteeing IRC-safety.
The s-PFN with experimental PID information performs comparably to
ParticleNet.

##  12 Conclusion