On training and application ofequivariant neural networks.
The University of Tokyo & RIKEN AIP
Mukuta Yusuke

1Invariance in pattern recognition• The input signal for recognition often has invarianceand equivariance with respect to transformations.
input image

horizontal flipping

rotation

Image rotationand flip

‘dog’

Co. A

Co. C

Co. D

Predictor 𝑓

𝐴𝑡+1
𝐵𝑡+1𝐶𝑡+1𝐷𝑡+1

equivalent

value

Permutation ofmulti-variate
time series

Co. B

𝐴𝑡
𝐵𝑡
𝐶𝑡
𝐷𝑡

time

𝐵𝑡
𝐶𝑡
𝐴𝑡𝐷𝑡

Predictor 𝑓

𝐵𝑡+1𝐶𝑡+1𝐴𝑡+1
𝐷𝑡+12

Equivariant Neural Network• Equivariant Neural Network constructs the network bycomposing multiple equivariant layers.
• Given transformation 𝑇, equivariant layer 𝑓 is a layerthat commutes with 𝑇 such that𝑓 𝑇𝑖𝑛 𝑥= 𝑇𝑜𝑢𝑡 (𝑓 𝑥 ).

• We apply pooling with respect to transformations toobtain invariant prediction.
Outputequivariant
prediction

Input 𝑥Equivariantlayer 𝑓 1… Equivariantlayer 𝑓 𝑙
the effect of transformation
is preserved3

Group Equivariant Convolutional
Networks [Cohen & Welling, 2016]• Apply all the transformations to the convolutional
filter and then apply convolution to the image.
• Then image transformations results in thepermutation of the filter response (equivariant).
convolution∗
Input image

∗

…

∗…
4

Today’s topic• Self-supervised learning on equivariant neural networks
https://arxiv.org/pdf/2303.04427.pdf• Time series prediction considering hierarchicalpermutation equivariance https://arxiv.org/pdf/2305.08073.pdf

5Self-Supervised Learning• First we pretrain the model with user defined pretext
task that does not use image label.
• Then we use the feature extractor for the target task.
unlabeled images

extract the knowledge
independent on category information

featureextractor

either target dataset
target dataset
and different
dataset is ok

pretext task

classifier

𝑙𝑝𝑟𝑒𝑡𝑒𝑥𝑡

initialize(freeze or finetune)

featureextractor

classifier

target task
𝑙𝑡𝑎𝑠𝑘

6

Self-Supervised Learning Methods• Hand-crafted tasks• Train the model to solve hand-crafted ill-posed problem.
• We assume that the feature extractor learn good image
prior while trying to solve the problem.

• Contrastive learning• Apply data augmentation and trains the model to makethe augmented images from the same image close.

7

Context prediction [Doersch et al.,2015]
• Predict the spatial relationship between two image
patches.

8Model• Cast the context prediction as 8 category prediction
from two image sub-region.

architecture

9

Contrastive Learning• Learn feature so that two augmented image arecloser than the other images by contrastive loss
exp

𝑥𝑡 𝑥𝑝
𝜏

𝑥𝑡 𝑥𝑛
𝑥𝑡 𝑥𝑝
exp 𝜏 +σ𝑘 exp 𝜏 𝑘

𝑥1𝑛

𝑥2𝑛

𝑥3𝑛

𝑥4𝑛𝑥5𝑛augmentation
close𝑥𝑝

𝑥

distant

10Motivation• Combine the idea of• Exploiting the prior knowledge as group equivariant
architecture.
• Exploiting the prior knowledge as pretext task.

11Difficulty• The function learned by equivariant neuralnetworks 𝑓𝑁𝑁 is restricted to equivariant such that𝑓𝑁𝑁 𝑇𝑖𝑛 𝑥 = 𝑇𝑜𝑢𝑡 (𝑓𝑁𝑁 𝑥 ).
• We cannot learn the task if the pretext label
violates this equivariance.
input image 𝑥transformed image 𝑇(𝑥)Needs to be consistent𝑝𝑟𝑒𝑡𝑒𝑥𝑡
Pretext label 𝑦𝑥

𝑝𝑟𝑒𝑡𝑒𝑥𝑡

Pretext label 𝑦𝑇(𝑥)12Proposed: equivariant pretextlabel• Restrict the pretext label space so that satisfies
𝑇out 𝑔 𝑦 pretext 𝑥 = 𝑦 pretext 𝑇in (𝑔)(𝑥) .
Input image 𝑥
assignlabel

𝑇in (𝑔)Transformed
image 𝑇in (𝑔)(𝑥)commuteassignlabel

Pretext label𝑦 pretext (𝑥)𝑇out (𝑔)Pretext label
𝑇out 𝑔 𝑦 pretext 𝑥= 𝑦 pretext 𝑇in (𝑔)(𝑥)13

Proposed: invariant contrastiveloss• Invariant contrastive loss is the loss function thatsatisfies
Loss function 𝑙equal

Loss function 𝑙𝑓𝜃 (𝑥1 )𝑓𝜃 (𝑥𝑚 )𝑓𝜃 (𝑥𝐵 )𝑓𝜃 (𝑥1 )𝑇out (𝑔)𝑓𝜃 (𝑥𝑚 )𝑓𝜃 (𝑥𝐵 )

𝑓𝜃

𝑓𝜃

𝑓𝜃

𝑓𝜃

𝑓𝜃

𝑓𝜃

⋯
𝑥1

⋯
𝑥𝑚

⋯
𝑥𝐵

𝑥1

𝑇in (𝑔)(𝑥𝑚 )

⋯𝑥𝐵

14

Context prediction [Doersch et al.,2015]
• Predict the spatial relationship between two imagepatches.

1 2 345
6 7 8
𝑋=,,𝑌 = 415Equivariant Context Prediction• The label space ℝ8 of context prediction tasksatisfies 90∘ rotation equivariance.

1 2 3
4
5
6 7 8
𝑋=,
𝑝𝑟𝑒𝑡𝑒𝑥𝑡𝑦𝑥

1 2 3456 7 890∘ rot,𝑌 = 4= (0,0,0,1,0,0,0,0)𝑋=𝑝𝑟𝑒𝑡𝑒𝑥𝑡𝑦𝑇(𝑥),,𝑌 = 7

= (0,0,0,0,0,0,1,0)

16

Invariant Contrastive Learning• We average the output feature asexp
exp

𝑡 11σ𝑔∈𝐺 𝑇𝑜𝑢𝑡 𝑔 𝑥σ𝑇𝑔 𝑥𝑝|𝐺||𝐺| 𝑔∈𝐺 𝑜𝑢𝑡𝜏𝑡 11σ𝑔∈𝐺 𝑇𝑜𝑢𝑡 𝑔 𝑥σ𝑇𝑔 𝑥𝑝|𝐺| 𝑔∈𝐺 𝑜𝑢𝑡𝐺𝜏+σ𝑘 exp𝑡 11σ𝑔∈𝐺 𝑇𝑜𝑢𝑡 𝑔 𝑥
σ𝑇𝑔 𝑥𝑛𝑘|𝐺| 𝑔∈𝐺 𝑜𝑢𝑡𝐺𝜏to satisfy

𝑥1𝑛𝑥2𝑛𝑥3𝑛𝑥4𝑛𝑥5𝑛

augmentation
close𝑥𝑝𝑥

distant

17

Experiment• Evaluation:• Pretrain on ImageNet (1,300,000 images, 1,000 labels)• Apply linear classifier on top of the pretrained model.

• Architecture: ResNet50• Compare• Standard non-equivariant model• Group equivariant model with the proposed loss
• Group equivariant model with standard non-equivariant loss

18

Result19

Multi variate time series prediction• Predict the future of multi-variate time series• The predictor should be equivariant to thepermutation of the time series.

value

Co. A

Co. B

Co. C

Co. D𝐴𝑡
𝐵𝑡
𝐶𝑡
𝐷𝑡

𝐴𝑡+1
𝐵𝑡+1
Predictor 𝑓
𝐶𝑡+1
𝐷𝑡+1
equivalent

𝐵𝑡
𝐶𝑡
𝐴𝑡
𝐷𝑡

𝐵𝑡+1
𝐶𝑡+1
𝐴𝑡+120𝐷
𝑡+1

Predictor 𝑓

Self-attention• Calculate the output by the weighted average ofthe input, whose weights are calculated by thesimilarity of the inputs.
• We can preserve permutation equivariance byapplying self-attention between the time series.

21[1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.

Hierarchical permutationequivariance• We consider the case that time series arehierarchically grouped by such as sector, class.
• We want to restrict the equivariance to thepermutation that considers hierarchy.

2222

Feature extractor• Use 3D self-attention within class 𝑆𝑖 , between
classes 𝐶, time 𝑇.

Within-class
information
Input time series 𝑋Series 𝑆Time 𝑇𝑆𝑖 self-attention

Add & 𝑇
self-attention EquivariantfeatureBetween-class
information

Feature 𝐷
𝐶 self-attention23

𝑆𝑖 self-attention• Split the time series according to the class andapply self attention within class.
Class 1

Input time series 𝑋
Time 𝑇Series 𝑆
split

Class 2

Class 1

Selfattention

Class 2merge

Feature 𝐷
Class 𝑀Class 𝑀24

𝐶 self-attention• Summarize the features of each class and then calculate
the self-attention between the summarized features.
Class 1

Class 1

Class 1attention

nput time series 𝑋
Time 𝑇Series 𝑆
split

Class 2

Class 2

Class 2poolmerge

Feature 𝐷
Class 𝑀Class 𝑀Class 𝑀25Experiment• NBA: Trajectory of players and the ball in the basket
game.
• In: 40 steps, Out: 10 steps• 11 agents, 3 classes (ball, team A, team B)• Train: 80,000, Validation: 48,299, Test: 13,464• Evaluate the prediction accuracy by reducing the number ofteam A/B players at test time.

26

Result

Left: without class information, Right: with class information
0~4 indicates the number of reduced players from each team27

Conclusion• We introduced two recent works relating equivariant
neural networks.
• Propose the idea of equivariant pretext labels and invariant
contrastive loss to combine equivariant neural networks andself-supervised learning https://arxiv.org/pdf/2303.04427.pdf• Propose the multi-variate time series prediction method
considers hierarchical permutation equivariancehttps://arxiv.org/pdf/2305.08073.pdf

28