Title: Rethinking SO(3)-equivariance with Bilinear Tensor Networks
Venue: NeurIPS 2023
Accepted: no
Keywords: equivariance

so(3) symmetry

tensor data

physics
Abstract: Many datasets in scientific and engineering applications are comprised of objects which have
specific geometric structure. A common example is data which inhabits a representation of the
group SO(3) of 3D rotations: scalars, vectors, tensors, etc. One way for a neural network to
exploit prior knowledge of this structure is to enforce SO(3)-equivariance throughout its layers, and
several such architectures have been proposed. While general methods for handling arbitrary SO(3)
representations exist, they computationally intensive and complicated to implement. We show that
by judicious symmetry breaking, we can efficiently increase the expressiveness of a network operating
only on vector and order-2 tensor representations of SO(2). We demonstrate the method on an
important problem from High Energy Physics known as b-tagging, where particle jets originating
from b-meson decays must be discriminated from an overwhelming QCD background. In this task,
we find that augmenting a standard architecture with our method results in a 2.3× improvement in
rejection score.

Abstract

Many datasets in scientific and engineering applications are comprised of objects which have specific geometric structure. A common example is data inhabiting a representation of SO(3) scalars, vectors, and tensors. One way to exploit prior knowledge of the structured data is to enforce SO(3) -equivariance. While general methods for handling arbitrary SO(3) representations exist, they can be computationally intensive and complicated to implement. We show that by judicious symmetry breaking, we can efficiently increase the expressiveness of a network operating on these representations. We demonstrate the method on an important classification problem from High Energy Physics known as b-tagging . In this task, we find that our method achieves a 2 . 7 × improvement in rejection score over standard methods.

1 Introduction

In many Machine Learning (ML) applications, at least some of the data of interest have specific geometric structure. For example, position measurements from LiDAR imaging, the configuration of atoms in molecular potentials, and measurements of particle momenta are all cases where the data are naturally represented as spatial 3-vectors. However, classical Neural Network (NN) architectures are not well suited to this sort of data; for instance, the standard Multi Level Perceptron would require that all information, spatial or otherwise, must be collapsed into a flat list of features as input to the network. In this case, the spatial nature of the data, while not lost, is not communicated a priori nor enforced post hoc .

More recently, developments in the field of Representation Learning have shown that equivariant 21 NNs are a natural way to accommodate structured data, and in many cases lead to substantially 22 improved algorithms. Very informally, a function (such as a NN) is called equivariant if the output 23 transforms similarly to the input. 24

25

26

27

28

29

30

31

32

Convolutional Neural Networks (CNNs) are the prototypical example of this. CNNs exploit the fact that image data can be most naturally represented as data on a discrete 2-dimensional grid. This data structure is associated with the representation of the group of discrete translations. The standard CNN layer takes advantage of this by operating on input grid (pixel) data with discrete translation operations, and returning outputs on a similar grid structure. Because the output of each layer has the same representational structure as the input, it is straightforward to build very deep representations without destroying the prior spatial structure of the data, simply by stacking CNN layers. The result, of course, is that CNNs have completely revolutionized the field of computer vision.

We specifically consider the case of continuous scalar and 3-dimensional vector point data, as may 33

be encountered in many point-cloud datasets. For these data, the natural group associated with their 34

35

representation is

SO(3)

, the set of 3D rotations. Therefore, one strategy to incorporate this structure

into a neural architecture is to enforce equivariance w.r.t. SO(3) , and several such architectures 36 have been proposed [1, 2, 3]. In general, these approaches achieve equivariance either by defining 37 a spherical convolutional operation [3, 1], or by constraining the network's operations to maintain 38 strict representational structure [2, 4]. 39

Our method follows the latter approach, but in a much simpler way. Rather than concerning 40 ourselves with arbitrary (2 ℓ +1) -dimensional representations, we consider only a few physically 41 relevant representations: scalars, vectors, and order-2 tensors. For these three representations, it 42 is straightforward to enumerate the options for linear neuron layers. We also want our network 43 to be able to exchange information between different representations. The Clebsh-Gordon theory 44 prescribed in other methods provides the most general method for projecting arbitrary tensor products 45 between representations back into irreducible representations, However, once again we take a similar 46 approach, and instead introduce a simple Tensor Bilinear Layer , a subset of the CG space that consists 47 of commonly known and physically intuitive operations, such as the vector dot product and cross 48 product. 49

50

51

52

53

54

55

56

Importantly, we propose a novel method that allows us to relax equivariance requirements when an axial symmetry is present, by allowing the global SO(3) symmetry to be locally broken down to SO(2) . These looser conditions allow us to design of models that enforce only the instantaneously relevant equivariance, and allows the network to learn more expressive functions at each layer. We show that this kind of equivariant neuron is generally only possible with the introduction of order-2 tensor representations, but we provide an efficient implementation for vector-valued networks that constructs only the minimal tensors required.

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

To illustrate a real-world application to data with an axial symmetry, we introduce a common problem from the field of High Energy Physics (HEP), described in Sec. 2. In Sec. 3, we describe the modular elements of our method, from which a wide variety of neural architectures may be composed. In Sec. 4, we describe a specific architecture based on Deep Sets [5] which will serve as a baseline model, and we illustrate how to adapt this architecture using our approach. In Sec. 5, we describe the simulated data used for training and evaluation, and describe the results of a progressive implementation of the modules developed herein. Finally, we offer concluding remarks in Sec. 6.

1.1 Related Work

From the field of High Energy Physics, there has been much work in applying various DL approaches to jet tagging [6, 7, 8, 9] in general and b-tagging in particular [10, 11]. The present work seeks to build on this effort by offering novel neural architectures that can be adapted into next-generation applications.

From the field of Machine Learning, there have been numerous prior works on SO(3) equivariant models [1, 2, 3, 4]. In general, these approaches depend on Clebsh-Gordon (CG) decomposition and/or sampling from spherical harmonics. While our approach is more similar to the CG method, it is simpler and more relevant for the task at hand. Moreover, we innovate on the the space of allowed equivariant operations by relaxing the global SO(3) symmetry which is relevant for our particular application.

1.2 Novel Developments

The main innovation of this paper is to expand the set of linear equivariant maps in the special case where there is a 'special' direction in space, which may change from sample to sample. In this case, it is possible to maintain global SO(3) equivariance, while breaking the per-layer equivariance condition down to a locally-defined SO(2) symmetry, which is parameterized by the special direction. We also innovate by introducing a simpler method of forming SO(3)-equivariant nonlinearities, by simply introducing familiar bilinear operations on spatial representations such as scalars, vectors, and tensors. In addition to the nonlinearity provided by the bilinear operations, we also introduce simple nonlinear activation functions on the vector and tensor representations, which we find helps stabilize training and improve performance.

Lastly, from the physics perspective, we propose a significant departure from standard practice, by stipulating that our b-tagging should be provided with raw 3-dimension position and momentum

information, as this is the only way to ensure that SO(3)/SO(2) equivariance is exploited.

While we demonstrate our methods using a specifc architecture based on Deep Sets [5], we expect 88 these innovations can be useful in many other applications. Given the modularity and strictly defined 89 input and output representations of each layer, these elements could be used to augment other neural 90 architectures such as convolutional, graph, and transformers as well. 91

92

2 B-jet Identification at LHC Experiments

In HEP experiments, such as ATLAS [12] and CMS [13] at CERN, b-jets are a crucial signal for 93 studying rare phenomena and precision physics at the smallest scales of nature. A jet is a collimated 94 spray of hadronic particles originating from energetic quarks or gluons produced in high energy 95 particle collisions. A b-jet is a jet which specifically originates from a b -quark; when these quarks 96 hadronize, they form metastable B-mesons which travel some distance from the collision origin 97 before decaying, emitting particles from a secondary, displaced vertex. 98

99

100

101

102

103

Charged particles originating from these vertices are measured with tracking detectors and are often referred to as tracks . Due to the displacement of the secondary vertex, when track trajectories originating from B-meson decays are extrapolated backwards, they are generally not incident to the origin. Therefore, we instead measure the distance to the point of closest approach; this is often referred to as the track impact parameter , which is a 3-vector quantity that we denote with a .

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

In most applications, only the transverse and longitudinal components, d 0 and z 0 , of this impact parameter are examined [14]. The magnitude of these projections is the most distinctive feature that indicates whether a particular jet originated from a b -quark.

The inspiration for this work was the observation that the physical processes which govern how particles within a jet are produced and propagated are largely invariant with respect to rotations about the jet axis , denoted ˆ j . This is the unit vector in the direction of the aggregate jet's momentum vector. On the other hand the standard b -tagging observables d 0 and z 0 have no well-defined transformation rule under rotations, i.e. they are not part of a covariant representation.

Previous works [8] have demonstrated that networks which exploit this natural SO(2) symmetry can greatly improve performance, but these methods all essentially rely on reducing the problem to vectors in a 2-dimensional plane. In order to obtain an equivariant representation in the case of b -jets, we must consider the full 3-dimensional structure of the impact parameter, which transforms as a vector under general rotations a R → R a . In addition to the 3-dimensional impact parameter a , we also have information about the track's momentum p and various scalar quantities such as the particle's charge, energy, and a limited identification of the particle type.

In the next section, we will describe modular neural elements that can solve this problem, by allowing a network to admit a global SO(3) symmetry which preserves the scalar and vector representations, while also breaking SO(3) down to the more physically appropriate SO(2) whenever possible.

3 Network Elements

Our proposed method depends on three modular elements, described in detail in the following subsections. The overall strategy begins by mirroring what has proved to work for NNs in general: we interleave simple linear (or affine) layers with nonlinear activation functions, in order to learn powerful models. For an equivariant network, we first need to identify a set of linear equivariant maps suitable for the symmetry at hand. In our case, we come up with two sets of such maps: a global SO(3) -equivariant affine layer, and a locally SO(2) ˆ j -equivariant linear layer.

Since we also require our network to mix between its scalar, vector, and tensor representations, we introduce an equivariant bilinear layer. Lastly, we define SO(3) equivariant nonlinear activations for each output representation.

In Sec. 4, we demonstrate how to combine these elements into a complete neural architecture. This

133

architecture is based on the Deep Sets [5] architecture suitable for variable-length, permutation-

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

3.1 SO(2) ˆ j -equivariant Linear Layers

A well-known way to ensure equivariance w.r.t. any group is to broadcast the neural action across the representational indices of the data [15, 16]. That is, the neural weight matrix simply forms linear combinations of the features in their representation space. In general, it is helpful to add a bias term, but care must be taken to select one that preserves equivariance.

The simplest example of this is for a collection of F scalar input features, { s i } , mapping to a collection of K output features. The scalar has no representational indices, so this simply amounts to the standard affine 1 network layer

where the learnable parameters W ij and b i are the neural weights and bias terms, respectively. In the vector case, we may generalize to

Note that the equivariance condition for vector-valued functions f ( R v ) = Rf ( v ) implies that R b = b for arbitrary rotation R ; hence, the bias vector must be zero. Finally, the analogous case for order-2 tensors is:

where again we have learnable scalar parameters b i . In this case, the equivariance condition is f ( RTR T ) = Rf ( T ) R T , which implies that RBR T = B , i.e. B must commute with arbitrary R . Therefore, B must be proportional to the identity tensor I .

The above neurons are purely isotropic in SO(3) . However, as discussed in Sec. 1, for our problem we have prior knowledge that the distribution is symmetric about a specific axis. At worst, having only isotropic operations can over-regularize the network by imposing too much structure, and at best it might be harder for the network to spontaneously learn about the axial symmetry. We therefore consider the most general linear map is equivariant w.r.t. the axial symmetry. Since this is a lesser degree of symmetry, the network should have greater freedom in choosing linear maps.

3.1.1 Vector Case

Let ˆ j be a unit vector (in our application, the jet's momenutm axis) which is instantaneously fixed per batch input. The rotations about this axis define a proper subgroup S ˆ j ⊂ SO(3) where we identify S ˆ j ∼ = SO(2) . We therefore refer to this subgroup as SO(2) ˆ j ⊂ SO(3) ; the distinction being that SO(2) ˆ j fixes a representation on R 3 which depends on ˆ j .

The set of all linear SO(2) ˆ j -equivariant maps is exactly the set of matrices A which commute with arbitrary R ˆ j ∈ SO(2) ˆ j , which are of the form

for arbitrary learnable parameters ¯ θ = ( a, b, ϕ ) . The first two terms represent anisotropic scaling 164 in the directions parallel and perpendicular to ˆ j , respectively. The third term represents any other 165 arbitrary rotation about the ˆ j axis, parameterized by a single angle ϕ . 166

167

168

169

170

171

172

Because A commutes with all R ˆ j ∈ SO(2) ˆ j , the linear layer defined by

is SO(2) ˆ j -equivariant, for arbitrary parameters ¯ θ ij . The isotropic linear neuron of Eq. 1 corresponds to the special case a ij = b ij , ϕ ij = 0 .

3.1.2 Tensor Case

In order for a tensor-valued linear map L to be equivariant, we require that L ( R ˆ j TR T ˆ j ) = R ˆ j ( LT ) R T ˆ j . Note that in the case of full SO(3) equivariance, the only option is for L to be proportional to the

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

identity. Without loss of generality, we may assume the order-4 tensor L can be written as a sum 173 of terms A ⊗ B for some order-2 tensors A,B . The tensor product acts on an order-2 tensor T as 174 ( A ⊗ B ) T := ATB T . Taking L to be of this form (up to linear combinations), the equivariance 175 condition reads A ( R ˆ j TR T ˆ j ) B T = R ˆ j ( ATB T ) R T ˆ j . This is satisfied when both A and B commute 176 with R ˆ j ; we have already identified the set of such matrices in Eq. 4. Therefore, we define the action 177 of the tensor-valued SO(2) ˆ j linear layer by: 178

where the parameters ( ¯ θ ij , ¯ φ ij ) are the neural connections and we also allow for an affine bias term parameterized by b i , which is proportional to the identity tensor and hence also equivariant.

3.2 Tensor Bilinear Operations

So far we have provided two means for working with data in the SO(3) scalar, vector, and order-2 tensor representations. However, we also desire a means for allowing information between the different representations to be combined and mixed.

The most general approach to this is addressed by Clebsh-Gordon theory [2, 4]. But we adopt a simpler approach, wherein we take advantage of the familiar representations of our data and employ common bilinear operations such as dot products and cross products for vectors 2 . This allows the network to create a mixing between different representations. The operations considered are enumerated schematically in Fig. 1. In order to form these terms, the bilinear layer requires that the scalar, vector, and tensor inputs ( s, v , T ) all have the same size, 2 F , in their feature dimension, and that the size is a multiple of two. We then split the features into groups of two: s a = { s i } i =1 ..F , s b = { s i } i = F +1 .. 2 F , and define similarly v a,b and T a,b .

After effecting all of the options from Fig. 1, the layer returns scalar, vector, and tensor outputs with 3 F features each.

3.3 SO(3) -equivariant Nonlinear Activations

For the scalar features, any function is automatically equivariant. Therefore, for these features we use the well-known ReLU[17] activation function, although any alternative nonlinearity would also work.

In the vector and tensor cases, care must be taken to ensure equivariance. For the vector case, we state a simple theorem[18]:

Theorem 3.1 For any vector-valued function f : R 3 → R 3 which satisfies f ( R x ) = Rf ( x ) for all R ∈ SO(3) , there exists a scalar function ˜ f such that

where ˆ x = x / | x | when | x | > 0 and ˆ x = 0 otherwise.

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

In other words, we may chose an arbitrary, nonlinear function ˜ f which acts only on the vector magnitude, and the layer must leave the direction of the input unchanged. This leaves many possibilities; after some experimentation, we found the following activation, which we call Vector ReLU ( VReLU ), works well:

The VReLU activation is analogous to the standard rectified linear unit, except that the transition from linear to constant happens at a fixed positive magnitude rather than zero. We found that in particular, the saturating aspect of VReLU greatly helps to stabilize training, as otherwise the vector features tend to coherently grow in magnitude, leading to exploding gradients.

For the order-2 tensor case, we note here that the tensor analog to Theorem 3.1 is much more nuanced[18], and in general depends on three principal invariants I 1 , I 2 , I 3 . For simplicity, we define the Tensor ReLU ( TReLU ) completely analogously to the vector case, and leave a more complete analysis of tensor nonlinearities to future work:

4 Benchmark Architectures

We now have defined the four modular elements which provide the appropriate equivariant operations. In order to evaluate the practical effects of these modules, we define a benchmark architecture that is based on the Deep Sets architecture[5], also referred to as a Particle Flow Network (PFN) [19] in the field of HEP. The PFN is a commonly-used architecture for this sort of problem in real-world applications such as at the ATLAS experiment[14].

We will first define the standard PFN architecture, which will serve as our baseline in experiments. Then, we describe a modified version at the module level using the analogous equivariant operations in place of the standard neural network layers.

4.1 Particle Flow Network

The basic structure of the PFN [19] is based on the Deep Sets [5] architecture, and will serve as our baseline. It is of the form:

where Φ : R F → R L and F : R L → Y are arbitrary continuous functions parameterized by neural networks. L is the dimension of the latent embedding space in which the particles are aggregated and P is the number of particles in an observed jet. Y represents the relevant output space for the task at hand; since our task is classification, we consider Y = [0 , 1] .

The input features { p k } represent the observed track particles within the jet. These features include:

The jet 3-momentum in detector coordinates, ( p ( J ) T , η ( J ) , ϕ ( J ) )

The 3-momentum of each particle track in relative detector coordinates, ( p k T , ∆ η k , ∆ ϕ k )

The track impact parameters of each particle ( d k 0 , z k 0 )

The particle's charge q and particle type {electron, muon, hadron}

For each jet, we allow up to P = 30 particle tracks; inputs with fewer than 30 particles are padded with zeros. We also repeat the jet 3-momentum over the particle axis and concatenate with the rest of the per-particle features. The discrete particle type feature is embedded into 3 dimensions. After concatenating all features, the input to the PFN is of shape ( ∗ , P, F ) where F = 12 is the feature dimension.

The subnetworks Φ and F are simple fully-connected neural networks. Φ consists of two hidden layers with 128 units each, and ReLU activation. The output layer of Φ has L units and no activation applied. The F network consists of three hidden layers with 128 units each and ReLU activations.

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

The final output layer has two units with no activation, in order to train with a categorical cross entropy objective.

4.2 Vector and Tensor PFN

We now adapt the basic PFN architecture and promote it to what we term a Vector PFN (VPFN) or Tensor PFN (TPFN), according to the highest representation included. The overall architecture is of the same form as Eq. 9; we will simply modify the detailed implementation of the Φ and F sub-networks.

The first change is that the input features now belong strictly to one of the three SO(3) representations: scalar, vector, or order-2 tensor:

In general, the number of features in any of the representation channels are independent. The features for the TPFN experiments include:

The jet 3-momentum in Cartesian coordinates ( p ( J ) x , p ( J ) y , p ( J ) z )

The 3-momentum of each particle track p k

The 3-position of the track's point of closest approach to the origin a k

The charge and particle type of each track, as described in Sec. 4.1

As before, we replicate the jet momentum across the particle index, and we embed the particle type into 3 dimensions, resulting in F s = 4 scalar and F v = 3 vector features. Since there are no observed tensor features for particle tracks, we synthesize an initial set of features to act as a starting point for the tensor operations. This is done by taking the outer product between all combinations of the three available vector features, resulting in F t = 9 features.

We now have Φ : R F s × 3 F v × 9 F t → R L × 3 L × 9 L , where F s , F v , F t are the number of scalar, vector, and tensor inputs, respectively. A single layer of Φ is formed as shown in Fig. 2, by combining in sequence the Affine, SO(2) ˆ j -Linear, Bilinear, and Nonlinear modules described in Sec. 3. The network consists of two hidden and one output layer. Each hidden Affine layer of the Φ network contains 2 F = 128 features per representation, which results in 3 F = 192 features after the Bilinear layer. The output of the Φ sub-network had L features, and there is no Bilinear or Nonlinear layers applied.

The F network is built similarly to the Φ network, except that it has three hidden tensor layers. In lieu of an output layer, after the hidden tensor layers, the F network computes the square magnitude of each vector and tensor feature, in order to create a final set of 3 × 3 F scalar invariant features. The scalar features are concatenated, passed through two more hidden layers with 128 units each and ReLU activations, and finally to an output layer with two units and no activation.

5 Experiments

To train b-tagging algorithms, we must use Monte Carlo simulations of particle collision events, as this is the only way to get sufficiently accurate ground truth labels. The optimization of these algorithms is commonly studied by experiments such as ATLAS and CMS, which use highly detailed proprietary detector simulation software, and only limited amounts data are available for use outside of the collaborations. [20] There are also some community-generated datasets available for benchmarking [7], however none of these publicly-available datasets contain the key information that our method leverages for its unique equivariant approach. Specifically, our model requires the full 3-dimensional displacement vector of each track's impact parameter, whereas the existing datasets only retain the transverse and longitudinal projections d 0 and z 0 . Therefore, we have created a new dataset for b-jet tagging benchmarks, to be made publicly available. The data is generated using standard Monte Carlo tools from the HEP community.

We begin by generating inclusive QCD and t ¯ t events for background and signal, respectively, using PYTHIA8[21]. PYTHIA handles sampling the matrix element of the hard processes at √ s = 13 TeV ,

Tensor Bilinears

TReLU

ReLU

SO(2) Linear

VReLU

Tensor Bilinears

SO(2) Linear

F

Tensor Bilinears

TReLU

ReLU

SO(2) Linear

VReLU

NN Layer(s)

Covariant Inputs

Invariant Output

the parton shower, and hadronization. The hadron-level particles are then passed DELPHES[22], a 290 fast parametric detector simulator which is configured to emulate the CMS[13] detector at the LHC. 291

After detector simulation, jets are formed from reconstructed EFlow objects using the antik T [23, 24] 292 clustering algorithm with radius parameter R = 0 . 5 . Only jets with p T > 90 GeV are considered. For 293 the signal sample, we additionally only consider jets which are truth-matched to a B-meson. Finally, 294 the highestp T jet is selected and the track and momentum features are saved to file. 295

The training dataset consists of a balanced mixture of signal and background with a total of 1M jets. 296 The validation and test datasets contain 100k signal jets each. Due to the high degree of background 297 rejection observed, we must generate a substantially larger sample of background events for accurate 298 test metrics, so the validation and test datasets contain 300k background jets each. 299

300

301

302

303

304

305

306

307

308

309

310

311

5.1 Results

To quantify the performance of our model, we consider the following metrics in our experiments. First, the loss function used in the training is sparse categorical cross entropy, which is also used in the validation dataset. We also consider the area under the ROC curve (AUC) as an overall measure of the performance in signal efficiency and background rejection. We also consider the background rejection at fixed efficiency points of 70% and 85% , labeled by R 70 and R 85 , respectively. Background rejection is defined as the reciprocal of the false positive rate at the specified true positive rate.

A summary of a variety of experiments is given in Table 1. The numbers in the table represent the median test score over 10 training runs, where the test score is always recorded at the epoch with the lowest validation loss. The quoted uncertainties for the rejections are the inter-quartile range.

5.2 Discussion

Table 1 shows that the family of models with only vector representations can indeed improve over 312 the baseline, provided that we include at least the bilinear layer allowing the vector and scalar 313 representations to mix. Moreover we find that adding the the SO(2) linear operations gives the vector 314 network access to a minimal set of order-2 tensors, R ˆ j , ˆ j ˆ j T , and I to enable it to exploit the axial 315 symmetry of the data. 316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

)

In the case of the tensor family of models, there is a less substantial improvement when adding the SO(2) linear layer. We expect that this is because the network with only bilinear operations is, at least in theory, able to learn the relevant operations on its own. Nonetheless, there is some improvement when adding this layer, so it would be reasonable to include both unless computational constraints are a concern.

Finally, we note that neither family of models performs even as well as the baseline, when no bilinear operations are allowed. This clearly demonstrates the effectiveness of a network which can mix SO(3) representations.

6 Conclusion

In this work, we have introduced four modules of neural network architecture that allow for the preservation of SO(3) symmetry. The Tensor Particle Flow Network (TPFN) shows promising results in our dataset, yielding up to 2 . 7 × improvement in background rejection, compared to the simple Particle Flow baseline model. We emphasize that the overall architecture of the PFN and TPFN are nearly identical; the improvement is entirely due to a drop-in replacement of standard neural layers with our covariant and bilinear layers. We also note that in our approach, the TPFN outputs a scalar which is rotationally invariant. However, it is also possible to obtain a covariant output by simply not apply the scalar pooling operations. This could be useful for many other applications, such as regression tasks, where covariant predictions are desired.

Moreover, we show that second-order tensor representations are required in order to exploit a locallyrestricted class of equivariance with respect to the axial rotations SO(2) ˆ j . When computational constraints are a concern, it is possible to recover most of the performance of the Bilinear Tensor Network, by restricting it to a faster Bilinear Vector Network with the appropriate SO(2) equivariant linear layer.

While the example application demonstrated here is of particular interest to the field of HEP, we 340 expect our method can have great impact in other ares where vector-valued point cloud data is used. 341 Finally, we note that while we demonstrated the modular elements of the TBN on a simple Deep Sets 342 / PFN type network, it should also be possible to use these modules for creating equivariant Graph 343 and attention based networks. 344

345

346

347

348

349

350

351

352

353

354

355

References

Nathaniel Thomas, Tess E. Smidt, Steven M. Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick F. Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds. ArXiv , abs/1802.08219, 2018.

Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-gordan nets: a fully fourier space spherical convolutional neural network. In Neural Information Processing Systems , 2018.

Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. In Neural Information Processing Systems , 2018.

Brandon M. Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In Neural Information Processing Systems , 2019.

Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, 356 and Alex Smola. Deep sets. ArXiv , abs/1703.06114, 2017. 357

Anja Butter et al. The Machine Learning landscape of top taggers. SciPost Phys. , 7:014, 2019. 358

Huilin Qu, Congqiao Li, and Sitian Qian. Particle Transformer for Jet Tagging. 2 2022. 359

Chase Shimmin. Particle Convolution for High Energy Physics. 7 2021. 360

Chase Shimmin, Peter Sadowski, Pierre Baldi, Edison Weik, Daniel Whiteson, Edward Goul, 361 and Andreas Søgaard. Decorrelated Jet Substructure Tagging using Adversarial Neural Networks. 362 Phys. Rev. D , 96(7):074034, 2017. 363

Daniel Guest, Julian Collado, Pierre Baldi, Shih-Chieh Hsu, Gregor Urban, and Daniel Whiteson. 364 Jet Flavor Classification in High-Energy Physics with Deep Neural Networks. Phys. Rev. D , 365 94(11):112002, 2016. 366

Graph Neural Network Jet Flavour Tagging with the ATLAS Detector. Technical 367 report, CERN, Geneva, 2022. All figures including auxiliary figures are available 368 at https://atlas.web.cern.ch/Atlas/GROUPS/PHYSICS/PUBNOTES/ATL-PHYS-PUB-2022369 027. 370

G. Aad et al. The ATLAS Experiment at the CERN Large Hadron Collider. JINST , 3:S08003, 371 2008. 372

S. Chatrchyan et al. The CMS Experiment at the CERN LHC. JINST , 3:S08004, 2008. 373

Georges Aad et al. Configuration and performance of the ATLAS b -jet triggers in Run 2. Eur. 374 Phys. J. C , 81(12):1087, 2021. 375

Jeffrey Wood and John Shawe-Taylor. Representation theory and invariant neural networks. 376 Discret. Appl. Math. , 69:33-60, 1996. 377

Marc Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing 378 equivariant multilayer perceptrons for arbitrary matrix groups. ArXiv , abs/2104.09459, 2021. 379

Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann 380 machines. In International Conference on Machine Learning , 2010. 381

C. S. Jog. Introduction to Tensors , volume 1, page 1-136. Cambridge University Press, 3 382 edition, 2015. 383

Patrick T. Komiske, Eric M. Metodiev, and Jesse Thaler. Energy flow networks: deep sets for 384 particle jets. Journal of High Energy Physics , 2019(1), jan 2019. 385

Kimmo Kallonen. Sample with jet properties for jet-flavor and other jet-related ml studies, 386 2019. 387

388

389

390

391

392

393

394

395

396

397

Christian Bierlich, Smita Chakraborty, Nishita Desai, Leif Gellersen, Ilkka J. Helenius, Philip Ilten, Leif Lonnblad, Stephen Mrenna, Stefan Prestel, Christian T. Preuss, Torbjorn Sjostrand, Peter Skands, Marius Utheim, and Rob Verheyen. A comprehensive guide to the physics and usage of pythia 8.3. SciPost Physics Codebases , 2022.

J. de Favereau, C. Delaere, Pavel Evgen'evich Demin, Andrea Giammanco, Vincent Lemaître, Alexandre Mertens, Michele Selvaggi, and The Delphes 3 collaboration. Delphes 3: A modular framework for fast-simulation of generic collider experiments. Journal of Physics: Conference Series , 523:012033, 2013.

Matteo Cacciari, Gavin P. Salam, and Gregory Soyez. FastJet User Manual. Eur. Phys. J. C , 72:1896, 2012.

Matteo Cacciari and Gavin P. Salam. Dispelling the N 3 myth for the k t jet-finder. Phys. Lett. 398 B , 641:57-61, 2006. 399

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes n permutation and SO(3) equivariant network for b-tagging in high energy physics. Permutation equivariance is required since the input data is a _set_ of track particle features, while SO(3) equivariance is required since these features are geometrically scalars or vectors whose rotational transformation law should be respected. The basic architecture is a particle flow network (PFN), which is based on deep sets. It consists of 1) a first SO(3)-equivariant subnetwork, applied to each track particle individually, 2) permutation invariant pooling via summation, 3) a second SO(3)-equivariant subnetwork, and 4) a final layer which extracts scalars. Internally, the network operates on scalars, vectors, and Cartesian 2-tensors. The first two are irreps of order $0$ and $1$, respectively, the third one is reducible (it would in principle decompose according to $1\\otimes 1 \\cong 0\\oplus 1\\oplus 2$).\n\nThe equivariant subnetworks employ a range of different SO(3)-equivariant mappings:\n- Affine layers, i.e. linear maps followed by bias summation. The linear maps are made equivariant by broadcasting weights over the full representation dimension. Biases are only summed to scalars and 2-tensors since equivariant bias summation is impossible for vectors.\n- Linear layers that are applied in local frames that are aligned along each individual jet's momentum axis $j$. This alignment is only specified up to rotations around the momentum axis, which is addressed by making the layers $\\mathrm{SO}(2)\\_j$ _gauge_-equivariant (the subscript labels the axis along which the subgroup is taken). Technically, the network seems to operate on restricted representations $\\mathrm{Res}^{\\mathrm{SO}(3)}_{\\mathrm{SO}(2)} \\rho$. Since the momentum axes $j$ are moving along with SO(3) rotations, the operation is as a whole still SO(3)-equivariant. Two explicit constructions of such SO(2) equivariant maps for vectors and 2-tensors are proposed. Due to the restricted equivariance requirement, these maps are less constrained than fully SO(3)-equivariant layers.\n- Nine different _bilinear_ maps which map pairs of features of different types again to scalars, vectors and 2-tensors. This step allows (in contrast to the others) to transition between different representations.\n- SO(3)-equivariant nonlinearities. For scalars, the authors use conventional ReLUs, while the nonlinearities for vectors and 2-tensors are acting on their norm (a common approach).\n\nThere is a single experiment, in which the models are trained as binary classifiers for b-tagging. The full equivariant model improves significantly upon a non-equivariant baseline and a version which ablates 2-tensor features.\n\nSTRENGTHS:\nThe paper is well written and generally easy to follow. While I am not familiar with the b-tagging task and competing approaches, the empirical improvements over the baseline model seem very significant. Another strength is the use of bilinear mappings - most equivariant networks utilize only linear maps.\n\nI really liked the idea of using locally $SO(2)\\_j$ gauge equivariant operations besides fully SO(3) equivariant layers. The authors identified the additional geometric structure given by the momentum axes $j$ and addressed it appropriately.\n\nWEAKNESSES:\nMy main concern with the approach is that the linear layers are not shown to be _complete_: in principle one could derive a basis of the most general equivariant linear maps (intertwiners) between the given representations. The authors show only the sufficiency of their layers regarding equivariance, but not the necessity. Indeed, I believe that quite some maps are in fact not the most general ones\n(more details listed below).\n\nTo address this issue it is most convenient to work in the basis of irreducible representations, which the authors consciously avoid. For the following comments, note that scalars and vectors are irreps of order $0$ and $1$, while Cartesian 2-tensors are an irrep tensor product which decomposes according to $1\\otimes 1 \\cong 0\\oplus 1\\oplus 2$. Furthermore, intertwiners exist by Schur's lemma only between irreps of the same order, and are for SO(3) scaled identity matrices $\\lambda\\mathbb{I}$. For all non-trivial SO(2) irreps over $\\mathbb{R}$ the spaces of irreps are 2-dimensional and are spanned by the irrep-endomorphisms ((1,0),(0,1)) and ((0,-1),(1,0)).\n\n- The intertwiners for scalars and vectors (presented as \"broadcasting\") are complete since these are irreps. However, the broadcasting for $1\\otimes 1$ tensors is overly restrictive, and there are actually three parameters, one for each irrep in the 2-tensor. The claim that the intertwiner space for 2-tensors is one-dimensional is repeated in line 172.\n- The solutions for biases are indeed complete: biases can only be summed to trivial irreps, which exist with multiplicity 1 in scalars and 2-tensors, but not in vectors.\n- It is furthermore possible to have intertwiners between scalars or vectors and 2-tensors, since the latter contain irrep orders 0 and 1 as subspaces (trace and antisymmetric part). These solutions are not used.\n- The SO(2)-intertwiners between SO(3)-vectors in 3.1.1 seem complete, however, this is not proven but just claimed (\"The set of all SO(2)-equivariant maps is _exactly_ ... [proposed parametrization]\"). One can easily prove the completeness by observing that order 1 SO(3)-irreps decompose under restriction into the direct sum of an order zero and order 1 SO(2)-irrep, whose intertwiner spaces are 1 and 2-dimensional, respectively. The proposed parametrization has the same dimensionality.\n- The SO(2)-intertwiners between Cartesian SO(3) 2-tensors are again not proven to be complete (\"[equivariance] is satisfied when ... [proposed parametrization]\" just claims sufficiency). Going to the irrep basis shows again that there are more possible solutions.\n- In the case of bilinear maps, there are again some missing operations, e.g. combinations of two 2-tensors that result in a vector. All possible solutions follow directly from Clebsch-Gordan decompositions, which are well known for SO(3).\n- An alternative approach to TReLU would be to apply three independent nonlinearities to the irreps contained in the 2-tensor. The equivariance of TReLU is not explicitly shown (this might be quite trivial to show).\n\nAs mentioned above, addressing these concerns would be easiest by switching to the irrep basis. As this would require a major revision, I am not sure whether this is the right way forward for this submission, or should rather be addressed in follow-up work. An alternative way to address these concerns would be to explicitly discuss completeness of intertwiner bases in general and prove it for each operation in which it holds. The equivariance of operations like TReLU or Eq. (4) should also be proven.\n\nAnother issue is that it is not well explained how the overall network remains SO(3)-equivariant despite intermediate operations only being SO(2)-equivariant. This is one of the most interesting aspects of the paper and deserves more attention.\n\nIt should also be discussed how this relates to other _gauge equivariant_ / _coordinate independent_ networks. The alignment of frames along the $j$-axis with remaining SO(2) ambiguity seems very similar to the SO(2)-structure (SO(2)-bundle of frames) considered by (Weiler et al. 2021), specifically their figure 53 (right).\n\nThe transition from SO(3) to SO(2) features is currently not sufficiently explained. I believe that the authors assume the restriction function $\\mathrm{Res}^{\\mathrm{SO}(3)}\\_{\\mathrm{SO}(2)}$ - please clarify this!\n\nThe group actions and the domains and codomains of A and L in sections 3.1.1 and 3.1.2 are nowhere defined. One can read them off between the lines, but they should be stated more clearly.\n\nI am somewhat worried about the extent of the experiments. It would be nice to have a more thorough empirical analysis, for instance giving more ablations, showing training curves, investigating whether the equivariance is exact or due to numerical errors only approximate. Not all ablations discussed in the text of section 5.2 are shown in the table.\n\nFinally, I wondered about the input features of the baseline, are they the same as for vector and tensor PFN? There should really be two baselines, with either set of features, to clarify that the improvement is really coming from the architectural changes instead of the input alone.\nHow exactly are the different models made comparable? Do they have the same number of parameters or the same computational cost to ensure a fair comparison?\n\nQUESTIONS:\nMost of my questions and suggestions are given in the \"weaknesses\" section. In addition I wondered about the following points:\n- How can the SO(2)-equivariant layers be applied to the summed features, i.e. in the subnetwork $F$? As I understand, each jet comes with its own momentum axis $j$, but which axis is used after the summation?\n- Why are the 2F features in the bilinear layers split in two groups? Other models consider all possible pairings of features (which scales quadratically instead of linearly in F).\n- Is there a specific reasoning behind the norm-capping of VReLU and TReLU? Prior work like (Worrall et al., 2017) or (Weiler&Cesa, 2019) used ReLU-based \"norm-nonlinearities\" which do not saturate - this seems more intuitive and closer to the usual behavior of ReLUs to me.\n- In how far is the saturation of VReLU and TReLU required to stabilize the training? Could you analyze or substantiate this claim with experiments?\n\nLIMITATIONS:\nThe main limitation is in my opinion that the solutions are incomplete - which should be addressed more clearly. Other limitations, like the limitations of the bilinear ops due to not using Clebsch-Gordan decompositions or the lack of investigating the space of tensor nonlinearities are  adequately addressed.\n\nNegative societal impacts are not to be expected.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a method to handle complex geometric structured data, specifically SO(3) representations, through SO(3)-equivariance and judicious symmetry breaking. The technique improves computational efficiency and enhances the performance of a network operating on these representations. When applied to b-tagging, a High Energy Physics classification problem, it yielded a 2.7x improvement in rejection score over conventional methods.\n\nSTRENGTHS:\n1. It makes sense in machine learning to explore more efficient representation for data with symmetric structures. It is particularly important in many scientific fields such as HEP and material.\n2. Using the proposed method, it shows a significant performance improvement compared with the baseline method.\n\nWEAKNESSES:\n1. Some explanation in the paper is difficult to be followed by ML researchers. This work deeply involves the task of B-jet identification at LHC experiments, but I’m not sure if these are sufficiently interesting for the ML and AI community.\n2. In the related works, the authors mentioned that there had been numerous prior works on SO(3) equivariant models, but in the experiment only PFN with simulated datasets is implemented for comparison.  The numerical evidence in this paper looks not sufficient.\n\nQUESTIONS:\nSome comments:\n1. It would greatly enhance the clarity of the paper if the authors provided an initial explanation of the concept of SO(3) symmetry. While the authors may be well-versed in this notion, it is essential to remember that many readers within the machine learning community may not be familiar with it. By providing a concise and accessible explanation, the authors can bridge the knowledge gap and ensure that readers grasp the significance of SO(3) symmetry in the context of their work.\n\n2. In the introduction section, the authors make a claim that the proposed method is simpler compared to existing approaches. To strengthen this claim, it would be beneficial for the authors to support it with numerical or theoretical evidence. By presenting concrete results or theoretical analyses, the authors can substantiate their assertion and enhance the persuasiveness of their argument. This evidence will enable readers, including myself, to be convinced of the method's simplicity and its potential advantages over existing techniques.\n\n3. The design of the new layers appears to be somewhat ad-hoc. While the paper demonstrates the equivariance of these layers, it would be valuable to explore whether alternative designs can achieve the same goal and what differentiates them. By discussing potential alternative designs or comparing the proposed layers with existing approaches, the authors can provide a more comprehensive understanding of the design choices made. This analysis will strengthen the paper's contribution by highlighting the unique aspects of the proposed layers and providing insights into their advantages over other potential designs.\n\n4. If the main contribution of the paper lies in the introduction of the new layers, it would be beneficial to provide additional numerical evidence using different datasets. This would further validate the effectiveness of the proposed layers across a range of scenarios and reinforce their potential applicability beyond specific domains. Alternatively, if the focus of the paper is primarily on the performance improvement in the context of the LHC experiment, it would be essential to demonstrate the significance and relevance of this task. By providing additional context, motivation, and potentially exploring the broader implications of the performance improvement, the paper can better engage readers and establish the attractiveness of the addressed problem.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a SO(3)-equivariant network operating on scalars, vectors and 2-tensors. They identify the corresponding equivariant linear layers and come up with a mixing strategy to mix different representations. Furthermore, SO(2)-equivariant linear layers are proposed to allow the scenario that SO(3) symmetry breaks into a subgroup SO(2) axial symmetry along a specific axis $\\hat{j}$.  Finally, they demonstrate how it is appllied to b-tagging in High Energy Physics (HEP), where the data is rotational symmetric w.r.t. jet axis $\\hat{j}$.\n\nSTRENGTHS:\n- The authors provide a good and simple implementation of SO(3) equivariant networks on scalars, vectors and 2-tensors. The weight matrix is carefully designed to preserve the symmetry. The way to mix different representations is also intuitive and easy to understand. \n- The discussion on axial symmetry is well motivated and easy to follow, and the analysis is solid and sound.\n\nWEAKNESSES:\n- [Novelty] The idea of tensor-product-based representations is not new [1]. The main method (except the SO(2) part) looks a simple variation and the technique involved is pretty standard. Add discussion and comparison with existing tensor-product-representation-based methods could make the work more solid.\n- [Evaluation] To show the effectiveness of mixing and SO(2) linear layers, I think it is better to put more intermidiate results (e.g., w/ and w/o SO(2) linear layers) in the main table. \n- [Minor issues]: Eq. (1, 2, 3) use Einstein summation without declaration, which may cause confusion to readers without physics background. Line 168 should be \"isotropic linear neuron of Eq. (2)\" instead of \"Eq. (1)\".\n\n[1] Finkelshtein, Ben, et al. \"A simple and universal rotation equivariant point-cloud network.\" _Topological, Algebraic and Geometric Learning Workshops 2022_. PMLR, 2022.\n\nQUESTIONS:\n- How to understand Line 54 \"We show that this kind of equivariant neuron is generally only possible with the introduction of order-2 tensor representations\"? For me it seems everything should work well even if we only use scalars and vectors (like VectorPFN).\n\nLIMITATIONS:\nLimitations are not included in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents a lightweight architecture based on scalars, vectors, and tensors for learning in 3D, subject to $SO(2)$-equivariance in a known direction that varies by sample. They are motivated by the jet-tagging problem in high energy physics, in which a given batch is equivariant with respect to a certain axis (which may vary between batches), and test their framework on this problem.\n\nSTRENGTHS:\nRestricting to cross products and matrix products is new relative to previous work, which tends to focus more on the expressivity via irreducible representations of higher orders. The jet HEP dataset is also not commonly used in similar papers, and could provide a useful dataset for future papers. The proposed operations do indeed seem to be equivariant, and they are written out explicitly.\n\nWEAKNESSES:\n1. The proposed architecture is a simple restriction of many other existing architectures. The “tensor bilinear layer”, is a special case of the more general CG product that is now standard practice in other architectures, and it is not clear what benefits it has over a more general CG architecture. The nonlinearity of scaling by the vector or tensor norm is also not new: it is subsumed by the nonlinearity of applying an arbitrary function to the norm, and then scaling the vector or tensor by this value, which similarly was widespread in foundational works such as Tensor Field Networks (Thomas et al 2018) and subsequent works.\n2. It does not seem like this architecture is particularly expressive, due to the use of low order features and the simple nonlinearity. (Note that prior work on the universality of point cloud architectures, and equivariant architectures more generally, usually requires making statements about polynomial approximation, where higher order tensor products are required to approximate higher degree polynomials — see e.g. Lorentz nets, Bogatskiy et al 2020, or Dym et al 2020 on the universality of point cloud architectures. Therefore, it seems to me that these simpler layers will likely have worse approximation properties.)\n3. Based on the paper’s description (but the authors can correct if this is not the case), the motivating problem is really SO(2)-equivariant about a known but sample-dependent axis. Therefore, the discussion of the SO(3) CG product and other SO(3) architectures is somewhat misleading/confusing. It is also not explained clearly how using this paper’s SO(2)-equivariant architecture compares to the standard approaches in HEP based on the “transverse and longitudinal projections”.\n4. The baselines and experiments are not sufficiently developed. For example, the only baseline is a DeepSet-style permutation invariant architecture, presumably on the raw 3D coordinates. However, the authors should compare to approaches using the apparently more standard “transverse and longitudinal projections”, as well as to SO(2)- or SO(3)-equivariant baselines on the raw coordinates. The paper claims that using a physically intuitive restriction of the space of possible operations is beneficial, but does not demonstrate this with an ablation study or by using e.g. a full CG product.\n\nQUESTIONS:\n1. What exactly are the “transverse and longitudinal” projections that other methods employ on this data, and why do other methods use them? Is returning to the full 3D coordinates with SO(2)-equivariance likely to be more efficient than reducing to these coordinates in the first place, from a sample complexity perspective? \n2. Line 113 says “these methods all essentially rely on reducing the problem to vectors in a 2-dimensional plane”. Can the authors elaborate on what this means? If these methods use the jet direction to canonicalize the coordinate system, for example, they could easily use frame averaging (Puny et al 2021) to obtain equivariant outputs — but I am not sure if this is what they are doing or not.\n3. What’s the difference between this approach and using a standard method for SO(2)-equivariance, such as circular convolutions about the direction j?\n4. Line 48 claims that the tensor bilinear layer consists of “commonly known and physically intuitive operations”. Can the authors elaborate on the problem-specific physical intuitions that come with these operations? Is there some intuitive benefit to using physically intuitive operations as opposed to a more general framework relying on the full CG product?\n5. Does the HEP problem enjoy translational symmetry as well?\n\nSuggestions:\n1. A diagram of the experimental set-up and the individual particle directions would make the paper and its motivation much more clear.\n2. Order 2 tensors can be defined somewhere early in the paper.\n3. The abstract refers to this method as “judicious symmetry breaking”. I think this is a misleading phrase, when really it is just that the symmetry group is a subgroup of SO(3) (rather than all of SO(3)). In other literature, symmetry-breaking refers to situations more complex than simply having (consistency) a smaller group than SO(3); for example, addressing the ambiguity between 6s and 9s in the MNIST digit dataset under rotation.\n\nLIMITATIONS:\nThere is no potential negative societal impact. The paper does not address a potential lack of expressivity of the architecture.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBefore starting, I must mention that I am not a physicist, and therefore, I have focused on the machine learning aspects of the paper.\n\n### Summary\n\nThis paper proposes the use of SO(3) equivariant neural networks for B-tagging. The method proposed builds upon Deep Sets and provides a quite constrained formulation of SO(3) equivariance. The paper further indicates that breaking SO(3) equivariance locally can be used to get better representations for the task at hand, while maintaining global equivariance to that group.\n\nSTRENGTHS:\nThe paper is presented very clearly, well-structured, and it is in general easy to follow and understand.\n\nWEAKNESSES:\n### Concerns\n* My main concern is that, to the best of my knowledge, it is not possible to obtain a certain equivariance without all networks being equivariant to that group. In particular, I do not understand how global SO(3) equivariance can be locally broken into SO(2)_j equivariance while preserving global SO(3) equivariance. In the best of my understanding, as shown in roughly all the papers on group equivariance, in order to have a network be equivariant to a certain transformation, all layers need to respect that equivariance. I think that clarifying how this is possible is *crucial* for the paper. It is important to note that the paper simply states this and does not provide proofs or analyses regarding this statement. \n* I am very concerned with regard to the expressivity of the proposed algorithm. For instance, it has been shown in several works that equivariance can be obtained in expressive ways that do not have such hard restrictions as having biases be equal to zero –note that several similar very constraining restrictions are also defined for each of the mappings in the networks–, e.g., E(n)-equivariant Steerable CNNs, among many others. From what I understand, this work builds mostly upon the Deep Sets literature. However, this is by far not the most general way to obtain equivariance to a certain symmetry. I believe that the authors should at least state this clearly in the paper.\n* The paper performs multiple ablations that are not found in the submission other than by conclusive statements. For instance, in line 322, the authors state: “Finally, we note that neither family of models performs even as well as the baseline, when no bilinear operations are allowed”. I believe that clearly showing the results of these ablations will strengthen the contributions of the paper. In general, such conclusive statements in their own are often vague and not very informative.\n* In the final part of the conclusion it is stated that “it should be also possible to use these models for creating equivariant Graph and attention based networks”. Given that these families of networks –as stated earlier– are much more general than Deep Sets for equivariant formulations, I believe that this statement is not really easy to accomplish in practice.\n\nQUESTIONS:\n### Comments\n* Line 34 - 35. What about translations? \n\n### References\nThis paper misses several references in several parts of the paper. For instance:\n* The second paragraph should include citations. \n* Line 52. There are several works that perform equivariance relaxations.\n* This paper is very similar to recent papers on Clifford Algebras, e.g., Clifford Group Equivariant Neural Networks. A discussion wrt these methods should be included.\n* This paper also seems very similar to Vector Neurons. A discussion wrt such methods should be included.\n* I have seen several nonlinearity formulations that look very much like the VReLU / TReLU, e.g., in Tensor field networks, Clifford Group Equivariant Neural Networks, etc. Please cite and discuss how the proposed nonlinearities are better.\n\n### Introducing concepts\nThis paper does not introduce several things properly. For instance:\n- Equation 4 falls completely out of the blue. To me, it is not clearly how the authors have reasoned to arrive at this. \n- It is never explained that \\alpha and \\beta are in Figure 1.\n- Why is x in |x| not bold?\n\nLIMITATIONS:\nSee previous responses.\n\n### Conclusion\nIn conclusion, I believe that the paper needs some work before I am able to support acceptance. I believe that there are several factors that need to be clarified, .e.g, how to get SO(3) equivariance with only SO(2)_j layers, for this to be an strong submission.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-27"}