11institutetext: CERN – European Organization for Nuclear Research
22institutetext: Faculty of Physics, Warsaw University of Technology
33institutetext: Faculty of Electrical Engineering, Warsaw University of
Technology 44institutetext: Faculty of Electronics and Information Technology,
Warsaw University of Technology 55institutetext: IDEAS NCBR

# Particle identification with machine learning in ALICE Run 3

Maja Karwowska 1122<EMAIL_ADDRESS>Monika Jakubowska 33 Åukasz Graczykowski 22
Kamil Deja 4455 MiÅosz Kasak 44

###### Abstract

The main focus of the ALICE experiment, quark–gluon plasma measurements,
requires accurate particle identification (PID). The ALICE sub-detectors allow
identifying particles over a broad momentum interval ranging from about 100
MeV/$c$ up to 20 GeV/$c$. However, a machine learning (ML) model can explore
more detector information. During LHC Run 2, preliminary studies with Random
Forests obtained much higher efficiencies and purities for selected particles
than standard techniques.

For Run 3, we investigate Domain Adaptation Neural Networks that account for
the discrepancies between the Monte Carlo simulations and the experimental
data. Preliminary studies show that domain adaptation improves particle
classification. Moreover, the solution is extended with Feature Set Embedding
and attention to give the network more flexibility to train on data with
various sets of detector signals. PID ML is already integrated with the ALICE
Run 3 Analysis Framework. Preliminary results for the PID of selected particle
species, including real-world analyzes, are discussed as well as the possible
optimizations.

##  1 Introduction

ALICE (A Large Ion Collider Experiment) ALICE:2008ngc  is one of the four big
detectors located at the Large Hadron Collider (LHC) at CERN. ALICE studies
the properties of quark–gluon plasma (QGP), a hot and dense state of matter,
and the strong force that holds quarks together inside hadrons ALICE:2022wpn .
Detailed analysis of QGP requires accurate particle identification (PID),
i.e., the ability to discriminate between different particle species produced
during the collision. High PID precision distinguishes ALICE from other LHC
experiments and allows for selecting a subset of particles required for
specific analysis.

The ALICE experiment is composed of several sub-detectors, some of which
measure particle properties that can be used for identification. Figure 1
presents a scheme of the detector in Run 1 and Run 2 LHC data-taking periods.

Figure 1: Components of the ALICE detector in its Run 2 configuration
botta:2017a .

The three detectors particularly useful for PID are: Time Projection Chamber
(TPC), Time-of-Flight (TOF), Transition Radiation Detector (TRD). TPC is one
of the most important ALICE detectors as it records 3D information of the
trajectory of the particles, and particle ionization energy loss essential for
particle identification. The Time-of-Flight detector measures particle travel
time from the collision vertex to the detector, from which particle velocity
and mass are calculated. TRD records transition radiation, that is, the
emission of a few photons by traversing electrons, which helps in
distinguishing electrons from other charged particles.

With the signals recorded by the detectors described above, particles are
chosen using selection criteria compared with theoretical calculations. The
traditional method compares the number of standard deviations from the
expected value for all detector signals. Particles falling outside the
selection region are rejected. However, when the characteristics of different
particle species overlap, combining information from multiple detectors
becomes difficult. Choosing selection regions by trial and error is less
effective in this case, lowering PID efficiency and limiting the statistical
significance of the final analysis.

These shortcomings can be addressed with more advanced classification methods
such as Bayesian models or neural-network-based (NN) approaches. The Bayesian
approach ALICE:2016zzl  is available in the new ALICE software framework O2,
but its flexibility is limited. For these reasons, we studied possible
improvements employing neural networks for PID.

##  2 Machine learning for PID

Particle identification is a standard classification problem for a machine
learning algorithm. The inputs to the model are various track parameters, and
from them, machine learning (ML) can utilize more information and learn more
complex relationships than what a human analyzer can observe. Thanks to the
present popularity of neural networks, there are many software libraries that
enable fast implementation of a given neural network model.

We start with a simple feed-forward neural network, which is trained and
applied on Monte Carlo simulated data. The network acts as a binary
classifier, outputting predictions for a given particle species. The single-
value output is normalized to the range $(0,1)$ by applying the logistic
function $f(x)=\frac{1}{1+e^{-x}}$. The classifier result approximates the
probability of the example corresponding to a specific particle type.
Additionally, the user can decide which detector information they want to use.
This has an impact on the size of the input vector, and requires the
development of different models for different input sizes. The networks for
each (anti)particle species and detector setup are trained independently.
Initial results are promising: a simple network using combined information
from TPC and TOF detectors can improve efficiency compared to traditional
method while not losing purity.

One must also consider the limitations that an ML model will have. Firstly,
the performance greatly depends on training data quality, so also on the
results of algorithms used for simulation, detector calibration, and track
reconstruction. Physics studies are usually performed on well-calibrated,
high-quality data, which further predisposes the use of ML. However, by
design, such data still contains many incomplete samples. The problem and our
solution are described in detail in Section 3. Another particularity is the
domain shift between the distributions of variables in a training dataset
(simulation) and a corresponding test/inference dataset (real experiment
data). This is mitigated with the domain adaptation technique explained in
Section 4.

Next, it is more difficult to obtain systematic uncertainties from ML results.
A few approximate methods exist. For example, in ghosh2021uncertainty , the
authors propose a model where classifiers are aware of systematic
uncertainties of input parameters. This approach improves the model’s
sensitivity. In englert2019machine , the authors use adversarial training for
this purpose.

###  2.1 Integration of PID ML with the O2 framework

Applying machine learning models in an analysis brings the technical challenge
of crossing different programming languages. Various ML frameworks used in
ALICE are implemented in Python, PID ML not being an exception. However, the
ALICE analysis framework, O2 o2physics , is written in C++ and heavily
utilizes the ROOT BRUN199781  library. Our solution makes use of the ONNX
(Open Neural Network Exchange) standard onnx , which defines a common file
format for storing machine learning models developed in various frameworks
such as Tensorflow tensorflow2015-whitepaper  and PyTorch pytorch . The
ONNXRuntime onnxruntime  can then be used for ONNX model inference in C++.

Figure 2 depicts the current workflow of PID ML in O2. The input data consists
of reconstructed collisions and tracks in AO2D files, which are stored on disk
in ROOT file format. A PIDML producer task takes the initial AO2D files,
applies a few rough preselections to exclude meaningless tracks, and produces
skimmed data that contains only track properties used by a neural network. The
data is then used for model training. Trained models are stored in the ONNX
format in the Condition and Calibration Data Base (CCDB). The database is
accessible by analyses running on the worldwide LHC computing GRID. An
analyzer can apply an ML model for particle identification by adding an
instance of the PidOnnxModel class or, in more complex use cases, the
PidOnnxModelInterface class. The latter enables handy management of multiple
ML models, each with different target particle species, detector setup, and
acceptance threshold.

The same solution, with ONNXRuntime and a C++ class to wrap management of a
single ONNX model, was applied in other machine learning projects in O2. A
common interface that would further unify various implementations of the ONNX
integration is being developed based on PidOnnxModel, and it is already used
for supporting calculations for TPC PID and selection of $\Lambda_{c}^{+}$ in
the $\rm{pK}^{-}\pi^{+}$ decay channel.

Figure 2: Scheme of the current ML particle identification workflow in O2.
More information about the workflow topology in the O2 Analysis Framework can
be found in Alkin:2021mfo .

##  3 Dealing with incomplete data

The number of inputs to the network can vary not only because of user
preferences but also at random, because of independent measurements of
particle properties by various detectors. A particle may be measured by a
subset of detectors while not being recorded by others. For instance, a
detector may malfunction or be switched off, or the particle can have
characteristics (i.e., too low transverse momentum) that do not match the
detector specification.

A detailed study was performed to find the best model for particle
identification with missing data. The simplest method is case deletion, that
is, removing incomplete samples from data. However, it prevents the
identification of samples with missing detector information, which are equally
important as the samples with a complete set of signals. On the other hand,
imputation techniques work by filling missing data with artificial values
computed as a mean or median of existing signals, or with values predicted
with a simple model such as linear regression. Both case deletion and
imputation distort the initial dataset, which can disturb predictions of the
neural network.

Instead of changing data, one can adjust the model architecture. A
straightforward solution is to use a neural network ensemble: a set of
classifiers, one per each subset of the training dataset without missing data.
In particular, in neural network reduction sharpe1995dealing , authors propose
to split the dataset into the largest possible complete subsets. Nevertheless,
this method can be computationally expensive, especially with a growing set of
attributes with missing values.

To overcome the aforementioned shortcomings, we introduce a novel method based
on the attention mechanism, similar to the method introduced for a medical
use-case in AMI-Net wang2019attention . The overview of the system is shown in
Figure 3.

Figure 3: The proposed model architecture. Layered blocks are applied
separately to each vector in a set. Single blocks are applied to their input
as a whole.

First, all data examples are encoded into a set of feature-value pairs
grangier2010feature , independently of their characteristics regarding missing
data. Each pair consists of a non-missing value in the input data and the
index of that value with one-hot encoding applied. Then, the feature-value
pairs are transformed by a neural network with a single hidden layer into
embeddings that place similar features close in the embedded space.

The Transformer’s vaswani2017attention  encoding module is applied to the set
of embedding vectors to connect the different features each vector represents
and find input patterns. For example, a measurement from a specific detector
could be used if the momentum is within a particular range. The softmax
function is applied to the attention output, a variable-size set of vectors.
An additional self-attention layer is used to merge these vectors into a
single one.

Finally, the pooled vector is processed by the simple classifier described in
Section 2.

###  3.1 Results

The data comes from a Monte Carlo simulation of proton–proton collisions at
$\sqrt{s}=13$ TeV with realistic simulation of the time evolution of the
detector conditions in the LHC Run 2 data-taking period. The simulation was
performed with Pythia8 Sjostrand:2014zea , the Geant 4 Brun:1994aa  particle
transport model, and general-purpose settings.

The dataset consists of 2 751 934 examples with track’s transverse momentum
$p_{\rm T}\geq 0.1$ GeV/$c$. 95% of examples fall into the range $[0.12,1.76]$
GeV/$c$. Each example contains 19 features, representing signals from TPC,
TRD, and TOF detectors as well as predetermined properties such as spatial
coordinates of a track reconstruction starting point ($x$, $y$, $z$, angle
$\alpha$), track momentum, type, charge, and the distances of closest approach
(DCA) of the track trajectory to the collision’s primary vertex measured in
the $xy$ plane ($d_{xy}$) and the $z$ direction ($d_{z}$). The examples are
labeled with ten different particle types. The distribution of particles and
their missing values is shown in Figure 4.

pion | kaon | proton | electron | muon
---|---|---|---|---
43.59% | 3.415% | 2.026% | 0.794% | 0.323%
antipion | antikaon | antiproton | antielectron | antimuon
43.66% | 3.288% | 1.819% | 0.762% | 0.321%

Figure 4: Left: Particle type distribution in studied dataset. Approximately
97.8% of the examples belong to the 6 most populous particle types. Right:
Missing data distribution. Over 62.8% of the examples is missing at least one
value.

For each compared method, we train identical binary classification models for
6 of the most abundant particle types: pions, protons, kaons, and their
respective antiparticles. All trained models are identical for imputation
methods and the neural network ensemble. They have three hidden layers of
sizes 64, 32 and 16, and a single output. Between the layers, we use Rectified
Linear Unit (ReLU) activation function $f(x)=\max({0,x})$. Dropout
regularization with a rate of 0.1 is applied after each activation layer. For
imputation techniques, all models process inputs of size 19, as all missing
features are imputed. For the neural net ensemble, the four networks have
input sizes 19, 17, 17, and 15, depending on the combination of missing values
from the TRD and TOF detectors (2 features per each detector). The tested
attention architecture described in Section 3 is implemented with ReLU
activation between neural network layers, and dropout regularization with a
rate of 0.1 at the output of the embedding layer and the output of each sub-
layer in the encoder. Model parameters are selected using a hyperparameter
sweep for a fair comparison of different network architectures.

Each model is trained using all complete and incomplete examples and tested in
two cases: (1) with all available test examples and (2) with the complete ones
only. For case deletion, no results are available for testing on the missing
data dataset. The models are evaluated on test-set with standard metrics –
precision (purity) and recall (efficiency). Since precision and recall are
antagonistic, also $F_{1}$ metric is considered, which is a combination of the
first two: $\rm{F}_{1}=\frac{precision\leavevmode\nobreak\
\cdot\leavevmode\nobreak\ recall}{precision\leavevmode\nobreak\
+\leavevmode\nobreak\ recall}$. The scored means and standard deviations of 4
runs with different seeds for classification of 3 particle species are shown
in Table 1. Analogous results are obtained for the respective antiparticle
species. The best-performing algorithm is highlighted in green, and the second
best – in blue.

Table 1: Classification results for the three most common particle species.

$\pi$ (pions), only complete data
---
model | precision | recall | $F_{1}$
Delete | 0.9908 $\pm$ 0.0007 | 0.9967 $\pm$ 0.0004 | 0.9937 $\pm$ 0.0001
Ensemble | 0.9911 $\pm$ 0.0004 | 0.9964 $\pm$ 0.0006 | 0.9938 $\pm$ 0.0001
Mean | 0.9885 $\pm$ 0.0009 | 0.9969 $\pm$ 0.0004 | 0.9927 $\pm$ 0.0004
Proposed | 0.9908 $\pm$ 0.0002 | 0.9964 $\pm$ 0.0003 | 0.9936 $\pm$ 0.0001
Regression | 0.9902 $\pm$ 0.0002 | 0.9949 $\pm$ 0.0014 | 0.9925 $\pm$ 0.0007

p (protons), only complete data
---
model | precision | recall | $F_{1}$
Delete | 0.9923 $\pm$ 0.0032 | 0.9963 $\pm$ 0.0005 | 0.9943 $\pm$ 0.0016
Ensemble | 0.9916 $\pm$ 0.0024 | 0.9976 $\pm$ 0.0007 | 0.9946 $\pm$ 0.0013
Mean | 0.9922 $\pm$ 0.0019 | 0.9972 $\pm$ 0.0004 | 0.9947 $\pm$ 0.0008
Proposed | 0.9928 $\pm$ 0.0010 | 0.9968 $\pm$ 0.0009 | 0.9948 $\pm$ 0.0002
Regression | 0.9910 $\pm$ 0.0009 | 0.9965 $\pm$ 0.0009 | 0.9937 $\pm$ 0.0007

K (kaons), only complete data
---
model | precision | recall | $F_{1}$
Delete | 0.9693 $\pm$ 0.0037 | 0.9698 $\pm$ 0.0026 | 0.9695 $\pm$ 0.0006
Ensemble | 0.9665 $\pm$ 0.0038 | 0.9782 $\pm$ 0.0031 | 0.9723 $\pm$ 0.0010
Mean | 0.9683 $\pm$ 0.0017 | 0.9533 $\pm$ 0.0067 | 0.9608 $\pm$ 0.0036
Proposed | 0.9603 $\pm$ 0.0098 | 0.9806 $\pm$ 0.0072 | 0.9704 $\pm$ 0.0017
Regression | 0.9427 $\pm$ 0.0098 | 0.9701 $\pm$ 0.0051 | 0.9562 $\pm$ 0.0039

$\pi$ (pions), all data
---
model | precision | recall | $F_{1}$
Ensemble | 0.9747 $\pm$ 0.0025 | 0.9946 $\pm$ 0.0021 | 0.9845 $\pm$ 0.0004
Mean | 0.9731 $\pm$ 0.0007 | 0.9952 $\pm$ 0.0007 | 0.9840 $\pm$ 0.0001
Proposed | 0.9749 $\pm$ 0.0006 | 0.9954 $\pm$ 0.0005 | 0.9850 $\pm$ 0.0002
Regression | 0.9733 $\pm$ 0.0006 | 0.9949 $\pm$ 0.0007 | 0.9840 $\pm$ 0.0004

p (protons), all data
---
model | precision | recall | $F_{1}$
Ensemble | 0.9716 $\pm$ 0.0046 | 0.9374 $\pm$ 0.0030 | 0.9542 $\pm$ 0.0012
Mean | 0.9785 $\pm$ 0.0041 | 0.9334 $\pm$ 0.0032 | 0.9554 $\pm$ 0.0006
Proposed | 0.9780 $\pm$ 0.0044 | 0.9386 $\pm$ 0.0027 | 0.9579 $\pm$ 0.0007
Regression | 0.9738 $\pm$ 0.0040 | 0.9367 $\pm$ 0.0038 | 0.9549 $\pm$ 0.0015

K (kaons), all data
---
model | precision | recall | $F_{1}$
Ensemble | 0.9118 $\pm$ 0.0200 | 0.8272 $\pm$ 0.0142 | 0.8674 $\pm$ 0.0016
Mean | 0.9083 $\pm$ 0.0171 | 0.8232 $\pm$ 0.0096 | 0.8636 $\pm$ 0.0034
Proposed | 0.9155 $\pm$ 0.0071 | 0.8368 $\pm$ 0.0082 | 0.8744 $\pm$ 0.0014
Regression | 0.9117 $\pm$ 0.0100 | 0.8178 $\pm$ 0.0021 | 0.8622 $\pm$ 0.0046

It is clear that the proposed attention architecture achieves very good scores
of each metrics, comparable with other analyzed models. It performs
particularly well with the realistic dataset containing incomplete samples. At
the same time, our model avoids the flaws of other solutions: artificial bias
in imputed and case-deleted data, and potentially larger complexity of neural
network ensemble.

##  4 Domain adaptation

In analysis, PID is used to select particles of desired types in both real
experimental data and Monte Carlo simulations. However, recorded signals in
physical detectors can differ from those produced in simulations. Standard PID
methods rely on partially automated processes for data domains alignment,
e.g., ALICE uses a tuning of simulated signals in such a way they reproduce,
on average, the collected data distributions of specific signals.

To circumvent the limitations of standard domain alignment methods, we propose
combining domain alignment with particle identification stating it as a known
problem of classification with unsupervised domain adaptation. The main idea
of this technique is to learn the discrepancies between two data domains, that
is, the labeled source domain (simulation data) and the unlabeled target one
(experimental data), and translate those to a single hyperspace. Classifiers
trained on top of features located in combined latent space should have
similar performance on both MC simulated and experimental data.

A Domain Adversarial Neural Network (DANN) applies domain adaptation in neural
networks. The architecture, presented in Figure 5, is composed of three neural
networks. The goal of the feature mapping network is to map original input
into domain invariant features. Those features serve as an input to the
standard particle classifier that outputs the particle type. Additionally, the
last model, known as domain classifier, enforces domain invariance of
extracted features through adversarial training. The training of the model is
divided into two steps. First, on top of current features from the feature
mapping network, the domain classifier is trained independently to classify
domain labels – whether data come from a real or a simulation source. Then,
the domain classifier is kept frozen so that the particle classifier and the
feature mapper can be trained jointly to predict accurate particle types while
fouling the domain classifier. With this approach, the feature mapper’s
weights are updated with a gradient from the particle classifier and reversed
gradient from the domain classifier. Training a domain-adaptation-based
classifier is more complex than the classical neural model. However, the
application performance of those two methods is similar and depends on the
complexity of the classifier and feature extractor.

The domain-adaptation technique is widely used in natural language processing
blitzer2007biographies ; glorot2011domain  and computer vision
gopalan2011domain ; fernando2013unsupervised . In the domain of high-energy
physics, the author of walter2018domain  presents how this method can improve
the quality of automatic jet tagging on real experimental data. The initial
tests with Domain Adversarial Neural Networks (DANN) ganin2016domain  show
that this technique improves the classification of particles in experimental
data.

Figure 5: Architecture of Domain Adversarial Neural Network.

##  5 Future developments and conclusions

The current priority is to test the framework in a real-world analysis on
simulation data from the present Run 3 data-taking period. Afterward, we will
extend the attention model with domain adaptation and test it on the new real
data. Finally, we will be able to achieve regular production of models for Run
3.

In order to estimate statistical uncertainties, we plan to include tests
analogous to jha2019impact  measuring how training dataset selection might
affect the model’s performance. We intend to follow the developments described
in gal2016dropout , where the authors show that dropout, a standard method for
reducing overfitting in neural networks, might approximate Bayesian
uncertainty in deep Gaussian processes. A method for calibrating Bayesian
uncertainties is also presented in kuleshov2018accurate .

## Acknowledgments

The research was funded by POB HEP of Warsaw University of Technology within
the Excellence Initiative: Research University (IDUB) program, by the Polish
Ministry of Education and Science in the framework of International Co-
financed Projects under agreements no. 2022/WK/01, 5236/CERN/2022/0 and by
National Science Centre under agreement UMO-2021/43/D/ST2/02214.

## References

  * (1) K. Aamodt et al. (ALICE), Journal of Instrumentation 3, S08002 (2008)
  * (2) S. Acharya et al. (ALICE), arXiv:2211.04384 [nucl-ex] (2022), 2211.04384
  * (3) E. Botta (ALICE), Proceeding of the Fifth Annual Conference on Large Hadron Collider Physics (2017), 1709.00288
  * (4) J. Adam et al. (ALICE), European Physics Journal Plus 131, 168 (2016), 1602.01392
  * (5) A. Ghosh, B. Nachman, D. Whiteson, arXiv preprint arXiv:2105.08742 (2021)
  * (6) C. Englert et al., The European Physical Journal C 79, 1 (2019)
  * (7) AliceO2Group, _O2 analysis framework documentation_ , https://aliceo2group.github.io/analysis-framework/
  * (8) R. Brun, F. Rademakers, Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 389, 81 (1997)
  * (9) ONNX Community, _ONNX_ , ONNX Community, 2021, https://onnx.ai/
  * (10) M. Abadi et al., _TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems_ (2015), software available from tensorflow.org
  * (11) A. Paszke et al., in _Advances in Neural Information Processing Systems_ (2019), Vol. 32, pp. 8024–8035
  * (12) ONNXRuntime Community, _ONNXRuntime_ , ONNXRuntime Community, 2021, https://onnxruntime.ai/
  * (13) A. Alkin, G. Eulisse, J.F. Grosse-Oetringhaus, P. Hristov, M. Kabus, EPJ Web Conf. 251, 03063 (2021)
  * (14) P.K. Sharpe, R.J. Solly, Neural Computing and Applications 3, 73 (1995)
  * (15) Z. Wang et al., _Attention-based multi-instance neural network for medical diagnosis from incomplete and low quality data_ , in _2019 International Joint Conference on Neural Networks (IJCNN)_ (2019), pp. 1–8
  * (16) D. Grangier, I. Melvin, Advances in Neural Information Processing Systems 23 (2010)
  * (17) A. Vaswani et al., Advances in Neural Information Processing Systems 30 (2017)
  * (18) T. Sjöstrand et al., Computer Physics Communications 191, 159 (2015), 1410.3012
  * (19) R. Brun et al., Tech. rep., CERN (1994)
  * (20) J. Blitzer, M. Dredze, F. Pereira, _Biographies, Bollywood, Boom-boxes and Blenders: Domain Adaptation for Sentiment Classification_ , in _Proceedings of the 45th Annual Meeting of the Association of Computational Linguistics_ (2007), pp. 440–447
  * (21) X. Glorot, A. Bordes, Y. Bengio, _Domain adaptation for large-scale sentiment classification: A deep learning approach_ , in _ICML_ (2011)
  * (22) R. Gopalan, R. Li, R. Chellappa, _Domain adaptation for object recognition: An unsupervised approach_ , in _2011 International Conference on Computer Vision_ (IEEE, 2011), pp. 999–1006
  * (23) B. Fernando et al., _Unsupervised visual domain adaptation using subspace alignment_ , in _Proceedings of the IEEE International Conference on Computer Vision_ (2013), pp. 2960–2967
  * (24) D. Walter, Master thesis (2018)
  * (25) Y. Ganin et al., The Journal of Machine Learning Research 17, 2096 (2016)
  * (26) A. Jha et al., Modelling and Simulation in Materials Science and Engineering 27, 024002 (2019)
  * (27) Y. Gal, Z. Ghahramani, _Dropout as a Bayesian approximation: Representing model uncertainty in deep learning_ , in _International Conference on Machine Learning_ (PMLR, 2016), pp. 1050–1059
  * (28) V. Kuleshov, N. Fenner, S. Ermon, _Accurate uncertainties for deep learning using calibrated regression_ , in _International Conference on Machine Learning_ (PMLR, 2018), pp. 2796–2804


# Concatenated Masked Autoencoders as Spatial-Temporal Learner

Zhouqiang Jiang1, Bowen Wang2, Tong Xiang1, Zhaofeng Niu3, Hong Tang4,
Guangshun Li3, Liangzhi Li1∗ 1Meetyou AI Lab, China.2Institute for Datability
Science, Osaka University, Japan.3Department of Computer Science, Qufu Normal
University, China.4Department of Information Engineering, East China Jiaotong
University, China.∗Corresponding author<EMAIL_ADDRESS>
###### Abstract

Learning representations from videos requires understanding continuous motion
and visual correspondences between frames. In this paper, we introduce the
Concatenated Masked Autoencoders (CatMAE) as a spatial-temporal learner for
self-supervised video representation learning. For the input sequence of video
frames, CatMAE keeps the initial frame unchanged while applying substantial
masking (95%) to subsequent frames. The encoder in CatMAE is responsible for
encoding visible patches for each frame individually; subsequently, for each
masked frame, the decoder leverages visible patches from both previous and
current frames to reconstruct the original image. Our proposed method enables
the model to estimate the motion information between visible patches, match
the correspondences between preceding and succeeding frames, and ultimately
learn the evolution of scenes. Furthermore, we propose a new data augmentation
strategy, Video-Reverse (ViRe), which uses reversed video frames as the
model’s reconstruction targets. This further encourages the model to utilize
continuous motion details and correspondences to complete the reconstruction,
thereby enhancing the model’s capabilities. Compared to the most advanced pre-
training methods, CatMAE achieves a leading level in video segmentation tasks
and action recognition tasks. Code is available at
https://github.com/minhoooo1/CatMAE.

##  I Introduction

In the realm of image pre-training, Masked Autoencoders (MAE) [1] have
demonstrated their effectiveness in learning visual representations by
reconstructing missing patches from randomly masked input images. Recent
research extends this paradigm to the pre-training for videos [2, 3, 4, 5].
Existing studies consider the spatio-temporal specificity of masking
strategies in order to reduce inductive bias. However, they either overlook
differences in spatial and temporal dimensions [6] or neglect the ability to
model continuous motion. This leads to the dilemma where learned
representations struggle to simultaneously incorporate both continuous motion
information and the correspondences between video frames.

(a) Chopping wood.

(b) Skateboarding.

Figure 1: Visualizations on the Kinetics-400 [7] validation set (masking rate
90%). For each video sequence, we sample 6 frames with a frame gap of 4. Each
subfigure displays the original frames (top), masked future frames (middle),
and CatMAE reconstruction results (bottom).

Properly handling the temporal dimension in videos is essential to effectively
utilize the masking and reconstruction paradigm. VideoMAE [2] and MAE-ST [3]
expand 2D image patches into 3D cubes in videos, applying a self-supervised
mask reconstruction pipeline on these cubes. However, the semantics of video
frames vary slowly over time [8], creating temporal redundancy which increases
the risk of reconstructing missing cubes solely from the spatial-temporal
neighborhood. Therefore, to prevent the model from learning shortcuts by
exploiting this leaked information during reconstruction, a very high mask
rate is considered a straightforward universal solution, which also
significantly reduces computational cost [3]. The pre-trained models from
these methods demonstrate outstanding transferability on action recognition
tasks.

Yet, utilizing 3D cube masking in the spatial-temporal dimension is sub-
optimal for learning inter-frame correspondences; for instance, Image MAEs
outperform VideoMAE in video segmentation tasks [9]. Therefore, SiamMAE [4]
proposes to use an asymmetric mask strategy for reconstructing future frames
[10, 11] to learn inter-frame correspondences, and shows superiority in fine-
grained correspondence tasks such as video object segmentation [9], video part
segmentation [12], and pose tracking [13].

Continuous motion information and long-term correspondences usually span an
extended period within a video. As such, predicting and reconstructing the
future frames over a longer interval using the asymmetric masking strategy in
SiamMAE is inherently hard due to their ambiguity. In extreme cases, the model
may not be able to model any effective motion information, as the scenes in
the two sampled frames could be entirely different.

To capture continuous motion information and long-term correspondences in
videos, we propose Concatenated Masked AutoEncoders (CatMAE), which uses a
concatenated information channel masking strategy to enhance the learning
ability of the encoder. In our method, we first chronologically select $N$
frames from a video clip, keeping all patches of the first frame visible and
then performing random masking on patches of subsequent frames with an
extremely high mask ratio. The encoder encodes the visible patches from $N$
frames separately, and then the decoder reconstructs the masked patches for
$N$-1 subsequent frames. Note that when reconstructing each frame’s missing
patches, the decoder utilizes cross-attention to receive information from the
visible patches of previous and current frames. This concatenated information
channel masking assists in modeling motion differences and correspondences,
eventually estimating the long-term dynamic evolution of video frames. As
shown in Fig 1, CatMAE is able to reconstruct a long video sequence.
Furthermore, inspired by time’s bidirectionality, we hypothesize that the
reconstruction of reversed actions can reinforce the understanding of actions,
and propose a video reverse data augmentation method, Video-Reverse (ViRe), to
enhance the representation learning.

In conclusion, CatMAE enriches the exploration of the masking and
reconstruction pipeline [1] in the field of self-supervised video
representation learning. In video object segmentation tasks CatMAE based on
ViT-S/16 improves by 2% in $\mathcal{J}\&\mathcal{F}_{m}$. Additionally, we
also achieve leading results in the action recognition task. Experiment
results show that our proposed CatMAE and ViRe arouses the full potential of
our paradigm of masking and reconstruction within the spatial-temporal domain.

Figure 2: Pipeline of our CatMAE. During pre-training, we chronologically
extract $N$ frames from a video clip, keeping all patches of the first frame
visible, and apply a very high masking ratio to mask the patches of the
subsequent $N-$1 frames randomly. The visible patches of $N$ frames are
independently processed by the ViT [5] encoder. The decoder reconstructs the
missing patches for the subsequent $N-$1 frames. Note that when reconstructing
each frame’s missing patches, the decoder receives information from all
visible patches of previous frames through a sequence of cross-attention
layers [4].

##  II Related Work

###  II-A Masked Visual Modeling

Masked visual modeling method [14] learn representations from images disrupted
by masking. Essentially, they function as denoising autoencoders, disrupting
input signals and reconstructing the original undamaged signals to learn
effective representations. This paradigm has spawned a range of derivatives,
such as reconstructing masking pixels [15, 16] or restoring lost color
channels [17].

The success of masked language modeling [18] in NLP’s self-supervised pre-
training, along with the popularity of ViT [5], has sparked extensive research
into the use of transformer-based architectures for masked visual modeling
within the field of computer vision [19, 1, 20, 21, 22, 23]. BEiT [19], PeCo
[20], and iBOT [23] naturally inherit the idea of BERT [18] and propose
learning representations from images by predicting discrete tokens.

Some other research [24, 1, 22] focuses on using pixels as prediction targets,
which is simpler. MAE [1] and SimMIM [22], by learning to reconstruct missing
patches from randomly masked input image patches, achieve good visual
representation. Moreover, MAE’s encoder only handles visible patches under a
high mask ratio, significantly speeding up training and achieving better
transfer performance. MAE has also been straightforwardly extended to the
video domain [2, 3], reconstructing masked cubes. However, cube masking is
sub-optimal for learning correspondences. SiamMAE [4] proposed an asymmetric
masking strategy that retains past frames while reconstructing future frames
containing a high proportion of masks, encouraging the model to model motion
and match inter-frame correspondences. However, this strategy struggle to
model long-term continuous motion information. Therefore, we propose a
concatenated information channel masking reconstruction strategy that extends
motion modeling to a theoretically infinite frame interval.

###  II-B Contrastive Based Self-supervision

Effectively utilizing the temporal dimension is crucial in self-supervised
video representation learning. Currently, there are various pretext tasks for
pre-training, including predicting the future [25, 26, 27], segmenting pseudo
ground truth [28], reconstructing future frames [10, 11, 29], tracking [30,
31], reference coloring [32], and temporal ordering [33, 34, 35, 36]. More
advanced contrastive learning methods [37, 38] have been developed, which
learn representations by modeling image similarities and dissimilarities [39,
40] or solely similarities [41, 42, 43]. However, these methods rely on large
batches [40], multi-crops [44], negative key queues [39], or custom strategies
to prevent representation collapse [42]. Their performance greatly depends on
the choice of image augmentation [40]. In contrast, our method is based on a
simple masking and reconstruction pipeline [1].

##  III Method

In this section, We introduce each key component of CatMAE. The overall model
architecture is shown in Fig 2.

Patch Embedding. First, we chronologically select $N$ frames from a video clip
as input sequence. The interval between frames is randomly chosen from a pre-
determined frame gap range. Then, following the process of ViT [5], we divide
each frame into non-overlapping patches. We apply linear projection and
flatten the patches, concatenate a [CLS] token to patches and eventually add
position embedding to them.

Concatenated Information Channel Masking. We keep the first frame completely
visible and apply masking with a high masking rate to $N-$1 subsequent frames,
preserving only a minimal number of visible patches. Given that video signals
are highly redundant [8], high masking rate prevents the model from utilizing
information from adjacent frames during the reconstruction process, thus
encouraging the encoder to capture motion information and correspondence.

Visible patches in subsequent frames provide information channels,
theoretically allowing the frames’ sequence to be reconstructed indefinitely.
To reconstruct the final frame, the model needs to reconstruct the continuous
motion and correspondence of the intermediate frames, ultimately realizing the
evolution between frames. (as shown in Fig. 1).

Encoder. We employ weight-shared vanilla ViT [5] to independently process all
frames, with each encoder only encoding visible patches. This design
significantly reduces temporal and memory complexity while achieving better
performance [1]. Furthermore, the patch embedding only adds a position
embedding in the spatial structure, hence, the encoder is unaware of the
temporal structure of the patches. However, the decoder’s cross-attention
layer needs to use the encoder’s output to propagate the visible content. Even
when the temporal structure is unknown, the encoder is forced to match the
correspondence between spatial-temporal patches before and after motion. This
correspondence assists the decoder in utilizing the information channels to
propagate visible content, ultimately achieving the reconstruction of
subsequent frames.

Decoder. We also employ a weight-shared decoder to separately predict the
reconstructed frames’ masked patches. Each decoder block consists of a cross-
attention layer and a self-attention layer [45]. Specifically, the visible
tokens of the frame to be reconstructed are projected via a linear layer and
combined with mask tokens to form a set of full tokens, to which spatial
position embedding is then added. Subsequently, the full tokens attend to all
previously visible tokens (tokens of the first frame and other visible tokens)
through a cross-attention layer, followed by mutual attention through a self-
attention layer. Our concatenated information channel masking enables the
decoder to enhance the encoder’s ability to estimate motion offsets and learn
correspondences. Finally, The output sequence of the decoder is used to
predict the normalized pixel values in the masked patches [1], with an L2 loss
applied between the decoder’s prediction and ground truth.

##  IV Experiments

###  IV-A Pre-training

Architecture. By default, we use ViT-S/16 as the encoder, with a depth of 12
and a width of 384, which is similar to ResNet-50 [46] in terms of the number
of parameters (21M vs. 23M). We use a small decoder with a width of 256 and a
depth of 4. The effectiveness of this setup is demonstrated in Table III(b)
and III(c).

Settings. We conduct self-supervised pre-training on the Kinetics-400 [7]
action recognition dataset, which encompasses a total of 400 classes,
comprising 240k training videos and 20k validation videos, and we only conduct
pre-training on the training set. Considering the cost of training, we default
to sampling three frames (224x224) from each video at certain frame intervals,
and employ two image augmentation methods: RandomResizeCrop and horizontal
flipping. A large portion (95%) of the patches in the last two frames are
randomly masked, and the batch size is 2048. The base learning rate is
initialized to 1e-4 and follows a cosine schedule decay. We use the AdamW [47]
optimizer, and the reconstruction losses for the second and 3rd frames are
scaled by 0.8 and 1.0, respectively. We adopt a repetition sampling [4, 48]
factor of 2 and finally conduct pre-training for 800 effective epochs.

###  IV-B Evaluation Method

We conduct two downstream experiments, video segmentation and action
recognition, to evaluate the correspondence and motion information of the
obtained representations.

Fine-grained Correspondence Task. We use the semi-supervised video
segmentation dataset DAVIS-2017 [9] to evaluate the quality of representations
in fine-grained correspondence tasks. We adopt the label propagation method
used in [49, 4] to implement semi-supervised video segmentation. Given the
ground truth labels of the first frame, we obtain k-nearest neighbor results
through a similarity function between frames. Afterward, the ground truth
labels of the first frame are propagated to subsequent frames as pixel labels,
based on the k-nearest neighbors’ results.

Action Recognition Task. We followed the evaluation protocol of VideoMAE [2]
for the motion evaluation on action recognition task by Kinetics-400 [7]. In
order to avoid a variety of meticulously crafted fine-tuning settings
obstructing fair comparisons between different pre-training methods, we used
the same pipeline to uniformly fine-tune all pre-trained models on the
training set for 150 epochs and report the results on the validation set. To
fairly compare with video MAEs [2, 3], we also will extend the 2D patch
embedding layer to a 3D embedding layer following the approach proposed in
[7].

###  IV-C Main Results and Analysis

Fine-grained Correspondence Task. Table I demonstrates all the methods using
480x880 images. We find that CatMAE-ViT-S/8 surpasses all other self-
supervised or supervised methods, and its performance is further enhanced when
ViRe is applied. This demonstrates that using reversed videos as
reconstruction targets can further enhance the model’s ability to match
correspondences. As analyzed before, the cube masking recovery strategy cannot
learn the inter-frame correspondence, resulting in MAE-ST and VideoMAE being
significantly lower than CatMAE. In addition, CatMAE also outperforms SiamMAE
by a large margin, owning a better segmentation performance with larger frame
intervals (Table III(a)), proving that using the concatenated information
channel masking strategy can capture longer and more accurate correspondences
in videos.

Method | Backbone | Dataset | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---|---|---
Surpervised [46] | ResNet-18 | ImageNet | 63.8 | 60.4 | 67.2
Surpervised [46] | ResNet-50 | ImageNet | 67.9 | 64.7 | 71.2
SimSiam [41] | ResNet-50 | ImageNet | 66.7 | 63.4 | 70.0
MoCo-v2 [50] | ResNet-50 | ImageNet | 66.5 | 63.2 | 69.7
DINO [43] | ResNet-50 | ImageNet | 66.3 | 63.0 | 69.5
VFS [51] | ResNet-50 | Kinetics | 65.0 | 62.0 | 68.1
CRW [49] | ResNet-18 | Kinetics | 67.4 | 64.2 | 70.6
MAE [1] | ViT-B/16 | ImageNet | 52.0 | 49.5 | 54.6
MAE-ST [3] | ViT-L/2x16x16 | Kinetics | 54.6 | 55.5 | 53.6
VideoMAE [2] | ViT-S/2x16x16 | Kinetics | 39.3 | 39.7 | 38.9
DINO [43] | ViT-S/16 | ImageNet | 61.5 | 58.1 | 64.9
SiamMAE [4] | ViT-S/16 | Kinetics | 61.3 | 58.4 | 64.2
CatMAE | ViT-S/16 | Kinetics | 62.0 | 59.0 | 65.0
CatMAE + ViRe | ViT-S/16 | Kinetics | 62.5 | 59.4 | 65.5
DINO [43] | ViT-S/8 | ImageNet | 70.3 | 66.7 | 73.9
SiamMAE [4] | ViT-S/8 | Kinetics | 69.0 | 65.6 | 72.4
CatMAE | ViT-S/8 | Kinetics | 71.1 | 67.1 | 75.1
CatMAE + ViRe | ViT-S/8 | Kinetics | 72.3 | 68.0 | 76.6
TABLE I: Video object segmentation results on DAVIS 2017 validation set. Comparison of our method (ViRe indicates that ViRe is utilized during the pre-training phase), with the state-of-the-art self-supervised methods and pre-trained feature baselines. $\mathcal{F}$ is a boundary alignment metric, while $\mathcal{J}$ measures region similarity as IOU between masks. Method | Backbone | Epoch | Top-1 | Top-5
---|---|---|---|---
VideoMAE [2] | ViT-S/2x16x16 | 10 | 47.70 | 73.99
SiamMAE [4] | ViT-S/16 | 10 | 40.68 | 67.39
CatMAE | ViT-S/16 | 10 | 46.41 | 72.87
VideoMAE [2] | ViT-S/2x16x16 | 150 | 72.54 | 89.60
SiamMAE [4] | ViT-S/16 | 150 | 66.42 | 85.96
CatMAE | ViT-S/16 | 150 | 68.50 | 87.63
TABLE II: Action recognition results during the fine-tuning stage on the
Kinetics-400 validation set. Comparison of our method, with MAE variants.

Action Recognition Task. We implement this experiment with two competitive
video representation methods VideoMAE and SiamMAE. All methods undergo 800
epochs pre-training. In Table II, it is observed that during the early stage
of fine-tuning, CatMAE exhibits a significantly stronger transfer capability
than SiamMAE, with an enhancement in top-1 accuracy by 5.73%. This indicates
that our proposed concatenated information channel masking reconstruction pre-
training strategy can capture more motion information than the asymmetric mask
reconstruction strategy of SiamMAE. Even when the fine-tuning saturates,
CatMAE still leads SiamMAE by 2.08% in top-1 accuracy. However, the gap
between CatMAE and VideoMAE widen, which we believe mainly stems from the
method of modeling tokens. VideoMAE directly models token embedding using a 3D
cube, which essentially makes VideoMAE a type of 3D network. VideoMAE
naturally contains more motion information than CatMAE which encodes 2D patch
as tokens. Yet, in the initial stage of fine-tuning, CatMAE manages to
maintain comparable transfer ability to VideoMAE (only 1.29% lower in top-1
accuracy), performing far better than the SiamMAE (7.02% lower). Our CatMAE
maximally exploits the motion information on the 2D patch.

rec. num | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
1 | 54.3 | 52.7 | 55.9
2 | 56.0 | 54.2 | 57.8
3 | 57.5 | 55.3 | 59.7
4 | 58.2 | 56.2 | 60.2
(a) Reconstruction number. More reconstructions can improve performance.

dim | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
192 | 56.0 | 54.2 | 57.8
256 | 59.9 | 57.4 | 62.5
512 | 57.3 | 55.0 | 59.6
1024 | 55.6 | 53.2 | 58.0
(b) Decoder width. A narrower decoder can achieve the best performance.

blocks | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
1 | 51.7 | 49.7 | 53.7
2 | 56.0 | 54.2 | 57.8
4 | 60.0 | 57.3 | 62.6
8 | 57.4 | 54.0 | 60.8
(c) Decoder depth. 4 blocks of decoder achieve the best performance.

mask ratio | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
0.90, 0.90 | 54.5 | 52.8 | 56.2
0.95, 0.90 | 54.6 | 53.0 | 56.2
0.90, 0.95 | 55.6 | 53.7 | 57.5
0.95, 0.95 | 56.0 | 54.2 | 57.8
0.99, 0.99 | 42.1 | 40.0 | 44.2
(d) Mask ratio. A high masking ratio (95%) works well, but not to extremes.

scale | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
0.5 1.0 | 54.6 | 52.5 | 56.7
0.8 1.0 | 56.0 | 54.2 | 57.8
1.0 0.8 | 55.0 | 53.4 | 56.7
1.0 1.0 | 55.8 | 54.0 | 57.6
(e) Loss scale. slightly emphasizing the reconstruction of the third frame is
optimal.

frame gap | $\mathcal{J}\&\mathcal{F}_{m}$ | $\mathcal{J}_{m}$ | $\mathcal{F}_{m}$
---|---|---|---
[4,16]-[4,16] | 56.0 | 54.2 | 57.8
[4,48]-[4,16] | 55.2 | 53.3 | 57.1
[4,16]-[4,48] | 56.8 | 54.9 | 58.6
[4,48]-[4,48] | 55.6 | 53.9 | 57.3
(f) Frame gap. Using a smaller sample rate between the first and 2nd frames
works better.

TABLE III: CatMAE ablation experiments with ViT-S/16 on DAVIS 2017 validation
set. If not specified, the default is: the number of reconstructions is 2, the
decoder has a width of 256 and a depth of 4, the Mask ratio and Loss scale for
the two reconstruction frames are [0.95, 0.95] and [0.8, 1.0], respectively,
the frame gap between frames 1-2 and 2-3 are both [4,16]. Default settings are
marked in grey. Figure 3: Impact of pretraining epochs and propagation FPS on
video segmentation. Comparison of Top-1 accuracy for VideoMAE, CatMAE, and
SiamMAE under different fine-tuning epochs.

###  IV-D Ablation Experiments

We conduct an ablation experiment to understand each component’s contribution.
Details can be found in Table III.

Number of Reconstructions. Fig. 1 shows the model’s reconstruction results for
future frames. As can be seen, the model is highly sensitive to motion
elements in video frames. In Fig. 1(a), the model can distinguish between
static backgrounds and accurately reconstruct the motion of waving hammers.
Beyond its motion modeling capability, the model can copy the background from
the first frame and paste it into future frames when dealing with static
backgrounds (such as the static woodpile). As shown in Table III(a), the video
segmentation performance steadily improves with the increase in the number of
reconstructions. When the number of reconstructions is 4, the performance is
higher in comparison with the number being 1 ($\uparrow$3.9%). This validates
the benefits of using a concatenated information channel masking strategy to
extend multi-frame reconstruction.

(a) Horse.

(b) Seabird.

(c) Black swan.

(d) Orca.

Figure 4: Cross-attention maps from a ViT-S/16 decoder (samples from
DAVIS-2017 [9]). We visualize the attention of the mask patch to be
reconstructed in the third frame towards the visible patches in the first and
second frames at the average head in the last layer of the decoder.

Decoder Capacity. In Table III(b) and III(c), we respectively examine the
width and depth of the decoder. We found that relative to the encoder (384-d,
12-block), a smaller decoder (192-d, 4-block) is critically important for
fine-grained correspondence tasks. This can be explained by the relationship
between pixel reconstruction and fine-grained feature correspondence: during
pixel reconstruction, the encoder is sensitive to motion offsets, thus
obtaining a more reasonable correspondence representation. However, when the
decoder capacity is too large, it weakens the encoder’s sensitivity to motion.
Therefore, in terms of the width and depth settings for the decoder, an
optimal width of 256 is 4.3% better than the maximum width of 1024 (59.9% vs.
55.6%). Additionally, an optimal depth of 4 is 2.6% better than the depth of 8
(60.0% vs. 57.4%). These optimal decoder capacities are much smaller than the
encoder, underscoring the importance of utilizing a smaller decoder.

Mask Ratio. In Table III(d), we compare different mask ratio. Based on
previous experience with video MAEs, we default to ablating from a high mask
ratio (90%), but under such a high mask ratio, increasing the mask ratio of
the last two frames to 95% still shows performance improvement (54.5% vs.
56.0%). However, when using an extreme mask ratio of 99% (Only one patch is
visible), the performance drops sharply. An information channel with limited
capacity is crucial for the model to learn correspondences. However, when the
channel capacity is reduced to its extreme limit, it becomes difficult for the
model to match correspondences. Hence, we default to using a mask ratio of
95%.

Loss Scale. In Table III(e), we compared the reconstruction loss scale between
the second and third frames. We found that slightly emphasizing the
reconstruction of the third frame is better than that of the second (compare
the second and third rows in Table III(e)). We believe that focusing more on
the third frame’s reconstruction facilitates the progressive transmission of
continuous motion information to the final frame. However, neglecting the
reconstruction of the second frame by reducing its loss scale to 0.5 leads to
performance degradation. This might weaken the model’s capability to capture
the intermediate motion, thus affecting the propagation of continuous motion
information.

Frame Gap. In Table III(f), we compare different frame gaps. While a larger
frame gap sampling can encompass richer motion information, reconstructing two
frames with large gaps does not bring performance improvement to the model
(compare the first and fourth rows). This may be because when the frame gap
becomes too large, the scene changes are too dramatic, which hinders the
model’s ability to model motion information. Additionally, we found that
having a smaller gap between the first and second frames and a larger gap
between the second and third frames (see the third row in Table) performs
better than the reverse configuration (see the second row in Table) (56.8% vs.
55.2%). This observation can also be interpreted as the propagation of
continuous motion being hindered: a large gap between the first and second
frames makes capturing intermediate motion difficult, thereby blocking the
transmission of continuous motion information within the information channels.

Training Epochs and Propagation FPS. In Fig. 3, we analyze the impact of
training epochs and propagation FPS on video segmentation, discovering that
optimal performance is achieved at pre-training 400 epochs, after which the
performance gradually declines, possibly due to overfitting caused by small
decoder. We observe that under various low FPS conditions, CatMAE consistently
outperforms SiamMAE, demonstrating that CatMAE possesses a stronger ability to
match long-term correspondences. Additionally, we delve into the influence of
fine-tuning epochs. CatMAE initially exhibits robust transferability but the
gap widens with VideoMAE over time due to differences in token embedding.
Simplistically speaking, CatMAE can be seen as a ”additional frame
reconstruction” version of SiamMAE. Therefore, according to Action Recognition
figure, we estimate that using CatMAE to pre-train weights for reconstructing
future four frames holds promise for reaching the fine-tuning performance of
VideoMAE.

###  IV-E Motion and Correspondence Analysis

In Fig. 4, we visualize the cross-attention map averaged across different
heads in the final layer of the ViT-S/16 decoder. Specifically, we input the
first frame and the subsequent two frames with masks, followed by
reconstructing the third frame. We collect the cross-attention from the mask
patch in the third frame towards the visible patches in the first and second
frames and then visualize this attention map. We find that the model is
sensitive to motion, focusing on the correspondence between patches under
continuous motion. For instance (Fig. 4(a)), when calculating the
correspondence of the horse’s head patch in the third frame, the model can
clearly distinguish the semantics of the brown region and attention accurately
focused on the head of the horse in previous frames.

Even when the background color is similar to the possible corresponding area
(such as the white head of the water bird and the reflective water surface
Fig. 4(b)), the model’s corresponding capability remains robust. Certainly,
the model can also establish correspondences by leveraging unique color
features (such as the red bill of a swan Fig. 4(c)) or specific semantic
structures (such as the fin of an Orca Fig. 4(d)). These provide direct
evidence of the model’s ability to estimate motion and learn correspondences.

##  V Conclusion

In this paper, we proposed CatMAE for self-supervised video representation
learning. It leverages a concatenated information channel masking strategy to
address the limitations posed by cube masking and enhances the capability to
capture continuous and long-term motion compared to asymmetric masking. Our
experimental results demonstrate superior performance in comparison to state-
of-the-art methods across both video segmentation and action recognition
tasks. One distinctive feature of our training pipeline is the propagation of
reconstruction information from the initial frame throughout the entire video
sequence. This theoretically unlimited propagation showcases CatMAE’s
potential to learn long-term video representations. Our future work focuses on
extending the application of CatMAE to real-world scenarios involving embodied
agents, such as robots.

## References

  * [1] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick, “Masked autoencoders are scalable vision learners,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2022, pp. 16 000–16 009.
  * [2] Z. Tong, Y. Song, J. Wang, and L. Wang, “Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training,” _Advances in neural information processing systems_ , vol. 35, pp. 10 078–10 093, 2022.
  * [3] C. Feichtenhofer, Y. Li, K. He _et al._ , “Masked autoencoders as spatiotemporal learners,” _Advances in neural information processing systems_ , vol. 35, pp. 35 946–35 958, 2022.
  * [4] A. Gupta, J. Wu, J. Deng, and L. Fei-Fei, “Siamese masked autoencoders,” _arXiv preprint arXiv:2305.14344_ , 2023.
  * [5] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly _et al._ , “An image is worth 16x16 words: Transformers for image recognition at scale,” _arXiv preprint arXiv:2010.11929_ , 2020.
  * [6] E. H. Adelson and J. R. Bergen, “Spatiotemporal energy models for the perception of motion,” _Josa a_ , vol. 2, no. 2, pp. 284–299, 1985.
  * [7] J. Carreira and A. Zisserman, “Quo vadis, action recognition? a new model and the kinetics dataset,” in _proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2017, pp. 6299–6308.
  * [8] Z. Zhang and D. Tao, “Slow feature analysis for human action recognition,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 34, no. 3, pp. 436–450, 2012.
  * [9] J. Pont-Tuset, F. Perazzi, S. Caelles, P. Arbeláez, A. Sorkine-Hornung, and L. Van Gool, “The 2017 davis challenge on video object segmentation,” _arXiv preprint arXiv:1704.00675_ , 2017.
  * [10] W. Lotter, G. Kreiman, and D. Cox, “Deep predictive coding networks for video prediction and unsupervised learning,” _arXiv preprint arXiv:1605.08104_ , 2016.
  * [11] M. Mathieu, C. Couprie, and Y. LeCun, “Deep multi-scale video prediction beyond mean square error,” _arXiv preprint arXiv:1511.05440_ , 2015.
  * [12] Q. Zhou, X. Liang, K. Gong, and L. Lin, “Adaptive temporal encoding network for video instance-level human parsing,” in _Proceedings of the 26th ACM international conference on Multimedia_ , 2018, pp. 1527–1535.
  * [13] H. Jhuang, J. Gall, S. Zuffi, C. Schmid, and M. J. Black, “Towards understanding action recognition,” in _Proceedings of the IEEE international conference on computer vision_ , 2013, pp. 3192–3199.
  * [14] P. Vincent, H. Larochelle, Y. Bengio, and P.-A. Manzagol, “Extracting and composing robust features with denoising autoencoders,” in _Proceedings of the 25th international conference on Machine learning_ , 2008, pp. 1096–1103.
  * [15] P. Vincent, H. Larochelle, I. Lajoie, Y. Bengio, P.-A. Manzagol, and L. Bottou, “Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion.” _Journal of machine learning research_ , vol. 11, no. 12, 2010.
  * [16] D. Pathak, P. Krahenbuhl, J. Donahue, T. Darrell, and A. A. Efros, “Context encoders: Feature learning by inpainting,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 2536–2544.
  * [17] R. Zhang, P. Isola, and A. A. Efros, “Colorful image colorization,” in _Computer Vision–ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part III 14_. Springer, 2016, pp. 649–666.
  * [18] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” _arXiv preprint arXiv:1810.04805_ , 2018.
  * [19] H. Bao, L. Dong, S. Piao, and F. Wei, “Beit: Bert pre-training of image transformers,” _arXiv preprint arXiv:2106.08254_ , 2021.
  * [20] X. Dong, J. Bao, T. Zhang, D. Chen, W. Zhang, L. Yuan, D. Chen, F. Wen, N. Yu, and B. Guo, “Peco: Perceptual codebook for bert pre-training of vision transformers,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 37, no. 1, 2023, pp. 552–560.
  * [21] C. Wei, H. Fan, S. Xie, C.-Y. Wu, A. Yuille, and C. Feichtenhofer, “Masked feature prediction for self-supervised visual pre-training,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 14 668–14 678.
  * [22] Z. Xie, Z. Zhang, Y. Cao, Y. Lin, J. Bao, Z. Yao, Q. Dai, and H. Hu, “Simmim: A simple framework for masked image modeling,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 9653–9663.
  * [23] J. Zhou, C. Wei, H. Wang, W. Shen, C. Xie, A. Yuille, and T. Kong, “ibot: Image bert pre-training with online tokenizer,” _arXiv preprint arXiv:2111.07832_ , 2021.
  * [24] M. Chen, A. Radford, R. Child, J. Wu, H. Jun, D. Luan, and I. Sutskever, “Generative pretraining from pixels,” in _International conference on machine learning_. PMLR, 2020, pp. 1691–1703.
  * [25] J. Walker, C. Doersch, A. Gupta, and M. Hebert, “An uncertain future: Forecasting from static images using variational autoencoders,” in _Computer Vision–ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11–14, 2016, Proceedings, Part VII 14_. Springer, 2016, pp. 835–851.
  * [26] A. Gupta, S. Tian, Y. Zhang, J. Wu, R. Martín-Martín, and L. Fei-Fei, “Maskvit: Masked visual pre-training for video prediction,” _arXiv preprint arXiv:2206.11894_ , 2022.
  * [27] C. Vondrick, H. Pirsiavash, and A. Torralba, “Anticipating visual representations from unlabeled video,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 98–106.
  * [28] D. Pathak, R. Girshick, P. Dollár, T. Darrell, and B. Hariharan, “Learning features by watching objects move,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017, pp. 2701–2710.
  * [29] N. Srivastava, E. Mansimov, and R. Salakhudinov, “Unsupervised learning of video representations using lstms,” in _International conference on machine learning_. PMLR, 2015, pp. 843–852.
  * [30] X. Wang, A. Jabri, and A. A. Efros, “Learning correspondence from the cycle-consistency of time,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 2566–2576.
  * [31] X. Wang and A. Gupta, “Unsupervised learning of visual representations using videos,” in _Proceedings of the IEEE international conference on computer vision_ , 2015, pp. 2794–2802.
  * [32] C. Vondrick, A. Shrivastava, A. Fathi, S. Guadarrama, and K. Murphy, “Tracking emerges by colorizing videos,” in _Proceedings of the European conference on computer vision (ECCV)_ , 2018, pp. 391–408.
  * [33] B. Fernando, H. Bilen, E. Gavves, and S. Gould, “Self-supervised video representation learning with odd-one-out networks,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017, pp. 3636–3645.
  * [34] H.-Y. Lee, J.-B. Huang, M. Singh, and M.-H. Yang, “Unsupervised representation learning by sorting sequences,” in _Proceedings of the IEEE international conference on computer vision_ , 2017, pp. 667–676.
  * [35] I. Misra, C. L. Zitnick, and M. Hebert, “Shuffle and learn: unsupervised learning using temporal order verification,” in _Computer Vision–ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11–14, 2016, Proceedings, Part I 14_. Springer, 2016, pp. 527–544.
  * [36] D. Wei, J. J. Lim, A. Zisserman, and W. T. Freeman, “Learning and using the arrow of time,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2018, pp. 8052–8060.
  * [37] R. Hadsell, S. Chopra, and Y. LeCun, “Dimensionality reduction by learning an invariant mapping,” in _2006 IEEE computer society conference on computer vision and pattern recognition (CVPR’06)_ , vol. 2. IEEE, 2006, pp. 1735–1742.
  * [38] S. Becker and G. E. Hinton, “Self-organizing neural network that discovers surfaces in random-dot stereograms,” _Nature_ , vol. 355, no. 6356, pp. 161–163, 1992.
  * [39] K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick, “Momentum contrast for unsupervised visual representation learning,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2020, pp. 9729–9738.
  * [40] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, “A simple framework for contrastive learning of visual representations,” in _International conference on machine learning_. PMLR, 2020, pp. 1597–1607.
  * [41] X. Chen and K. He, “Exploring simple siamese representation learning,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2021, pp. 15 750–15 758.
  * [42] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar _et al._ , “Bootstrap your own latent-a new approach to self-supervised learning,” _Advances in neural information processing systems_ , vol. 33, pp. 21 271–21 284, 2020.
  * [43] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin, “Emerging properties in self-supervised vision transformers,” in _Proceedings of the IEEE/CVF international conference on computer vision_ , 2021, pp. 9650–9660.
  * [44] M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin, “Unsupervised learning of visual features by contrasting cluster assignments,” _Advances in neural information processing systems_ , vol. 33, pp. 9912–9924, 2020.
  * [45] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [46] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778.
  * [47] I. Loshchilov and F. Hutter, “Decoupled weight decay regularization,” _arXiv preprint arXiv:1711.05101_ , 2017.
  * [48] E. Hoffer, T. Ben-Nun, I. Hubara, N. Giladi, T. Hoefler, and D. Soudry, “Augment your batch: Improving generalization through instance repetition,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 8129–8138.
  * [49] A. Jabri, A. Owens, and A. Efros, “Space-time correspondence as a contrastive random walk,” _Advances in neural information processing systems_ , vol. 33, pp. 19 545–19 560, 2020.
  * [50] X. Chen, H. Fan, R. Girshick, and K. He, “Improved baselines with momentum contrastive learning,” _arXiv preprint arXiv:2003.04297_ , 2020.
  * [51] J. Xu and X. Wang, “Rethinking self-supervised correspondence learning: A video frame-level similarity perspective,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 10 075–10 085.