Masked Capsule Autoencoders
Authors: Miles Everett
Mingjun Zhong
Georgios Leontidis
Venue: Accepted by TMLR ()

Abstract: We propose Masked Capsule Autoencoders (MCAE), the first Capsule Network that utilises pretraining in a modern self-supervised paradigm, specifically the masked image modelling framework. Capsule Networks have emerged as a powerful alternative to Convolutional Neural Networks (CNNs). They have shown favourable properties when compared to Vision Transformers (ViT), but have struggled to effectively learn when presented with more complex data. This has led to Capsule Network models that do not scale to modern tasks. Our proposed MCAE model alleviates this issue by reformulating the Capsule Network to use masked image modelling as a pretraining stage before finetuning in a supervised manner. Across several experiments and ablations studies we demonstrate that similarly to CNNs and ViTs, Capsule Networks can also benefit from self-supervised pretraining, paving the way for further advancements in this neural network domain. For instance, by pretraining on the Imagenette dataset---consisting of 10 classes of Imagenet-sized images---we achieve state-of-the-art results for Capsule Networks, demonstrating a 9\% improvement compared to our baseline model. Thus, we propose that Capsule Networks benefit from and should be trained within a masked image modelling framework, using a novel capsule decoder, to enhance a Capsule Network's performance on realistically sized images.

Published in Transactions on Machine Learning Research (01/2025)
Masked Capsule Autoencoders
Miles Everctt
liteessceerel@@aadnn.cc..c
Deportment of Computing Science University of Aberdeen, UK
Mingjun Zhong
minggxon.zhong@@bbdn..cc
Deportment of Computing Science University of Aberdecn, UK
Georgios Leontidis
georgjos.conttidss@abbd..ccc
Interdisciplinory Institute Department of Compuling Science University of Aberdeen, UK
Reviewed on OpenReview: https: // openreview. net/ forum? id= JHarh00W11
Abstract
We propose Masked Capsule Autoencoders (MCAE), the first Capsule Network that utilises protraining a modern self-supervised paradigm, specifically the masked image modelling framework. Capsule Networks have emerged as a powerful alternative to Convolutional Neural Networks (CNNs). They have shown favourable properties when compared to Vi- sion Transformers (VIT), but have struggled to effectively learn when presented with more complex data. This has led to Capsule Network models that do not scale to modern tasks. Our proposed MCAE model alleviates this issue by reformulating the Capsule Network to use masked image modelling as a pretraining stage before finetuning in a supervised manner. Across several experiments and ablations studics we demonstrate that similarly to CNNs and ViTs, Capsule Networks can also benefit from self-supervised pretraining, paving the way for further advancements in this neural network domain. For instance, by pretraining on the Imagenette dataset. consisting of 10 classes of Imagenet-sized images we achieve state-of-the-art results for Capsule Networks, demonstrating 9% improvement compared to our baseline model. Thus, we propose that Capsule Networks benefit from and should be trained within a masked image modelling framework, using a novel capsule decoder, to enhance Capsule Network's performance on realistically sized images.
1 Introduction
Capsule Networks are an evolution of Convolutional Neural Networks (CNNs) which remove pooling opera- tions and replace scalar neurons with a fixed number of vector or matrix representations known as capsules at each location in the feature map. At each location, there will multiple capsules, each theoretically representing a different concept. Each of these capsules has corresponding activation value between and 1. The activation value measures the probability (from to 1) that the concept represented by the capsule exists. Capsule Networks have shown promising signs, such as being naturally strong in invariant and equivariant tasks (Sabour et al. 2017 Hinton et al. 2018 De Sousa Ribeiro et al. 2020 Hahn et al. 2019 Ribeiro et al. 2020 2022) while having low parameter counts, but have yet to scale to more complex datasets with realistic resolutions that CNNs and Vision Transformers (ViTs) are typically benchmarked on.
Masked Image Modelling (MTM) is Self Supervisod Learning (SSL) technique with roots language modelling (Devlin al.) 20188. In language modelling, words are removed from passages of text, the network is then trained to predict the correct words to fill in the gaps. This technique be extended to

---

Published in Transactions on Machine Learning Research (01/202)
Car
Rander
Cipark
Fadero Map
Plotee
niraaeied
Rastape
Figure 1: Our Masked Capsule Autoencoder architecture. During pretraining we randomly select number of patches from the original image to be processed. The Capsule Network will then create representation for each patch. Masked patch capsule representations are then re-added before the capsule decoder, where the unmasked capsules can contribute to the masked positions, which are finally decoded by a single linear layer to the original patch dimensions. The pretraining objective is the mean squared error between the reconstructed patches and the target patches. The dog image used is sourced from the Imagewoof validation set (Howard 2019b).
image modelling by splitting an image into equal regions called patches, randomly removing some of these patches and then requiring the network to predict the pixel values of the removed patches. This has been shown to require the network to have an improved world model in both Vision Transformers (ViTs) (He ct al. 2021) and CNNs (Woo et al. 2023). which is strong enough to reconstruct occluded areas from the remaining visible areas. Combining this technique with supervised finetuning, accuracy can be significantly improved compared to not using any pretraining (He et al., 2021 Woo et al, 2023).
We propose incorporating MIM pretraining into the training process of Capsule Networks to address its shortcomings. By doing so, Capsule Networks can learn more refined representations for each region of the image, leading to improved reconstruction accuracy. These refined local representations employed during the finetuning phase to effectively activate the appropriate global class capsules, which are introduced after pretraining. These proposals are based on the observed significant improvements in CNNs and ViTs, particularly in terms of better training regimes with minimal architectural changes. Instead of solely focusing on enhancing the routing algorithm, we demonstrate that Capsule Networks can improved through training modifications.
subsequently be
The main contributions of this work can be summarised follows:
We propose a novel adaption of Capsule Networks to accommodate masked image modelling.
2. We demonstrate that the classification accuracy of a Capsule Network improves when self-supervisel pretraining is introduced.
3. We achieve state-of-the-art performance multiple benchmark datasets for Capsule Networks, including those with realistically sized images where Capsule Networks typically perform poorly.
4. We implement fully capsule-based docoder layer, replacing the CNN decoders commonly used for reconstruction tasks Capsule Networks, ensuring that our proposed MCAE model does not require handcrafted decoder.
2

---

Published in Transactions on Machine Learning Research (01/2025)
5. We conduct the first investigation into the use of ViTs as replacement for the traditional convolu- tional stem.
The rest of this paper presents the necessary background on Capsule Networks, highlighting previous re- scarch that has inspired the work presented here. We then formally define our new self-supervised capsule formulation called Masked Capsule Autoencoders and present several experiments and ablation studies on benchmark datasets. We conclude the paper by highlighting the main advantages of our new methods, with some key takeaway messages and some future directions that could further support the future developments of large-scale self-supervised Capsule Network models.
While interest in Capsule Networks may have waned in recent years, we argue that this is largely due research being overly focused on routing algorithms, leading to incremental improvements on toy datasets without addressing fundamental challenges such handling complex images. This work represents a signifi- cant. departure from this trend by proposing novel training regime for Capsule Networks, analogous to the advancements seen in ViTs and CNNs. We posit that this approach could reignite interest in capsule net- works and encourage rescarchers to explore alternative training paradigms beyond the standard supervised learning framework.
2
Related Works
2.1 Capsule Networks
Capsule Networks are a variation CNNs, which replace scalar neurons with vector or matrix capsules and construct a parse tree, representing part-whole relationships within the network. Each type of capsule in layer of capsules can be thought of as representing a specific concept. the current level of the parse tree which is part of bigger concept. Capsules in deeper layers are closer to the final class label than capsules in shallower layers. The capsules in the lowest layer, known primary capsules, detect the most basic visual features. These basic features serve as building blocks that can be combined to form any of the final object classes, forming the foundation of the parse tree. Capsules in lower layers decide their contribution to capsules in higher layers through a process called routing.
Capsule Routing, in brief, non-linear, cluster-like like process that takes place between adjacent capsule layers. This part of the network has been the predominant research focus for state-of-the-art Capsule Networks, to find better more efficient methods of finding ways to decide the contribution of lower capsules to higher capsules. In brief, the purpose of capsule routing is to assign parl capsules i= 1, in layer l to object capsules j = I....... M in layer 1+1, by adjusting coupling coefficients 7 E RR N iteratively, where 0 I "hy I 1. These coupling coofficients have similaritics to an attention matrix (Vaswaní et al. 2017) which modulates the outputs as a weighted average of the inputs. For more information on the numerous routing algorithms proposed for Capsule Networks, please see here (Ribeiro et al. 2022).
Dynamic Routing Capsule Networks are the original Capsulc Network architecture, as described in (Sabour et al.) 2017). DR Caps employs a technique called dynamic routing to iteratively refine the connec- tions between capsules. This approach introduces the concept of coupling coefficients which represent the strength of each connection and updates them using softmax function to ensure that each capsule in a lower layer must split its contribution amongst capsules that it deems relevant in the higher layer. The update process relies on agreement. values calculated as the dot product between a lower-level capsule's output and a predicted output from a higher-level capsule. After a pre-determined number of iterations, the activation of each higher-level capsule is calculated as the weighted sun of the lower-level capsule activations, where the weights are the final coupling coefficients.
Self-Routing Capsule Networks (SR-CapsNet) Hahn et al. 2019 address the computational burden of iterative routing algorithms in traditional Capsule Networks by introducing an efficient, independent rout- approaches, SR-CapsNet employs a dedicated
ing mechanism. Unlike conventional routing network for each capsule to compute its coupling coefficients directly, rather than relying upon iter- ative agreement methods, drawing inspiration from mixture of experts networks Masoudnia & Ebrahimpour 22044.
3

---

Published in Transactions on Machine Learning Research (01/2025)
Let us formally define layer in SR-CapsNet. Given input feature map with batch size b, the activation tensor E RbxAxxx represents capsule activations across a feature map, where A is the number of input capsule types and h, w are spatial dimensions of the feature map. Each capsule also maintains a pose matrix pose E ACxhxx, where represents the dimensionality of each capsule's pose vector.
The self-routing mechanism is parameterised by two learnable weight matrices: Wroute E RRKAABxC for computing routing coefficients and Wpose Ig for pose transformation, where k is the convolu- tional kernel size (kk: k²), B is the number of output capsule types, and D is the output pose dimension.
The routing process begins by applying convolutional operation with kernel size k, stride 8, and padding to transform the imput poses into po E [gxxxkkAxxx1, where represents the resulting spatial dimension. The routing coefficients are then computed as:
softmax Wroute posewnfolded!
(1)
where r E Rbx
The output activations Aout B are computed by weighting the input activations with the routing coefficients:
Eien, a,r; Eign, a;
aout
(2)
where Ω4 represents the set of input capsules in the local receptive field and h', w/ are the spatial dimensions after convolution.
For pose transformation, the network computes:
Eien, r,a, (WPose ( pose,)
pose mut Σien, r,a;
(3)
where pose out E RDxBDDX'xxx xwe represents the transformed pose matrices.
This formulation enables each capsule to independently determine its routing coefficients through learned parameters while maintaining the ability to model part-whole relationships. The elimination of iterative routing significantly reduces computational complexity compared to Capsule Networks with iterative routing.
While SR Caps perform competitively on standard benchmarks, their reliance on pre-learned routing network parameters hinders the network's ability to dynamically adjust routing weights based on the specific input, a key advantage of agreement- based routing approaches.
The Stacked Capsule Autoencoder (SCAE) Kosiorek al. 2019 is another Capsule Network archi- tecture which does not use supervised learning. Comprising of Part Capsule Autoencoder (PCAF) and Object Capsule Autoencoder (OCAE), the SCAE segments images into parts and organises these parts into coherent objects without labelled data. The model's unsupervised learning mechanism maximises the likelihood of reconstructing both the input image and inferred part poses, subject to sparsity constraints. By leveraging explicit geometric relationships between objects and their parts, SCAE achieves improved gener- alisation and viewpoint robustness. This approach demonstrates competitive performance in unsupervised classification tasks, particularly on MNIST and SVHN datasets, without relying on mutual information- based techniques or data augmentation. Consequently, SCAE represents substantial paradigm shift from the conventional Capsule Network training paradigm.
2.2 Masked Autoencoders
Masked Autoencoders (He al. 2021 are specific variant of ViTs which are pretrained via patch-specific reconstruction loss, tasking the network to reconstruct masked patches based upon the information which
4

---

Published in Transactions on Machine Learning Research (01/2025)
can be learnt from the visible patches, this can be visually in figure An image is first split into N x N patches of equal size and are flattened, allowing for the tokenisation of an image akin to text in a standard transformer [Vaswani et al. 2017). To mask specific areas (patches) in an image, tokens are randomly selected up to a specified percentage of the total tokens and removed from the sequence. This process eliminates the information from the feature map. The remaining visible patches are subsequently processed via a VIT. After the encoder has finished processing the visible patches, the masked patches are reinserted into the sequence, replacing the previously removed patches. The network now employs a ViT decoder to make predictions for these masked tokens utilising the attention mechanism and multilayer perceptrons within the standard VIT blocks. This process requires the network to learn how local areas correspond to their neighbouring patches by predicting the removed patches.
3 Masked Capsule Autoencoders
To create the MCAE we must first define how Capsule Networks can have their feature maps masked. CNNs, this task challenging and typically achieved by sctting specific regions of the feature map to zero. However, this method differs from the masked autoencoder (He al., 2021). Zero masking has been demonstrated alter the distribution of pixels in an image (Ballaubraamanian & Feizi 2023), which turn affects the performance. Therefore, in the subsequent section, will discuss the modifications have implemented to enable accurale masking within our MCAE.
3.1 Flattened Feature Map
patch feature map
capsule feature map
height
height
width
width
1d patch feature map
length
1d capsule feature map
length
visual representation of how a 2D patch feature map or capsule feature map with height and width is flattened into 1D feature map with length instead. At each location, there the same amount of different capsule types, each corresponding to different part concept in the part-whole parse tree. The dog image used is sourced from the Imagewoof validation set (Howard. 2019b).
Figure 2:
Vision Transformers can easily perform masking on a feature map, as patches of the image can removed from computation by simply removing selected patches from the flattened sequence of patches after the patch embedding layer. Capsule Networks on the other hand have traditionally used a 2D feature map, which comes with the drawback that masking can only be achieved either via replacing masked regions with
5

---

Published in Transactions on Machine Learning Research (01/2025)
Remove Patches
Patchify
Flatten
(Mask)
Figure 3: A visual representation of the masking process. An image is split into non-overlapping patches of N x N pixels. Randomly, a percentage, in this case, 50% of patches are removed to deprive the network of information available in these patches. The patches are then flattened into a 1D sequence of the remaining patches, ready to be processed by our encoder. The dog image used is sourced from the Imagewool valicattio set (Howard 2019b).
O's or utilising sparse operations (Woo et al. 2023 which come with their own drawbacks Balasubramanian & Feizi
2023
Tian et al. 2023
Thus we propose that by flattening the 2D feature map into a ID feature map, mimicking the design of ViT feature map, masking can be achieved in the same way as in the Masked Autoencoder (He et al., 2021). We thus achieve masking by simply removing all capsules at a specific location along the length dimension of our feature map.
3.2 Building Upon Self Routing Capsule Networks
We use SR Caps Network Hahn et al. (2019) its routing architecture to better suit our masked image modelling approach. The key architectural change lies in how we handle spatial relationships in routing.
starting point due to its simplicity and speed, but modify
The original SR Caps model uses a sliding window approach where each output capsule considers inputs from a local region & x of the feature map. In contrast, our MCAE uses point-wise routing where capsules only route the layer above at their exact spatial position in a flattened ID feature map shown in This creates local parse tree for cach patch position, which aims to determine the concept represented in the patch.
During pretraining, we do not include the final class capsule layer that would normally be part of supervised Capsule Network. Instead, after the encoding stage, we introduce a learnable masked token a shared vector of the same dimension as our capsule vectors. This masked token inserted all positions where patches were originally removed, returning the sequence its original length. Following the design principles of the original masked auloencoder He et al. (2021), this same learnable placeholder is used across all masked positions and is updated during training parameter of the model.
The decoder then uses a fully-connected capsule layer where cach capsule can route to all capsules in the subsequent layer regardless of spatial position. This is achieved by setting the kernel size to match the sequence length. While this global routing enables capsules at any position to influence the reconstruction of any masked patch, it comes with significantly incrcased computational cost compared to the position-wise routing used the encoder. However, this computational overhead only affects the pretraining phase. Once pretraining is complete, the decoder is discarded. This design choice trades higher pretraining computation

---

Published in Transactions on Machine Learning Research (01/2025)
for improved network accuracy during inference, where the network also maintains the fast inference speed of SR Caps.
When we switch to finetuning, we remove this decoder architecture and replace it with class capsule layer. This new layer averages the activations each capsule type across all spatial positions to make class predictions, is standard in the self routing training procedure. However, this layer now able leverage the improved representations required for reconstruction learned during pretraining.
3.3 Loss Function
Extract Reconstructions and Targets from masked patches only
Prediction
Full Target
Masked Prediction
Masked Target
Reconstructed Unmasked Token
Reconstructed Masked Token
Figure 4: A visual representation of how our pretrain loss function selects patches for the loss function defined in equation The dog image used is sourced from the Imagewoof validation set (Howard 2019b).
A crucial aspect of the pretraining stage of the MCAE involves training the network to accurately reconstruct the masked portions of the input image. To achieve this, we use the Mean Squared Error (MSE) loss, which quantifies the difference between the actual pixel values of the masked patches and the predicted pixel values generated by the capsule decoder. MSE loss is defined as:
1 MSE
Σ yt 9
N
where represents the total number of pixels across all masked patches the training batch, Yi is the actual value the 7th pixel the masked patch, and yi denotes the predicted value from our capsule decoder for the same pixel. A visual representation of patch selection from the target and prediction can be seen in figure
The MSE loss aligus with our objective to minimize the difference between the reconstructed and original patches, ensuring precise prediction of masked patch pixel values by the capsule decoder. It accentuates larger discrepancies by squaring errors, thereby pushing the model to improve on significant deviations and enhance reconstruction on cach masked patch.
When finetuning for classífication, the MSF loss is replaced with the cross entropy (CF) loss, defined by:
CE
Vio log(tie)
(5)
where N the number of samples, C the number of classes, lic indicates if class c is correct for sample i, and the is the average activation for each class of sample This loss encourages the model to activate the correct class capsules with high confidence.

---

Published in Transactions on Machine Learning Research (01/2025)
3.4 Backbone Selection
To ensure that information is completely masked out, replace standard ResNet (He el. al. 2016) ConvNet backbone with a ConvMixer (Trockman & Kolter 2022). This architecture's first layer uses a kernel size and stride of equal size, known as a patch embedding layer, allowing our feature map to contain no overlapping information. This ensures that when regions of the image are masked, information cannot be leaked through the overlapping sliding convolutional kernel.
We also provide set of architectures with a ViT backbone. This is achieved by setting the dimension of each token's is representation to Number of Primary Capsules Primary Capsule Embedding Dimension allowing for easy reshape into the primary capsule tensor dimensions. To create the activations for the primary capsules, we use simple linear layer with sigmoid activation to ensure that the value of the activation remains between 0 and
4 Experiments
Pretrain
Finetune
Pretrained Feature Extractor
Pretrained Capsule Encoder
Feature
Capsule
Capsule Decoder
Class Capsules
Extractor
Encoder
Figure 5: A visual depiction of the pretrain and finetuning components. We show how the feature extracting CNN and capsule encoder are kept. from the pretrain to finetune step. The capsule decoder is discarded after protraining and replaced with a class capsules layer that maps the capsule encoder network to a classification output.
To validate the effectiveness of our method, conducted numerous experiments with various ablations multiple datascts. These experiments demonstrate that masking is indeed a powerful technique for pushing the boundaries of Capsule Networks.
4.1 Experimental Setup
All of our experiments follow the same experimental setup. This involves lo optionally pretrain the network minus the class capsules for 50 epochs, with 50% of patches removed on either 'removed patch' or 'whole image roconstruction' as target. We then add the class capsules our network and fully finctune it for 350 epochs, following the supervised training settings of (Hahn al. 2019 Everett et al. 2023). visual depiction of the elements of the components of pretraining and finetuning be found in figure 5 All models use the SGD optimizer with default settings and the cosine annealing learning rate scheduler with 0.1 initial learning rate.
When a validation dataset has not been predefined, we randomly split 10% of the training dataset to act as our validation dataset. The best model is tested once on the test set our datasets, with the best model being chosen based on the epoch with the lowest validation loss.
4.2 Datasets
We validate our results on multiple datasets. For all of our benchmark datasets, we use the augmentation strategy proposed in (Everett et. al. 2023), which aligns with the augmentations used in other capsule

---

Published in Transactions on Machine Learning Research (01/2025)
papers, are the first to provide results Imagenette, we define the augmentations to be the same the augmentations for Imagewool.
Initially, we provide sanity check on the MNIST dataset (LeCun et al., 2010 to provide quick experi- mentation to ensure that our methods work at all. Next, we use both the FashionMNIST and CIFAR-10 datasets (Xiao et al. 2017 Krizhevsky et al.] 2200),, two datasets which are well within the abilities of a standard Capsule Network and allow us ensure that we are not limited to the simplest of experiments. For these two datasets, the augmentation that
is random resized crop of 0.8 to 1.0 during training. The SmalINORB dataset (LeCun et al., 2004) allows to ensure that we maintaining the equivari- ant propertics and generalisation abilities of Capsule Networks as the test set is specifically chosen to vary substantially from the train set while remaining within a similar distribution. In addition to standard classi-
fication accuracy on the Sma dataset, we also follow (Hahn et al., 2019 Ribeiro et al., 2020 Hinton ct al. 2018) and test our model on the novel azimuth and elevation tasks to verify generalisation capabilities.
The augmentations that we use for this dataset are that we standardise and take random 32x32 crops during
training. At. Lest. time, we centre crop the images to 32x32 defined in (Ribeiro al. 2020). Finally, we use the Imagenette and Imagewoof datasets (Howard) 2019a b to test network's performance larger, more realistic datascts. Imagenette and Imagewoof take 10 different classes from the Imagenct datasct (Deng et al. 2009). Imagenette is designed to be easily differentiable and tests our network's 's ability to process larger; more complex images. Imagewoof is ten classes of dogs and is designed to be more difficult to dif- ferentiate between classes due to the highly overlapping shared features. For Imagewoof and Imagenette we also resize the images to 64x64 for computational efficiency. We again use random resize crop and colour jitter for these datasets.
4.3 Results
4.3.1
Results on Image Classification:
CIFAR10 Validation Loss Comparison
CIFAR10 Validation Accuracy Comparison
3.0
CMAE No
CMAE No
2.5
80
2.0
Accuracy
70
1.5
60-
1.0
50-
200
Validation Epoch
Validation Epoch
Figure 6: Graphs comparing MCAE with vs MCAE w/o the pretraining stage. The graphs show the validation loss (left) and validation accuracy (right) on CIFAR10 over 350 epochs.
Jabbleppresennn the classification results of key state-of-the-art Capsule Networks compared our approach with no pretraining and with pretraining the datasets proposed in experimental design. MCAE with no pretraining is architecturally similar to SR Caps Networks, but with the ID modification to the feature map and 1 x kernels, along with the other required changes to the computation to allow for this. This method produces better results compared to SR-Caps, but does not. achieve stale-of-the-art in any dataset. However, when we apply the masked pretraining paradigm, our results improve on all datasets except MNIST. This pushes the MCAE with pretraining to be statc-of-thc-art for Capsule Networks in all datascts except SmaIINORB, which still dominated by Capsule Network models based on iterative routing methods. In addition to Capsule Network models, have trained three baseline models using the exact

---

Published in Transactions on Machine Learning Research (01/2025)
Table I: Main results for a number of foundational Capsule Network models compared to both the MCAE with masked pretraining and without. We show results on the four datasets that Capsule Networks are traditionally benchmarked on, as well as providing results for the Imagenette and Imagewoof datasets which are subsets of the Imagenet datasct. Unfortunately, it is computationally infcasible to train DR, EM or VB Caps these larger datasets due to their heavy VRAM requirements. Results for similarly sized CNNs and ViTs have been included using the results from Everett et al. (023)) and additional Imagenette results using our experimental settings for Imagewoof for fair comparison.
MNIST
FasHionMNIST CIFAR-10 SmalINORB
Imagenette
Imagewoof
ResNet18 He
al.
2016
99.1
89.3
78.4
89.8
68.9
52.7
MobileNetv3
Howard et al.
2019
99.5
98.7
77.1
83.3
62.4
45.9
ViT-S Dosovitskiy et al. 2020
99.4
99.3
80.0
91.1
70.1
57.3
DR Caps
(Sabour al.
2017
99.5
82.5
91.4"
97.31
EM Caps
Hinton et al.
2018
99.47
87.5
Caps
Ribeiro et
2020
99.7!
94.8
88.9
98.5!
SR Caps
Hahn et al.
2019
92.2
ProtoCaps (Everett et al.
2023
99.5
92.5
87.1
94.4
74.4
59.0!
ResNet20 SR*
(Hahn al.
2019
90.0
SR Cups SA* **
HHan et al.
2019
99.6*
91.5t
82.7
92
45. 45.2**
32.51
MCAE no PT
99.6
92.1
91.9
93.1
73.1
55.9
MCAE
99.6
95.0
92.8
95.0
82.1
61.8
*The ResNet20-SR results are the performance of the ResNet20 archilecture from the Self Routing Capsule Networks poper Hahn et al. (2019) which uses more augmentations, such as a random horizontal flip and normalisation, us to show the performance difference that including augmentations has on the convolutional baselines. **The SR Caps SA (Some Augmentations) results are taken from the ProtoCaps poper Everett et al. 2023) and show the performance of the SR Caps model using the same augmentations as us, note that the only overlapping dataset is Cifar10. ***The Imagenetle results for SR Caps are trained using the exact experimental sellings defined in ProloCaps for their Imagewoof resulls. i denoles to denole any resulls which were taken from their original papers and thus may not use our exact experimental setting, however, care has been taken follow conventions of training, as defined in section 4.2.
same experimental setting as our MCAE with no pretraining regarding epochs, hyperparameters and dataset augmentations. These models were trained using the Pytorch Image Library (TIMM) Wightman (2019) and are ResNet18 (He et al. 2016), MobileNetv3 (Howard et al. 2019) and ViT-S (Dosovitskiy al. 2020
4.3.2 Backbone Choice:
Leveraging ConvMixer backbone (Trockman & Kolter 2022 aligns with models' requirement of patch embedding layer to provide non-overlapping patches of the image. ConvMixer's feature maps are by default patchified, while ViIs (Dosovitskly et al. 2020) utilise a patch embedding layer. Prompted by this similarity, we explored this as an ablation study. Our observation reveals that ViT-based models underperform compared to those employing a convolutional backbone. Although ViT models yielded better performance than vanilla ViTs on smaller datasets, such as CIFAR-TO or SmallNORR, the overall results, shown in table 2 suggest that ConvMixers offer more suitable architecture for the MCAE and thus have been used for our results in 1
4.3.3 Reconstruction Target:
While the masked autocncoder (ee et al. 2021) framework that we build upon only roconstructs masked paiches, we also provide results where the reconstruction objective includes visible patches. Reconstructing based upon the whole image is inspired by DR Caps (Sabour et al. 2017) using a full image reconstruction objective along with the classification objective in order regularise the network. The results are shown in table 3 and show that reconstructing masked patches the best method, with reconstructing all patches providing significantly worse results.
10

---

Published in Transactions on Machine Learning Research (01/2025)
Table 2: Results of experimentation with a ViT (Dosovitskiy et al., 2020 with depth 4 backbone compared to Capsule Networks with standard CNN backbone. The specific CNN which we use is ConvMixer (Trockman & Kolter 2022) of depth due its easily scalable esoteric design being based on the presumption that the image has been patchified, ensuring no information leakage of masked regions due sliding window of overlapping convolutional kernels. In addition, show the mFLOPS (Total Floating Point Operations 10") to process one image through the backbone, showing that the ConvMixer is significantly more efficient. FLOPss are calculated using the FVCore library
FAIR
2023
Vision Transformer
Conv
mFLOPS
mFLOPS
Mixer
MNIST
80
99.6
0.5
99.6
FashionMNIST
80
91.1
0.5
95.0
SmallNORB
110
91.4
2
95.0
C1EAR-10
110
90.3
3
92.8
Imagenette
110
68.4
13
82.1
Imagewoof
110
55.4
13
61.8
Table 3: This table compares performance across our target datasets for MCAE pretraining based upon reconstructing both visible and masked patches those focusing on masked patches only, Results show equal or superior performance for models reconstructing masked patches only, across all datasets.
Visible and
Masked
Masked Patches
Patches Only
MNIST
99.6
99.6
PNSH
88.4
95.0
SmalINORB
82.0
CIFAR-10
84.8
92.8
Imagenette
45.1
82.1
Imagewoof
32.5
61.8
4.3.4 SmalINORB Novel Viewpoint:
Table 4: Comparing novel viewpoint generalisation on the SmalINORB novel azimuth and elevation tasks
[eCCnn et al. 2004). Results for DR, EM and SR Caps are from Hahn et al. 2019) and results for VB Caps are taken from Ribeiro et al. 2020).
Azimuth
Elevation
Familiar
Novel
Familiar
Novel
DR Caps
93.1
79.7
94.2
83.6
EM Caps
92.6
79.8
94.0
82.5
VB Caps
96.3
88.7
95.7
88.4
SR Caps 92.4
80.1
94.0
84.1
MCAE
93.2
85.6
95.3
86.1
In order to verify that retain the novel viewpoint generalisation capabilities of Capsule Networks, the novel azimuth and clevation tasks of the SmalINORB datasct. We replicate the experimental design of (Hahn et al. 2019 Hinton el al. 2018) and conduct two experiments. 1) Training only on azimuths in (300, 320, 340, 0, 20, 40) and test on azimuths the range 60 to 280. 2) Training on the elevations in (30, 35, 40) degrecs from horizontal and then testing on elevations the range of 45 to degrocs. In table we compare our accuracy on the test set on both the seen and unseen viewpoints. We pretrain for 50 epochs and finetune for 350 epochs, the same best model for SmalINORB in table
11

---

Published in Transactions on Machine Learning Research (01/2025)
We do not achieve state-of-the-arrs results this task, but do outperform all Capsule Networks except for VB Caps Ribeiro et al. 2020), showing that masked pretraining does not remove the generalisation capabilities of our network.
5
Conclusion
We have proposed the Masked Capsule Autoencoder model, the first capsule architecture trained in self- supervised manner. This could be step change in the development of scalable Capsule Network models. Extensive experiments demonstrate that MCAE outperforms other Capsule Network architectures al- most all datasets, with particularly favourable results on higher-resolution images. Considering the unique and well-establiseed advantages that Capsule Networks have in capturing viewpoint equivariance and view- point invariance compared with Transformers and CNNs (Riboiro et al. 2022). our model step towards developing large and scalable Capsule Network models.
We would consider the drawbacks of our method to be in the fully capsule decoder. In the Masked Autoen- coder paper (He et al. 2021) they state that the pretraining loss continued to decrease at the point which they stopped pretraining at 1600 epochs. While our reconstruction loss plateaus much quicker, to the point where it does not decrease any further after the 50 epochs for which we pretrain, this indicates that there is a point at which our model has reached the best reconstructions it can achieve. While we have shown that the pretraining stage improves the maximum classification accuracy for all datasets except MNIST (due to very fine margins for quantifiable improvement); an improved decoding mechanism can be found ben- efit from additional masked pretraining, the peak classification accuracy could likely be higher. Moreover, the decoder is computationally intensive because it needs to consider the entire feature map, significantly increasing training time and VRAM requirements compared to when no finet uning is used.
Acknowledgments
We would like to thank the University of Aberdeen's High Performance Computing facility for enabling this work and the anonymous reviewers for their constructive feedback.
References
Sriram Balasubramanian and Soheil Feizi. Towards improved input masking for convolutional neural net- works, 2023.
Fabio De Sousa Ribeiro, Georgios Leontidis, and Stefanos Kollias. Introducing routing uncertainty in capsule networks. Advances in Neural Information Processing Systems, 33:6490-6502 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. 2009 IEFE conference on computer vision and pottern recognition, pp. 248-255 Icco, 2009.
Jacob Devlin, Ming- Men Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirec- tional transformers for language understanding. CoRR, abs/1810.04805, 2018. URL http://arxiv..o org/
04805
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
16x16 words:
Miles Anthony Everett, Mingjun Zhong, and Georgios Leontidis. Protocaps: A fast and non-iterative capsule network routing method. Transactions on Machine Learning Research, 2023.
FAIR. fvcore: A light core library for computer vision, 2023. URL https://github.com com/
facebokreesarcc/Erccor.
12

---

Published in Transactions on Machine Learning Research (01/2025)
Taeyoung Hahn, Myeongjang Pyeon, and Gunhee Kim. Self-routing capsule networks. Advonces in neural information processing systems, 2019.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual lcarning for image recognition. In Proceedings of the IEEE conference on compuler vision and paltern recognilion, pp. 770-778 778, 2016.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners, 2021.
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. In International conference on learning representations, 2018.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chien Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Proceedings of the TEEE/CVP international conference computer vision, pp. 1314-1324, 2019.
Jeremy Howard. Imagenette: A smaller subset of 10 easily classified classes from imagenet, March 2019a.
URL htttps::/githbb.coott/t http://uuw.coom/ffsttaaaaaaaaans
Jeremy Howard. Imagewoof: a subset of 10 classes from imagenct that aren't so easy to classify, March
2019b. URL httppe//gggtubb..oo//aaaaaaaaats
Adam R. Kosiorek, Sara Sabour, Yee Whye Teh, and Geoffrey E. Hinton. Stacked capsule autoencoders, 2019. URL /tttps://axxiv..... hwwpaaawww.....gbbbaaaa/eeea http://www.rr//abb//......222 19tph//99960081811
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citesoer, 2009.
Yann LeCun, Fu Jie Huang, Leon Bottou, et al. Learning methods for generic object recognition with invariance to pose and lighting. In CVPR (2), pp. 97 104. Citeseer, 2004.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Avoilable: http://yann.lecun/ wwwpaanneeeuu.oom. coom/erdb/mmnn 2, 2010.
Saeed Masoudnia and Reza Ebrahimpour. Mixture of experts: literature survey. The Artificial Intelligence
Review, 42(2):275, 2014.
Proceedings of the AAAI Conference on Arlificial Intelligence, volume 34, pp. 3749-37-39 -7066,3766--60 2020. Fabio De Sousa Ribciro, Georgios Lcontidis, and Stefanos Kollias. Capsule routing via variational bayes. In
Fabio De Sousa Ribeiro, Kevin Duarte, Miles Everctt, Georgios Lcontidis, and Mubarak Shah. Learning with capsules: survey, arXiv preprint arXiv:2206.0260 2022.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. Advances in neural information processing systems, 30, 2017.
Keyu Tian, Yi Jiang, Qishuai Diao, Chen Lin, Liwei Wang, and Zehuan Yuan, Designing bert for convolu- tional networks: Sparse and hierarchical masked modeling, 2023.
Asher Trockman and Zico Kolter. Patches are all you need?, 2022. URL https://openreviee..ess
forum?id=TvHs5yYalee
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1706.03762,, 2017. URL
http:
//arxiv.org/artt wwwriv..org/bb//t/i 1006. 03762
Ross Wightman. Pytorch image models. https://github. http:://ggghuub.cc////rggaaaaaap http://tw.cccc//..igttaaaa//torr http::////.ttttio.oiiiiaaaaaaaaa http://yyooorc--mmmmm-mde......
2019. Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie.
Convnext v2: Co-designing and scaling convnets with masked autoencoders, 2023.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arxiv preprint 108.0777087777 2017.
13

---

Published in Transactions on Machine Learning Research (01/2025)
A Appendix
A.1 Decoder Ablations
Table 5: Table showing ablation results MLP and Convolutional decoder our MCAE model. Trained under the exact same experimental results as our MCAE model in [ MLP models are linear layer decoders
of dimension num num_caps * num dim * w, num caps * num dim h * w * 2, ieaai__eaaaaaeaaeeaa channels imgg h img w. Each layer has ReLU activation function. Convolutional models specified in Woo et al. 2023).
ConvNextvv2SS decoder
MNIST
FashionMNIST
CIFAR-10
SmalINORB
Imagenette
Imagewoof
MCAE Conv Decoder MCAE MLP Decoder
99.6
91.3
92.1
93.5
77.3
57.2
99.6
94.7
92.5
92.4
77.1
58.9
14

---

Published in Transactions on Machine Learning Research (01/2025)
A.2 Effect of Different Components in MCAE
MINIST
Pashio
Cuallaai Acstec
01.5
For teats Reornmartio
Road
Intontractinn
vi Bashane
HCd
200
SmaiNORB
citar10
Oasitaate Acturne3 110
65.0
Aecompnccinn
Bestoote
Aur Reconttiniio
No
wt Bestom
RCM
Imagenette
Imageweel
Costoat Acaato
33
40
Hoat
Figure 7: Graphs depicting how the top accuracy changes based on different ablations of the MCAE per dataset. Full Image Reconstruction refers to ConvMixer backbone MCAE pretrained for 50 epochs full image roconstruction. No PT refers to ConvMixer backbone MCAE with no pretraining epochs. VIT Backbone refers to ViT backbone MCAE pretrained for 50 epochs masked patch reconstruction. MCAE refers to our best-performing model which utilises a ConvMixer backbone and masked patch reconstruction. All models use the same lincar SR Caps model which contains 3 layers, with 16 Capsules per layer and are finetuned for 350 epochs.
15


Reviews:
[{"review_id": "Ukvaf93UUm", "reviewer": "Reviewer_iLVd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n\n* The paper has interesting ideas combining self-supervised learning with capsule networks. \n* The paper archives competitive results as compared to other capsule network architectures.\n* While a lot of algorithms have been proposed on how to improve the routing algorithm, this paper focuses on better and more improved training strategies.\n\n\n## Weaknesses:\n\n* The paper skips a lot of technical details regarding the actual architecture so it's quite difficult to judge the correctness of the architecture. Imho for acceptance, some parts of the paper requires a rewrite and explain the technical details \n* Some of the experiment results also need to be clarified.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces ideas from self-supervised learning as in masked autoencoders to the capsule architecture. They train an encoder-decoder architecture with patches as input on a reconstruction loss. The major claim of the paper is that this unsupervised training + supervised finetuning can unlock capsule networks for high-resolution images and complex datasets. Results are shown comparing to other Capsule Network architectures on traditional small datasets (MNIST, CIFAR-10, SmallNORB, FashionMNIST) and more complex datasets such as ImageNette and ImageWoof", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Here a list of required changes which are needed for acceptance.\n\n## Technical detials\n\n* The paper builds on top of Self-Routing Capsule Networks but it is not self-contained. The last paragraph of Section 2.1 is dedicated to this but I had to refer to (Hahn et. al 2019) to understand. The paper can use precise mathematical notation to describe the inputs, outputs, pose matrix and route matrix, activations with the precise dimensions. A good inspiration for this is the second paragraph of Section 4.2 in (Hahn et.al 2019).\n* In Section 3.2, the paper dives deeply into the intricacies of their modifications without describing the overall architecture first. The paper can first start with an overall description of the architecture used in Hahn et al, and then describe their proposed modifications to their architecture. Concretely, there are terms like “rather than merging local capsules within a HxW sliding kernel, we simply use a 1x1 region and route the capsules to the upper location in the same region.”, which are difficult to parse.\n* Why did the paper start with SR CAPS instead of ProtoCAPS, since ProtoCAPS has better performance? This decision needs a bit more justification.\n* “Resinsert a masked capsule placeholder” – What is the precise representation of the mask?\n* The major idea of the margin-based loss is that the length of the vector can encode the probability of an object being present and is more aligned with capsule networks. See Section 3 of https://arxiv.org/pdf/1710.09829. The paper replaces the standard margin loss used in capsule networks with a cross-entropy loss. What is the rationale behind this?\n* ConvMixer is composed of depthwise (sliding) convolutions at every layer. So the inputs to these are assumed to be spatially congruent, that is a 3-D Tensor (HWC). Throwing away patch information as done in Figure 1, removes this spatial congruence. So it’s quite unclear how authors integrated masking into the ConvMixer architecture. \n* The architecture of the capsule decoder is not clearly explained in the paper.\n\n## Experiments:\n\n* What is the resolution of ImageNette and ImageWolf? Is it 224x224\n* The ablation in Table 2 is quite difficult to assess without some numbers without  some sort of control such as wall-time or flops.\nThe paper claims that (MCAE no PT) is very similar to SR CAPS apart from the pretraining strategy. So it makes sense that the results between MCAE no PT and SR CAPS very similar on MNIST, Fashion MNIST, CIF. However, there exists a huge performance difference on ImageNette and ImageWoof between MCAE no PT and SR CAPS. Why is the case?\n* Please bold the best results in all tables.\n* The paper states that “but also a 9% improvement compared to purely supervised training”. The claim should be more specific. On which dataset and comparing to which “supervised training”?\n\n## Writing\n\n\nI found some sentences to be unreadable and unnecessary complex. Please split them into smaller, legible sentences.\n\n* Each of these capsules has a corresponding activation value between 0 and 1 which represents how strongly the network believes the concept which the capsule represents is present at the location in the feature map.\n* The capsules in the lowest layer can be thought of as the most basic parts which could be a part of any of the end classes, thus are denoted as the primary capsules, signifying that they are the base parts of the parse tree.\n* This feature map which now contains both encoded capsule representations and a random noise-masked capsule representation is now fed through a capsule layer which considers all capsules at all locations in the lower layer when creating the pose vector and activation values of all capsules at all locations in the higher layer, meaning that the encoded capsules can predict the values of the masked regions.", "broader_impact_concerns": "no", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "E6tCCSbS5o", "reviewer": "Reviewer_QHq8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n\nWhile existing work focuses on the routing mechanisms of capsule networks, this paper introduces MIM into the training pipeline, aiming to enhance scalability to more complex datasets.\n\n\n**Weakness/Clarifications:**\n\n\nThe authors state that capsule networks struggle to scale to more complex datasets, even with fewer parameters. However, why not increase the number of parameters to match larger models? Would that solve the scaling issue?\n\nAdditionally, the authors claim that MIM aims to address capsule networks' weaknesses by improving local representations. It is unclear how improving local representations resolves the challenge of scaling to larger datasets.\n\nA major concern relates to evaluation: ResNet18 achieves only 78.4% on CIFAR-10 and 68.9% on Imagenette. These results seem unusually low compared to expected performance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes incorporating masked image modeling (MIM) into the training pipeline of capsule networks. Comparisons with various capsule networks on image datasets show that MIM results in modest gains for small datasets like MNIST and CIFAR-10, while leading to significant improvements on two 10-class subsets of the ImageNet dataset.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Since a key metric for evaluating the proposed method and claims in this paper is image classification performance, would it be possible for the authors to provide the training and test loss/accuracy trajectories over the course of training? This would significantly aid in evaluating the results and supporting the claims made in this paper.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qpBttUY2ZT", "reviewer": "Reviewer_xJ4j", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThis paper does a great job integrating promising techniques and modeling approaches from across ML and bring them to bear toward the improvement of CapsNets. This is they type of synthesis research that we need more of in our community.\n\nThe claim of Masked-image modeling through pre-training providing a strong improvement over prior CapsNet papers is verified in the results presented in Table 1. There is some lingering concern that MCAE isn't a total improvement (considering the SmallNORB results) but the authors do a great job discussing this and analyzing performance on SmallNORB in Section 4.3.4. Overall, I was pleased with the experimental comparisons and ablations performed in this paper. I feel that they are sufficient to demonstrate the claims made about the MCAE improving over prior CapsNet approaches.\n\nI appreciated the work done to validate the architecture choices made with the MCAE through the ablation study performed through Section 4. As a result, I have increased confidence in the quality of the results and feel that the paper presents a nice foundation for future work to build from.\n\n## Weaknesses\n\nIn Section 4.3.2 and Figure 6 It would be clearer if it were made more explicit that “MCAE” is built upon ConvMixer. The current sentence “results suggest that ConvMixers offer a more suitable architecure from MCAE” could be re-written into something like, “Our default setting for MCAE uses ConvMixer as a backbone. We consider an alternative construction where ViTs are used instead as a potential ablation…”\n\nTheres is no reference to Table 2 in the text… Also, the information in tables 2 and 3 are pretty clearly identified from Figure 6. I’m not certain you need all three of these in the paper.\n\nThere are no direct comparisons to CNN or ViT based architectures.\n\nThere is little done to investigate the choice of using a \"fully capsule decoder\". This is discussed as drawback/weakness of the MCAE in the conclusion section. I feel that there is more to be done to justify this choice whether through an additional ablation or more focused writing/discussion beyond the single sentence at the end of the Introduction section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Masked Capsule Autoencoders (MCAE) that utilize pretraining with masked image modeling, a self-supervised learning technique, for later finetuning for image classification. It is demonstrated that pretraining a capsule network (CapsNet) in this manner enables better performance among similar CapsNet architectures. In addition to this, the MCAE is able to perform on larger, \"more realistic\", images where previous CapsNet architectures fail due to memory overhead. Finally, this paper investigates various architecture and training configurations of the MCAE to clearly establish which components are most beneficial for the presented results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I want to first explain my \"overall\" ratings below and at the outset say that I am very willing to change them throughout the review process of this paper. I will walk through my high-level reasons for saying \"no\" (at this point) for whether this paper satisfies the \"Claims and Evidence\" and \"Audience\" requirements of TMLR, before addressing specific changes that I feel that authors should make. Most of these opinions that I write here are not held strongly and I'm looking forward to discussing with the authors to come to an agreement. I am very comfortable with and want to change my assessment along these lines.\n\n### Claims and Evidence\n\nThe authors make an erroneous claim at the outset of the paper and in the abstract where they state that the MCAE is the first paper to use self-supervised learning with CapsNets. There have been previous works looking into the use of unsupervised/self-supervised learning with CapsNets. I feel that the major contribution in this paper is the use of Masked Image Modeling, similar to ViT architectures. This is a neat idea within the context of CapsNets and seems to differentiate it from the prior literature. However, I think that it's important that a more thorough discussion of CapsNet architectures is needed in this paper. Specifically this paper needs to discuss how it differs from work like \"Stacked Capsule Autoencoders\" from Kosiorek, et al in NeurIPS 2019. \n\nAn additional concern of mine along the \"Claims and Evidence\" dimension is that the paper claims that MCAE can compete on equal terms with Transformers and CNNs on the included image classification benchmarks. But there are no included direct comparisons with any architectures or published methods to this end. Without empirical evidence to support this claim, I would recommend removing this statement (which would significantly affect it's \"audience\" estimation). Alternatively, it is necessary to include additional experimental results demonstrating this comparison.\n\n### Audience\n\nI am uncertain whether there is a sufficient level of interest in CapsNet architectures among the TMLR community which would warrant there being an \"audience\" for this paper. This is perhaps my weakest held opinion. If there was an intermediary option between \"yes\" and \"no\" for this evaluation I would've chosen it with a slight lean toward \"yes\". The work presented in this paper is, from my estimation, good and interesting and certainly further advances the CapsNet architecture. I just worry that there really isn't much of a community that is building upon or furthering the use of CapsNets. Does this mean that research into their use and making them better is unnecessary? Certainly not. I am uncertain however that the paper as currently written will appeal to a broader audience. I think that this could be changed by a more thorough discussion of why CapsNets are appealing alternatives to standard CNN or Transformer architectures. There is an initial attempt at this in the first paragraph of the paper but it felt somewhat insufficient to wholly motivate why continued research in CapsNets is necessary. This could also be included in the concluding section where a more thorough discussion of what advances are enabled by better performing CapsNet architectures may provide over current standards.\n\nSome questions to consider: \n> A major concern with CapsNets is that they are slow and inefficient models with a high computational overhead. How do the authors justify this? Do they demonstrate significant improvements qualitatively and quantitatively that would assuage concerns about this overhead?\n\n## Requested changes\n - Adjust the claims made about the MCAE being the first/only unsupervised/self-supervised pretraining method for CapsNets\n - Include experimental comparisons with CNN/Transformer-based approaches.\n - Fix reference in text for Tables 2 and 3 if they are to be retained the paper. \n - Clarify what the final/intended configuration of MCAE to make the ablations more easy to follow/understand.\n - Improve discussion around the choice for the fully capsule decoder.\n - Improve discussion about how MCAE and CapsNets in general are situated in the overall ML landscape.", "broader_impact_concerns": "None that I can think of beyond the standard ethical practice of how image classification may be used to unfairly discriminate.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The idea is interesting enough to deserve publication, even if there remains a bit of skepticism on the reviewer's side.
Decision: Accept (poster)