Reconciling feature sharing and multiple predictions with   MIMO Vision Transformers
Authors: Rémy Sun
Clément Masson
Nicolas THOME
Matthieu Cord
Venue: Submitted to ICLR 2023 (2023)
Keywords: Deep learning
Computer vision
Vision transformers
Classification
Multi-input multi-output

Abstract: Multi-input multi-output training improves network performance by optimizing multiple subnetworks simultaneously. In this paper, we propose MixViT, the first MIMO framework for vision transformers that takes advantage of ViTs’ innate mechanisms to share features between subnetworks. This is in stark contrast to traditional MIMO CNNs that are limited by their inability to mutualize features. Unlike them, MixViT only separates subnetworks in the last layers thanks to a novel source attribution that ties tokens to specific subnetworks. As such, we retain the benefits of multi-output supervision while training strong features useful to both subnetworks. We verify MixViT leads to significant gains across multiple architectures (ConViT, CaiT) and datasets (CIFAR, TinyImageNet, ImageNet-100, and ImageNet-1k) by fitting multiple subnetworks at the end of a base model.

Under review as a conference paper at ICLR 2023
RECONCILING FEATURE SHARING AND MULTIPLE PRE- DICTIONS WITH MIMO VISION TRANSFORMERS
Anonymous authors Paper under double-blind review
ABSTRACT
Multi-input multi-output training improves network performance by optimizing multiple subnetworks simultaneously. In this paper, we propose MixViT, the first MIMO framework for vision transformers that takes advantage of ViTs' innate mechanisms share features between subnetworks. This is stark contrast traditional MIMO CNNs that are limited by their inability to mutualize features. Unlike them, Mix ViT only separates subnetworks the last layers thanks to a novel source attribution that ties tokens to specific subnetworks. As such, we retain the benefits of multi-output supervision while training strong features useful to both subnetworks. We verify Mix leads to significant gains across multiple architectures ViT, CaiT) and datasets (CIFAR, TinylmageNet, ImageNet- 100, and ImageNet- Ik) by fitting multiple subnetworks at the end of a base model.
INTRODUCTION
Training deep architectures has become commonplace in modern machine learning applications (Krizhevsky et al. 2012 He et al.) 2016) over the course of the last decade. Finding ways to better train these models has therefore become a question of paramount importance, and thus led to significant work in the literature (Zhang et al. 2018 (Shanmugam al., (022])
(MIMO) architectures provide promising (Havasi et al. 2021 Rame et al. 202T) technique to maximize model performance, These methods train 2 concurrent subnetworks within a base network by mixing 2 inputs into single one, extracting shared feature representation with core network and retrieving predictions, one for each input. At inference, we use the same input twice and get 2 predictions. This strongly benefits model performance both by emulating ensembling (Lakshminarayanan et al., 2017)) [Hansen & Salamon] 1990) with subnetwork predictions (Havasi | 2021) and by implementing partieular form (Rame al. 2021 [un al. 2022) of mixing samples ls data augmentation (MSDA) (Zhang et al.) 2018) Yun et al. (019)).
Multi-input
While one would hope the resulting subnetworks cooperate to learn and share strong generie features, this usually is not the case (Rame et al. 2021): if a feature is used by one subnetwork, the other does it (see the visualization given for reference at the bottom of Fig. []. This in turn causes issues for MIMO networks like forcing the use of large base networks. Although (Snn et al. (2022) shows this separation issue stems from difficulty in reconciling the shared representations with the need for distinct predictions, this is not a trivial problem to overcome in existing MIMO CNNs.
Vision Transformers (Dosovitskiy al. 202T Touvron al. 202 la) (ViT) offer a new solution to this problem with their token-basee representation and attention mechanism that could address these issues. Remarkably, MIMO networks' success has until now remained confined to Convolutionnal Neural Networks and the paradigm has yet to extend to these emerging vision transformers. This is all the more notable as they have started outperforming CNNs on vision benchmarks (Touvron et al. 202 Liu al. 202 and seem well suited to train more efficient MIMO transformers.
In this paper, we propose the first VIT based MIMO framework, ViT, that significantly improves upon the performance of standard single-input single-output models at minimal cost. Mix modifies the traditional MIMO structure to take full advantage of ViTs' propensity to mutualize features between subnetworks while still retaining the advantage of training distinct predictions (see Fig. [ therefore avoids common MIMO issues, and even benefits from strong implicit regularization (due feature sharing) that proves particularly useful smaller datasets. At training,

---

Under review as a conference paper at ICLR 2023
Shared Salt Attention blocks xL
Source attributin Attentior
Multi-oupuu supervision
Subnetwork
Subnetwork
Feature Visualization
MIMO CNN Reference
Classical MIMO problem
Figure I: MixViT trains efficient MIMO VITs. We aim mutualize features between subnetworks as in the visualization bubble (instead of disjoint features like in classical MIMO) by reconciling the need predict distinct outputs with the need for features to describe both inputs.
Mix mixes patch tokens from inputs before feeding them shared Self-Attentin blocks. To help extract 2 predictions from these shared features, we then introduce "source attribution" to associate each attended token the input/subnetwork they pertain and aggregate features with class token. Finally, we get 2 predictions from the class token features with 2 different dense layers. At inference, Mix ViT feeds unique input the model and retrieves two predictions thanks source attribution.
Our contributions are therefore as follows: 1) We leverage innate properties of Vision Transformers train efficient subnetworks that share features needed. ViT therefore adds only marginal costs at training and inference while strongly regularizing the model. 2) We introduce a novel source attribution mechanism to facilitate this late separation of subnetworks. 3) We propose the first working MIMO ViT as Mix solves issue of traditional MIMO frameworks on transformers.
Mix VIT sets new srateooftthe-as on the Tiny ImageNet dataset and shows strong benefits across multiple architectures (ConViT VIT (D' Ascoli et al. 2021), CaiT (Touvron et al. (021bb) and datasets (CIFAR 10/100 (Krizhevsky et al., 2009), Tiny IinylmageNer (Chrabaszez al. 2017). ImageNet- (Tian et al. 2020)), and ImageNet- 1k (Deng et al. 2009).
2 MIXVIT
We propose ViX (Fig. 2). MIMO framework for vision transformers that takes advantage of ViTs' innate properties to share features between subnetworks while still yielding diverse predictions. Mix adapts the classical MIMO workflow to ViTs by completely restructuring the model that subnetworks share early features and only specialize in the last layers thanks to our novel "source attribution" mechanism which allows the emergence of independent subnetworks in the last layers. This section starts by introducing Mix (Sec. 2.1 before discussing Mix ViT's s new MIMO structure in Sec. 2.2 our new source attribution (Sec. 2.3 mechanism, and added overhead in Sec. 2.4
ViTs are seen here separated into feature extraction blocks that use Self-Attenti blocks and classifier blocks that rely on Class-Attention blocks (see Fig. 2 number of modern vision transformers follow this structure like the CaiT (Touvron et al., 2021b) architecture. Mix is in itself ceatteccttrr-anossti and casily generalizes to other ViT architectures with minor adaptations.
2.1 VIX OVERVIEW
At training time, the N patches extracted from the inputs are mixed into N patches with binary masks [M,)] (so one token per position is picked shown on Fig. 2 with ratio À 3(a, a) (following a corrected CutMix (Yun et al.] 2019) scheme here). The N mixed patches are then encoded into d-dimensional tokens by linear layer e. Learnable positional embeddings
added to the tokens and the tokens are fed to the L Aiff-Atteniion blocks that serve as feature extraction blocks. This yields N attended tokens but indication which input/subnetwork each token belongs
We propose a novel source attribution mechanism s; to specify which input/subnetwork each patch pertains to. The source attribution adds subnetwork-speech information the tokens elaborate
2

---

Under review as a conference paper at ICLR 2023
dg
d
Co
dg
1-A
Fox
shared wsights
Classifier blocks
Atts--tttntionn
Co
attribution
Inference
Feature blocks Self-Attention
XL
Pos eerbelding
À (cxxoo)) Drose
Shared
M,
X
1 (A)
M.
1-A
Mixing
x
Mask generation
Training
Figure Overview of our MixViT framework. At training, we insert mixing and a source attribution step in the normal flow of the network. Source attribution associates tokens with the relevant input/subnetwork, and can be done through either source encoding or source embedding. At inference, we perform multiple passes the last layer using differently sourced tokens.
on in Sec. 2.3 The sourced tokens t* are passed along with learnable classification token co to the remaining classifier blocks. After the class-attention blocks, retrieve the classification token's features for classification. 2 predictions are then obtained from this feature representation through the of M different dense layers d, We then use the input targets yı to optimize the subnetworks' cross-entropy losses weighted by the mixing ratio À following standard MIMO practices:
LMIVVT = X
X (CE X- Po) + (1 X) (cE((.
(1)
At inference, MixViT only has to perform feature extraction on the image's patch tokens once as shown Fig. 2 mixing is needed since the tokens subnetwork-agnostic. We then collect predictions for each subnetwork i by using to specialize the general feature patch tokens for the subnetwork through source attribution, and then passing these specialized tokens the classification blocks along with ep (see Fig. 2). This new "parallel" inference scheme is made possible by Mix ViT's novel structure and solves issues with MIMO ViTs as discussed in Sec. 2.2
2.2 STRUCTURING MIXVIT FACILITATE FEATURE SHARING BETWEEN SUBNETWORKS
We structure Mix VIT to capitalize on ViTs' propensity to train subnetworks that share features while still retaining benefits training with multiple outputs. This restructuration classical MIMO hinges our "source attribution" mechanism (see Sec. 2.3). and explicitly separates the model between a core feature extraction block shared between subnetworks and specialized sub-classifiers.
feature extractor shared by all subnetworks. It takes a single (mixed) input - with no distinction between subnetworks and outputs single set of feature tokens: there classifier blocks take two source-attie aurree-atributed inputs (see Sec. 2..) and output two predictions.
More preciscly, deconstruct the traditional MIMO structure by training
subnetworks in the MIMO sense. Only the last
This structure both accentuates ViT's natural tendency share features and enables to parallel inference scheme. First, as all subnetworks must work from the same pool of feature tokens (after source attribution), they incentivized make of many useful features possible.
3

---

Under review as a conference paper at ICLR 2023
Moreover, by forcing the feature extraction blocks to process all patch tokens together give no information on their source we learn subnetwork agnostic features.
Secondly, this structure also improves the way the input is processed at inference which is long standing issue of MIMO frameworks. Indeed, Mix uses new "parallel" inference scheme instead of just summing two differently encoded versions of the input: after extracting generic features and attributing sources, we perform one evaluation per set of source attributed tokens "in parallel". This is fairly inexpensive since computations only have be performed once (see Sec. 2.4) but would be prohibitively expensive in a standard MIMO framework. Parallel inference is much preferable to the previous "summing" scheme it forced the use of sub-optimal summing during training and caused significant interference between subnetworks when they share features (Sun et al. 2022). The latter proves particularly problematic in ViTs (shown experimentally in Sec.
2.3 SOURCE ATTRIBUTION MECHANISMS
We move the subnetworks to the end of vision transformers by introducing source attribution mechanisms s; after the feature extraction blocks to specialize the generic patch tokens into subnetwork-specife tokens (*, As discussed by Rame et al. (2021), distinguishing inputs is neces- sary for subnetworks to work on their designated input. Standard MIMO frameworks rely on separate input encoders at the beginning of the network (see Fig. 3 to address this. Our source attribution mechanism extends this mechanism for ViTs and can be deployed anywhere the network.
t'
So
5o
S1
(a) Source encoding.
(b) Source embedding.
Figure 3: Source attribution helps the model identify information relevant to each input.
Source attribution is most naturally implemented with subnetwork-specifie source encoders (see Fig. 3a
We define source encoders s; such that Intuitively, each patch token is projected by the relevant source encoder 8;. In practice, we project the feature tokens with the Linear layers s; into specialized tokens, and mix these specialized tokens following the masks M, used on the inputs. The result is added to the original residual token to obtain subnetwork-speeifie token t* figured in Eq. ( is the Hadamard product):
Linear(W? b() are Linear layers.
= Σ M, =t Σ M. +))). (=0
We also propose an alternative lightweight source embedding mechanism (see Fig. 3bb to avoid incurring complications from training additional Linear layers (see Eq. Source embeddings 8; act similarly to positional embeddings, and are d-dimensional vectors that are added to feature tokens to specify which subnetwork the patch is relevant to. This therefore only adds 2d parameters and O(Nd) operations which is negligible in ViTs. Interestingly, source embeddings are particular case of source encodings where the weight matrix is null such that 8; = Linear (0axé, by) and
ヒ Σ M. M.
Σ Mi o (tx oxxd +b)).
(3)
2.4 REGARDING COMPUTATIONAL OVERHEAD
Deferring the separation of subnetworks to the last layers allows us to perform separate evaluations for each subnetwork at very low computational cost if we want train M subnetworks (usually
4

---

Under review as a conference paper at ICLR 2023
2). Consider transformer C-class problem with L self-attention blocks and 2 class-attention blocks taking N d-dimensional patch tokens. As a rule of thumb, one MixViT forward roughly incurs 1-2M X many computations normal model's during inference, and 1 1x at training.
Indeed, the most intensive computations are tied to Linear d x d weight multiplications: d is usually much larger than N, or M. While source encoding does incur M such computations per token, each Self- Aiff-Attention layer incurs 6 per token and Class-Attentin slightly more than 3 per token:
for our small 7-layer Con ViT (see Sec. 4 with M = 2 this represents about 6x5+3x242 = 1.06x 6x5+3x2
multtplicaatioss. The amount of additional computations is therefore bounded by the ratio of passes through attention blocks. This approximation could easily be lowered however: it assumes class- attention is as expensive as self-attention which is very pessimistic (see Appendix for tighter approximations and asymptotic complexity).
RELATED WORK
MIMO frameworks as data augmentations MIMO CNNs can be seen in- mixing samples data augmentation (Rame et al., 2021]: they mix two imputs 2o and 2ı in the first network manifold with ratio À before (roughly) optimizing mixed loss Alg + (1 (1-38 A)11 MIMO models however remain fundamentally different as they supervise multiple outputs.
Nevertheless, this similarity is even in patch/pixel space on a corrected mixed input (similarly to the MixToken scheme (Jiang et al. 202T)). Mix VIT's feature extractor is wholly shared between subnetworks and seems to benefit from this. Furthermore, much of Mix ViT's improvements on the single input ViT backbone must be credited to the model having very strong individual subnetworks due to inherent regularization.
pronounced in Mix ViT mixing is performed directly
Mixing in networks Mix ViT builds upon MIMO techniques (Havasi et al. 2021). Beyond simply extending MIMO input encoders, our source attribution also plays a similar role to the "unmixing" strategy proposed by Sun et al. (2022) to allow subnetworks to share features. Contrarily to unmixing however, Mix VIT does not require extensive tuning and brings significant performance improvements.
Data multiplexing methods (Murahari et al. 2022) use multiple inputs/outputs to accelerate model throughput to allow faster training and inference at a modest cost in performance. While Mix ViT shares similar structure those methods, it is noteworthy that it strongly improves on the base model instead of degrading performance like data multiplexing methods.
Finally, it must be noted that ViX emulates some ideas from the ensembling literature (Lak shminarayanan et al. 2017 Lee et al. 2015). Indeed, finding ways to train and evaluate multiple models a discount is central question of the field (Lee et al., 2015 Wenzel et al. 2020). Common solutions often resolve to share a core network and train a few specialized neural layers at the end of the network (Lee et al. 2015). MixViT closely matches this strategy, with the distinction that add specialized source attribution mechanisms instead of learning different parameters.
4 EXPERIMENTS
We first show our proposed ViT framework strongly improves on whichever backbone it builds upon in Sec. 4.1 in smaller datasets, before demonstrating in Sec. 4.2 that Mix VIT significantly improves accuracy on the large scale ImageNet-Ik Ik (Deng et al. (009) dataset. We then investigate feature sharing in ViX (Sec, 4.3) and its implications. Finally, Sec. 4.4 shows Mix ViT performs much better than traditional MIMO frameworks on ViTs.
Setting In this paper, we study the CIFAR-10/100 (Krizhessky et al., 2009). Tiny (Chrabaszez et al. (017), ImageNet- 100 (Tian et al. 2020). and ImageNet-1e 1k (Deng et al.) 2009) datasets. Most experiments use a variant of the ViT ViT (D' Ascoli et al., 2021 architecture with
5 GPSA blocks and 2 CA blocks with dimension 384 and 12 heads (see Appendix). On the more complex datasets, we also use deeper CaiT (Touvron et al. (0211) models. Our transformers work on 8 x 8 64 patches for CIFAR and TinylmageNet, and 14 14 196 patches on ImageNet. We report the best epoch overall ensemble accuracy as mean + std over three seeded runs for the best of the two Mix VIT variants unless specified otherwise.
5

---

Under review as a conference paper at ICLR 2023
Table Comparison of Mix against similarly sized transformer results reported in the literature.
CIFARIOD
CIFARIO Accuracy (56
TinyImugeNet
ImageNet
Models
Params
Accuracy (86)
Accuracy (%)
Accuracy (%)
Reported results (from (.)
CCT-7/3x (Hassani (Hassani
2021))
4M
77.1
94.8
DelT:T Touvion al 20cla Zhang al.
2022
5M
67.5
88.4
PVT- Wang di 2041
(022))
13M
69.6
90.5
WT
Disivilskyy el al 2021 LLe et (021)
3M
73.8
93.6
57.1
19I-X 2023
7M
79.0
95.9
60.3
SL-PII-XS (exxssLLysse of (0231
9M
75.0
94.2
62.9
CaiT-XXS-2a Toulyon orai 2021b
Lec ct. 202
9M
76.9
94.9
64,4
SL-CaiT-XX5-27 Le 202T1
9M
80.3
95.8
67.1
Dytox Douillard di puse [onn training)
10M
76.0
79.1
Swin-TTaireral 202010] (Tib
(222
29M
82.7
T2T-V/=11 Suunnst 2021 Car 2021a)
22M
65.2
84.2
82.7
CvT:13 aral (Lirer
20M
73.5
CVT:13 + Lpstor Liu 2021
20M
74.5
90.3
86.1
Our experiments
MixWI (ours) with Con backbone
12M
82. 0.1
96.5 0.2
70.2
88.8 0.3
Implementation details We mix inputs using CutMix, and start training subnetworks individually
towards the end of training. We largely adapt the DeiT Touvrro et al. widely used in the literature (Zhang et al. 2022 Touvron al. 2021b Lee al.. (0221) (see
2021a)
training settings that
Appendix). Unless specified otherwise, we therefore train with the Adam
(Loshchiloy Hutter
2019) optimizer for 150 epochs with weight decay, RandErase (Zhong et al. 2020), AutoAugment
Cubuk et al. (2019). stochastic depth (Huang et al., (016)). label smoothing Szegedy et al. (2016) and Batch augmentations (Hoffer al. (2020). We also apply CutMix Yun al. (2019) and Mix Up
Zhang et al.) (2188) but find applying them simultaneously instead of alternatively yields better result. Finally, we keep the 0.001 learning rate but adopt a step decay instead of cosine.
4.1 MIXVIT SIGNIFICANTLY IMPROVES THE PERFORMANCE OF VISION TRANSFORMERS
We start by showing how our ViX Cin fares against the current transformer state-of of-te-o-thh-ee of-thh-artt art in Sec. 1.1.1 across 4 datasets, before demonstrating ViT consistently improves over single-input
oingle baselines across multiple backbones in Sec. 4.1.2
4.1.1 MIXVIT RESULTS RELATIVELY TO SIMILAR TRANSFORMERS
We provide here some context on the performance of our models with regard to the existing literature.
Tab. gives comparison to the performance of normally trained transformers of similar sizes. Our results over 150 training epochs should allow a mostly fair comparison: vision transformers are typically trained for comparable longer amounts of time (100 (Liu al. 202 200 Hassani et al. (02T)) or more (Zhang et al. 2022) epochs) on CIFAR, TinylmageNer and ImageNet-
CIFAR Tab. I shows Mix largely outperforms small ViTs. In fact, our MiX 83.4% accuracy on the extended 300 epochs setting (see Appendix E is the best VIT on CIFAR-100. Indeed, the closest the literature are 82.7% 7% for CCT-7/3x1 (Hassani et 2021) trained 5000 epochs and 82.6% for a 90M parameters NesT-B ZZhang et al. 2022) trained 300 epochs.
TinyImageNe Tab. demonstratee Mix ViT largely outperforms transformers on the more complex TinylmageNet dataset. As the best reported transformers on the dataset are in line with our single- input dingle-outur backbones, the gains on the sfate-of-the-ae tte--rtt can largely be attributed to Mix ViT. Our Ca Mix ViT even reaches 70.9% (see Sec. 4.1.2 which improves on the previous state- of-the-art held by MixMo's 's ResNet- 18-3 model (70.2 accuracy), to the best of our knowledge.
ImageNet-100 Tab. shows Mix VIT outperforms the competing architectures for comparable training times. This suggests Mix ViT has little trouble accommodating larger 224 x 224 images.
4.1.2 MIXVIT CONSISTENTLY IMPROVES MODELS ACROSS BACKBONES AND DATASETS
Tab. 2 shows both of our VIT variants provide significant improvements over their single- input single-output counterpart across multitude of architectures and datasets. For instance,

---

Under review as a conference paper at ICLR 2023
Table 2: Mix ViT significantly improves the performance vision transformers
(a) Tiny ImageNet
(b) ImageNet-100
Backbone
Method
Accuracy (%)
Backbone
Method
Accuracy (%)
Single-ipputoourput
65.4 +:
tnngle-inpu//oututt
86.6 0.2
Cin
Mix Vit-emhedding (ours)
70.0 0.1
Cin
Mix Viteembebgg (ours)
88.8 +
Mix Vit-encoding (ours)
70.2 0.2
Mix Vit-enccdg (ours)
88.6 + 0.5
Single-input/ottpu
66.2 + 0.3
Single-inputootppu
85.8 = 0.2
CaiT XXS-24
CaiT XXS-24
Mix Vit-cmhedding (ours)
68.6 + 0.3
Mix Vit-embeddingg (ours)
87.2 + 0.3
Mix Vit-encoding (ours)
69.6 = 0.2
Mix Vit-encodinn (ours)
86.7 + 0.2
Single-input/ourpur
68.6 + 0.1
Single-input/oung Mix Vit-embedding (ours)
87.6 0.4
CaiT
CaiT XS-24
XS-24
Mix MixVit-embedding (ours)
69.4 0.1
89.7 + 0.3
MixVit-encoding (ours)
70.9 0.2
Mix MixVit-encoding (ours)
89.3 1 0.2
(c) CIFAR-100
(d) CIFAR-10
Backhone
Method
Accuracy (%)
Backbone
Method
Accuracy (%)
tunttl--pput/output Mix ViT-emmeedding (ours)
79.5 + 0.1
Single-input/outper Mix MixViT-emmeedding (ours)
96.1 + 0.1
ConViT
82.4 0.1
ConViT
96. 5 = 0.2
Mix ViT-encoding (ours)
82.1
Mix MixViT-encoding (ours)
96.5 = 0.2
TinyImageNet, our larger CaiT-XS model goes from baseline of 68.6% to 70.9% accuracy with Mix VIT. Interestingly, much of the gains are due to higher accuracy of the individual subnetworks: the framework also has strong regularizing effect, possibly due to the feature extractor being shared. This is particularly interesting considering ViTs' notorious difficulties with smaller datasets.
Our lightweight Mix ViT-embedding outperforms the complete Mix ViT-encoding CIFAR (Tabs. 2c and (d). Mix ViT-encoding however overtakes on the more complex IinylmageNet (Tab. 2a) Interestingly, this difference deepened by both model depth (our CiT has 7 layers against our CaiTs' 26) and model width (Cait-XXS's 192 against Cait-XS's s 288), As such, source embeddings are much faster and easier to train, but source encodings are more powerful overall.
4.2 PUSHING MODEL LARGE-SCALE DATASETS: IMAGENET-IK
Table 3: ImageNet- 1k scores.
Backbone
Method
Accuracy (%)
Tab. shows Mix VIT still improves ViT' on the large- scale ImageNet- Ik dataset (Deng et al. 2009): simply training subnetworks with Mix ViT improves a CaiT- M-24's performance by +0.8% which is significant given how challenging improving beyond 82% on ImageNet-Ik. 1k. This compares favorably with results reported by traditional MIMO frameworks.
ResNet 50
Single-input/ourpuut
76.1
MIMO (Havasi al.
2021
77.5
ResNet 18-7
Single-inpat/ourpur
77.2
MixMo (Rame et al.
2021
78.5
CaiT S-24
Single-input/ouppl
82.3
MixVit (ours)
82.8
CaiT M-24
Single-ipputoutput
82.7
Mix (ours)
83.5
4.3 FEATURE SHARING IN MIXVIT
We now verify that subnetworks do share most features in Mix in Sec. 4.3. before showing this
leads Vix accommodate simpler ViT training settings (see Sec. 4.3.2) and narrow networks (Sec. 4.3.3) contrarily to MIMO CNNs. Additionally, we show in Appendix Sec. B that Mix ViT is significant step towards training more than 2 subnetworks at the same time.
4.3.1 VISUALIZATION CLASSIFIER FEATURES
We check the influence of final features on subnetworks by considering the weight matrices W; e R# #Caassesx # Feats of our final classifiers d, (see Fig, 4). A feature's influence is approximated by the L1 norm of W, corresponding columm: this tallies how much a classifier depends on the feature. We then visualize whether the feature influences both subnetworks only on scatter plot.
Fig. shows Mix VIT features influence both subnetworks (close to the diagonal). On the contrary, most traditional MIMO features only influence one feature the other.
7

---

Under review as a conference paper at ICLR 2023
de SIS
L1 Influence d
MixMo features
MixviT features
Disjoint feats.
Shared featt.
L1 influence d,
Outgun
block
Figure 5: Visualization of final feature influence on subnet- works (by proxy on dg and dj) for MixMo (WideResNet-28-55) and Mix VIT (Con VIT), rescaled to range [0,1]. MixVil fea- tures influence both subnetworks (dark purple zone) contrarily to MixMo features.
Figure Feature influence on d, is given by L1 norm of classifier columns.
4.3.2 MIXVIT IS LESS RELIANT ON STRONG REGULARIZATION
Typical training schemes for vision transformers often require both heavy augmentation (Yun et al. 2019 Zhang et al. 2018) (in the form of mixing augmentations) and batch augmentations (Cubuk et al. 2019). This is concerning these procedures can significantly extend training costs and time.
Mix ViT is not as reliant as traditional ViTs on this strong regularization, possibly because feature sharing in ViT inherently provides some of the benefits of MSDA and batch augmentation. Indeed, we show mixed inputs to our shared feature extraction blocks which can emulate the effects of MSDA. Moreover, each sample in batch appears as part of an input pair: has be fed each of the subnetworks once. As such, each of the batch's 's sample is processed multiple times by the transformer's feature extraction blocks while mixed with other samples.
Tab. 4soows both Mix variants prove much more resilient than their single-input single-output backbone less regularized settings CIFAR-100. It is particularly interesting note that with no batch augmentations and MSDA augmentations, our Mix ViT-embeeddnn model matches the fully regularized variants of their single-input single-output backbone. This particularly important as batch augmentation linearly increases computations and has been a key fixture of vision transformer training since its introduction in DeiT Touvron al. 2021a). This linear cost has led some modern works to dispense with it (Liu et al. 2021b), but the procedure remains widely used in the literature (Zhang et al., 2022: Wang et al. (021). As such, MiX ability to provide strong performance with no batch augmentation could strongly simplify training schemes.
That is to say however that Mix not compatible with batch augmentation: Fig. 6 shows Mix still benefits significantly from it. As such, the regularization added by MixViT seems to complement regular batch augmentation well.
Table 4: CIFAR-100 accuracy with MSDA/BA.
34
Model
Strong reg.
Accuracy (%)
Single-input/output Con
72.3 + 0.3 79.5 + 0.1
Single-inputoottuu Con
Loss without strong regularization
-7.
Mix VIT (ours) MixVit (ours)
79.0 + 0.5
MOVIT Sum
82.4 + 0.1
IppacisENT Mage amaatcn
Loss without strong regularization
33.4
Figure 6: MixViT still scales well with batch augmentation.
8

---

Under review as a conference paper at ICLR 2023
4.3.3 MIXVIT ADAPTS TO SMALL MODELS BETTER THANN MIMO CNNS
Table 5: Mix outperforms CNN-based MixMo models on Tiny-ImageNet.
On the complex ImayIma dataset Mix ViT largely outperforms its much larger CNN-based competition (Tab. 5). Interestingly. this high- lights fundamental difference between tradi- tional MIMO models and our Mix VİT. MixMo learns subnetworks that share no features (Rame) et al. 2021] Sun al., (022) and therefore strug- gles to fit good subnetworks
Rackhone
Mathod
Parim
Width
Accimey (9)
ResNet 18
Single inpurooupuut
11M
512
65.1
MicMo Rame et al. 2021
1IM
512
64.4
ResNet 18-3
Single appulouren
100M
1536
69.0
100M
70.2
Singl--irqunkuuggu MiXVit (ours)
08.6++0..1
XS-34
+ 0.2
narrow architectures. As such, MixMo needs to consider wide ResNet-18-3 (He et al. 2016) backbone to start showing improvements. Mix VIT simply does not have this issue as the subnetworks share features and large part of the gains come from how the subnetworks learn features they can both use.
4.4 COMPARISON AGAINST VIT-BASED MIMO IMPLEMENTATIONS
Mix VIT is the first VIT based framework, and was developed account for the architecture's characteristics. Here, we provide comparisons against applications of previous MIMO methods to ViTs. While these methods were developed for CNN based frameworks, is possible to transpose the proposed frameworks to Vision Transformers with minimal alterations as detailed in Appendix C
Tab. shows our direct transpositions of the seminal MIMO and MixMo frameworks fail to perform for ViTs: both frameworks fail to even reach the same performance as the single-input singlc-output baseline which is consistent with observations in Allingham et al. (202T). MixMo performs much better with our parallel inference scheme, at the cost of significantly more overhead. Even better results are obtained if we only use CutMix mixing MixMo during training. This strongly suggests issues with sum mixing in MIMO ViTs, possibly due to VIT's propensity share features.
Mix ViT strongly improves over the base model with minimal overhead. Even without source attribution and therefore identical subnetworks, it brings strong improvements which clearly shows the benefits of multi-output supervision in a properly designed ViT. Adding source attribution, we recover distinct predictions and ensembling benefits reach our final 82.4% accuracy.
Table 6: CIFAR- 100 accuracy and estimated FLOPS of the backbone and MIMO formulations.
Backbone
Method
Inference Mixing
FLOPS
Accuracy (%)
Single-input
N/A
1x
79.5 + 0.1
MIMO
Sum
1x
77.1 0.1
MixMo
Sum
1x
77.0
Von
Mix Mo w/ parallel inference
Parallel
2x
78.5 0.2
Cut-MixMo only w/ parallel inf.
Parallel
2x
80.1 + 0.2
Mix VIT w/o source attribution (ours)
Parallel Parallel
1.1x
81.5
Mix (ours)
1.1x
82.4 1 0.1
DISCUSSION
We propose MixViT the first MIMO framework for ViTs that leads to significant improvements over the single-inpo input single-output baseline. We take advantage of ViT's innate properties to train subnetworks that share features, thereby solving a longstanding scaling issue in ViT' and adding implicit regularization between subnetworks. We achieve this by restructuring the standard MIMO framework and introducing a novel source attribution mechanism that helps identify input information.
The induced regularization and implicit ensembling obtained by fitting subnetworks leads MixViT to strongly improve on its base model across multiple architectures VCT, CaiT) and datasets (CIFAR, TinylmageNet, ImageNet- 100, and ImageNet-Ik). MixVIT models set a new state-of- the-art Tiny ImageNet and a new state-of-the -oce--of-hettttt art for transformers CIFAR-100. Furthermore, the combination of Mix ViT's feature sharing and multi-output supervision brings strong implicit regularization which can greatly simplify training settings.

---

Under review as a conference paper at ICLR 2023
Ethics statement At the least, Mix ViT aims to improve classification systems, say nothing of possible applications to object detection and other complex tasks. As such, must be acknowledged that it could lead to significant job automation and nefarious applications. While the implications of job automation can be debated to an extent, it will undeniably lead to a fundamental change in our society. It is also the unfortunate reality that deep learning frameworks he used by bad actors to harmful ends, with applications like Deepfakes already causing real problems today.
Nevertheless, it is our belief that these neural frameworks can and will have overall positive impact on the world by easing daily life and improving response rates to critical application. More precisely, we hope MixViT will serve to provide more reliable decision support systems in difficult applications so that human operators can take better and faster critical decisions.
Reproducibility statement We outline all significant components of the Mix ViT framework in Sec.2/ For contractual reasons it is not possible to give access to the full codebase, although this is still being discussed with the concerned organization at the moment. We have included detailed summary of experimental settings in Appendix H and given pseudo-code of the key components as well as indications concerning the codebases we adapted in Appendix
REFERENCES
James Urquhart Allingham, Florian Wenzel, Zelda Mariet, Basil Mustafa, Joan Puigcerver, Neil Houlsby, Ghassen Jerfel, Vincent Fortuin, Balaji Lakshminarayanan, Jasper Snoek, Dustin Tran, Carlos Riquelme Ruiz, and Rodolphe Jenatton. Sparse moes meet efficient ensembles. Arxiv preprint, 2021.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. ArXiv preprint, 2017.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation strategies from data. In Computer Vision and Pattern Recognition, 2019.
Stéphane D' D'Ascoli. Hugo Touvron, Matthew L Leavitt, Ari S Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. In International Conference on Machine Learning, 2021.
J. Deng. W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognition, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Arthur Douillard, Alexandre Ramé, Guillaume Couairon, and Matthieu Cord. Dytox: Transformers for continual learning with dynamic token expansion. Conference on Computer Vision and Pattern Recognition, 2022.
Diego Granziol, Stefan Zohren, and Stephen Roberts. Learning rates as function of batch size: A random matrix theory approach neural network training. Arxiv preprint, 2020.
Lars Kai Hansen and Peter Salamon. Neural network ensembles. IEEE transactions on pallern analysis and machine intelligence, 1990.
Ali Hassani, Steven Walton, Nikhil Shah, Abulikemu Abuduweili, Jiachen Li, and Humphrey Shi. Escaping the big data paradigm with compact transformers. In Learning on Limited and Imperfect Data Workshop at CVPR, 2021.
Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshmi- narayanan, Andrew Mingbo Dai, and Dustin Tran. Training independent subnetworks for robust prediction. In International Conference on Learning Representations, 2021.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European Conference on Computer Vision, 2016.
10

---

Under review as a conference paper at ICLR 2023
Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon
Oh. Rethinking spatial dimensions of vision transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), pp. 11936-11945, October 2021.
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. In Computer Vision and Pattern Recognition, 2020.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q. Weinberger. Deep networks with stochastic depth. In European Computer on Computer Vision, 2016.
Zi-Hang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Yujun Shi, Xiaojie Jin, Anran Wang, and Jiashi Feng. All tokens matter: Token labeling for training better vision transformers. In Advances in Neural Information Processing Systems, 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in Neural Information Processing Systems, 2012.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in neural information processing systems, 2017.
Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets. ArXiv preprint. 2021.
Stefan Lee, Senthil Purushwalkam, Michael Cogswell, David J. Crandall, and Dhruv Batra. Why M heads are better than one: Training diverse ensemble of deep networks. Arxiv preprint, 2015.
Yahui Liu, Enver Sangineto, Wei Bi, Nicu Sebe, Bruno Lepri, and Marco De Nadai. Efficient training of visual transformers with small datasets. In Advances in Neural Information Processing Systems, 2021a.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In International Conference on Computer Vision, 2021b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Vishvak Murahari, Carlos E. Jimenez, Runzhe Yang, and Karthik Narasimhan. Datamux: Data multiplexing for neural networks. Arxiv preprint, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Veto,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, gigh-pefformance deep learning library. In Advances in Neural Information Processing Systems, 2019.
Alexandre Rame, Remy Sun, and Matthieu Cord. Mixmo: Mixing multiple inputs for multiple outputs via deep subnetworks. In International Conference on Computer Vision, 2021.
Divya Shanmugam, Davis Blalock, Guha Balakrishnan, and John Guttag, Better aggregation in test-time augmentation. In International Conference on Computer Vision, 2021.
Remy Sun, Alexandre Rame, Clement Masson, Nicolas Thome, and Matthieu Cord. Towards efficient feature sharing in mimo architectures. Efficient Deep Learning for Computer Vision Workshop (ECV) CVPR, 2022.
C. Szegedy, V. Vanhoucke, S. loffe, J. Shlens, and Z. Wojna. Rethinking the inception architecture for computer vision. In Computer Vision and Pattern Recognition, 2016.
Yonglong Tian, Dilip Krishnan. and Phillip Isola. Contrastive multiview coding, In European Conference Computer Vision, 2020.
11

---

Under review as a conference paper at ICLR 2023
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers and distillation through attention. In International Conference Machine Learning, 2021a.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In International Conference on Computer Vision, 2021b.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In International Conference Computer Vision, 2021.
Florian Wenzel, Jasper Snoek, Dustin Tran, and Rodolphe Jenatton. Hyperparameter ensembles for robustness and uncertainty quantification. In Advances Neural Information Processing Systems, 2020.
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions to vision transformers. In International Conference on Computer Vision, 2021.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis E.H. Tay, Jiashi Feng. and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In International Conference on Computer Vision, 2021.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In International Conference on Computer Vision, 2019.
Sergey Zagoruyko and Nikos Komodakis. Conference, 2016.
Wide residual networks. In British Machine Vision
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, and Tomas Pfister. Aggregating nested transform- ers. AAAI, 2022.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmenta- tion. AAAI, 2020.
12

---

Under review as a conference paper at ICLR 2023
APPENDIX
This appendix contains the following information:
In Sec.
we provide further notes on the overhead incurred by MixViT.
In Sec.
show ViX suffers much less from training more than 2 subnetworks.
In Sec. C provide detailed descriptions of naive implementations for standard MIMO
methods on ViT's used in Sec.
In Sec.
D
we provide additional ablation experiments on MixVIT.
In Sec.
we provide additional comparison of Mix ViT with MixMo on the CIFAR-100 (which is much harder for vision transformers than CNNs).
dataset
In Sec.
we provide additional baseline comparisons for TinylmageNet and ImageNet- 100.
Sec.
we provide information on the licenses of some assets used in this paper.
In Sec.
we provide details regarding our experimental protocol.
In Sec.
provide information our codebase and pseudo-code.
13

---

Under review as a conference paper at ICLR 2023
A FURTHER NOTES ON MIX OVERHEAD
We provide in Sec. 2.4 a very rough upper bound on the additional overhead incurred by Mix at inference for one forward evaluation. We provide here more nuanced discussion on the computations at play and show that for M = or at least bounded - the asymptotic complexity in L, N and d of Mix is the same as that of a normal model.
Similarly to Sec. 2.4 we consider a transformer on C-class problem with L self-attention blocks and 2 Mass-attention blocks taking d-dimensional patch tokens. We neglect the computational overhead of induced by adding bias terms, embedding the patches into tokens, classifying from the classification tokens features, normalization and source attribution.
How many computations do class-attention and self-attention need We assumed previously class-attention and self-attention are as expensive computationally which leads to a very pessimistic upper bound on the overhead.
Let us count the important operations in a Self-Attention layer:
1. We have to extract key, query and value representations for each token which causes ddd operations due to 3N weight multiplications.
2. We have to compute the similarity scores between N tokens which means 2 operations.
3. We must compute the attended representation for all tokens which costs approximately 22d operations as attended token requires adding N d-dimensional tokens (we take the cost of multiplying the value representation by the attention weight to be atomic, but is technically d).
4. Most implementations include projection step properly aggregate the multiple attention heads which means an additional Nd cost.
5. Each token is then processed by the same hidden layer multi-layer perceptron. While the hidden dimension of the perceptron can vary it is standard in the literature to take 4d. Therefore we have to consider 2 x 4 x Nd operations as we consider 2 weights multiplications of dimensions d x 4d.
This yields a cost of approximately 3Nd N Nd = 13Nd +2N2 in Self- Attention layer whereas the cost of a Class-Attention layer is about 3Nd + Nd +d 8d (3N + 9)d + 2N d operations. Indeed, Class-Attentis follows a similar outline with the following differences: in 2) we only have Nd operations (similarity to the class token), in 3) we only compute one attended representation for cost Nd. in 4) we only compute the projection for one token at cost d² and in step 5) we only apply the MLP to one token for sdl operations.
If we throw out the terms linear in d (as it is usually the largest term by far), Class-Attention therefore 3N+9, x as many operations as Self-Antention layer. Note that to the best of our knowledge-
costs 13N
N 64 at the least in the literature, so this ratio is closer at maximum 0.25x.
Tighter approximation of MixViT's inference overhead We can now revise the very pessimistic approximation given Sec. 2.4 provided the estimations that Class-Attention costs most 0.25x as many computations as Self-Attention. One forward evaluation of Mix MixViT ViT costs approximately as many operations with this. For our modified Con this indicates a ratio of only 1.1x.
L+0.5MM L+0.5
Asymptotic complexity in L, N and d Asymptotically, for fixed number of subnetworks (M=2 usually), the complexity of MixViT and the underlying transformers are identical and equal to ( + N2d)). Indeed, as can be seen from the previous calculations, Self-Attention has a complexity of O(N2d + Ndd) and Class-Attention ONN2²). A normal forward pass is therefore on the order df + d2) as we have + 2 attention layers. Mix ViT-encoding induces M computations the same order in the form (NNd)) source encodings. Mix ViT-ennonodin complexity is therefore O(L+ dd + dN or O(L(Nd + N2d)) if we consider M bounded.
14

---

Under review as a conference paper at ICLR 2023
B MIXVIT ACCOMMODATES MULTIPLE SUBNETWORKS MUCH BETTER
The M>2 subnetworks case has been an open problem ever since the seminal MIMO (Havasi al. 2021) pa- per. We believe Mix VīT is a significant step towards solving this problem, which has been a pain point for all existing MIMO techniques (Havasi et al., 2021
Rame al.) 2021 Sun al., 2022). For instance, with single input baseline around 80% accuracy, a WedeRes MixMo network shows strong deterioration of performance over CIFAR 100 (batch repetition 2) whereas a ViT MixViT does not as shown Fig. 7
seems the feature sharing in VIT likely helps the network train more than independent subnet- works since our performance for M=4 is about 9% better with MixViT (over MixMo), and remains steady 3% above the single-input baseline. While this beyond the scope of this paper, we believe prop- erly designed MSDA scheme could allow Mix ViT to see further improvements with M>2 subnetworks.
Number sabnetwarts
Figure 7: MixVIT deals much better with more than 2 subnetworks.
As discussed in MixMo RRame et al. 2021). MIMO frameworks can generalize to more than 2 inputs by simply using a MSDA method that mixes more than 2 imputs. The training objective this case is simply the sum of the subnetwork losses ponderated by the (reweighted) mixing ratios. While such methods have been attempted the literature, they work fairly poorly for MIMO techniques.
A simple MSDA for this open problem could be to pick at random 2 inputs and perform cut-mixing on those two inputs (the mixing ratio for the other inputs would be 0). This is the scheme we use for the previous results on MixMo and Vi..
15

---

Under review as a conference paper at ICLR 2023
C
NAIVE IMPLEMENTATION OF MIMO AND MIXMO TRANSFORMERS
Feature blocks
Reature blocks
Classifter blocks
Clacsifler bloces
MIMO Transformer.
(b) MixMo Transformer
Figure 8: Transposition of traditional MIMO architectures to vision transformers.
At first blush, MIMO architectures easily be transposed vision transformers by analogy with convolutional neural networks. We outline below how vision transformers can replicate the steps of MIMO models, and provide illustration of MIMO transformers in Fig.
We consider M (inputs, labels) pairs ((r,, yi ) Josiem during training: we feed M inputs to the model as a combined input. To this end, we break down each image into small 4 patches and embed the patches into N 984-dimensional tokens using M linear embedding layers (one for each input). Corresponding tokens (at the same position the two images) are then mixed into one common set of tokens (either through sums Havasi et al. (202T) or other mixing schemes Rame et al. (202T)). The core network (feature blocks + classification blocks) then processes this representation until final attended classification token is obtained. MM dense layers {d: Jo:josi one for each subnetwork then yield M predictions from this attended classification. At test time, M predictions can be obtained for image by embedding the image patches with the M linear embeddings and summing the sets of embedded patch tokens.
In the seminal MIMO framework, mix the embedded tokens by simply summing them. The CutMix based MixMo formulation interpolates between the tokens following either a randomly drawn square mask with ratio À 3(a,a) a), linear scheme (similar summing).
Closs blocks
(a) Sum.
(b) Parallel (ours).
Figure 9: Inference schemes for MIMO architectures.
16

---

Under review as a conference paper at ICLR 2023
D FURTHER ABLATIONS
We verify in Tab. 7 the benefits of our source attribution design (Sec. D.I). our inference scheme (Sec. D.2 and our choice to mix inputs (Sec. D33 on CIFAR-100.
Table 7: CIFAR-100 accuracy and estimated
FLOPS for ablations on components and imple- mentation choices in the Mix framework. De-
fault settings are highlighted in grey for reference.
D.1 EFFECT OF SOURCE ATTRIBUTION
Method
ILOPS
Accuracy (%)
Tab. 7 shows MIMO training seems to inherently benefit the model: networks trained without source attribution
Importance of source attribution (Sec.
D.T
Late source embedding
x1.1
82.4+0.11
significantly better than standard CinViT networks. While such models learn completely identical subnetworks (as there no source attribution), having to learn features useful for the classification of multiple images simultaneously provides ample regularization.
Late source encoding
x1.1
82.1 0.1
source altribution
x1.1
81.5 $1.5.:0.2
Early input source embedding
x2
+2.2+00.2
Early source encoding
x2
82.2 0.2
Inference scheme (Sec.
D.2
Sepurate inference
x 1.1
0.14
Sum inference
x1
81.9 0.1
We separate subnetworks in the last layers in- stead of separately encoding the inputs early in the network avoid having to perform multi-
Allemating inference
82.0
Single subnetwork inferenee
x1
82.0 0.2
Mixing scheme (See.
D33
ple forward passes for inference (see Sec. [.4). We however find in Tab. 7 that displacing the separation back to early in the model slightly de- teriorates performance: the model benefits from wholly sharing feature blocks between subnetworks.
Binarized CutMix All tokens
x x1.
82.4 +
x2
81.5 + 0.2
D.2 INFERENCE SCHEME
In stark contrast to our findings in Sec. we find summing the 2 source attributed versions of each patch token at the end of our network before performing a single forward pass yields reasonably good results (Tab. 7). While these results are indeed worse than separate inferences, they are also far cry from the disastrous MixMo ViT results in Sec. We posit late patch representations are therefore significantly more suited than early patch representations to this sort of combination. Interestingly, we observe similar performance when drawing at random which patches should be attributed which subnetwork. also worth noting that the regularization induced by our Mix ViT scheme is strong enough for the performance of individual subnetworks to be reasonable on its own.
D.3 INPUT MIXING SCHEME
We initially chose to compress the inputs into a single set of N tokens through a CutMix scheme to avoid the quadratic cost that would come with using all M x N tokens from the M inputs. Tab. 7 actually indicates that far from deteriorating performance, our compressed scheme performs better than model trained with all tokens. We explain this by the fact the CutMix compression introduces additional regularization on the shared feature extractor and slightly more diverse subnetworks.
17

---

Under review as a conference paper at ICLR 2023
E FURTHER COMPARISON TO MIXMO ON CIFAR-100
Tab. 5 shows Mix matches the similarly sized WideResNett28-55 Zagoruyko & Ko- modakis 2016) on CIFAR- when given the same training budget. While MixMo's Rame
Table 8: Comparison between MixMo and ViT on CIFAR-100. indicates 300 training epochs with batch augmentations.
et al. (2021) reported WideRes.Net-28-5 perfor- mance better than our default setting Mix performance, MixMo models train over 300 epochs with batch repetition Batch repetition functions very similarly to Batch augmentation and is often needed to obtain good results with CNN-based MIMO models (which is less the case for Mix VIT as per Sec. 4.3.2). we reduce the learning rate and train with 4 batch augmen- tations, we close the gap in performance. This particularly noteworthy as vision transformers typically disadvantaged against ConvNets the CIFAR datasets as can be seen by the fact our baseline single-input single-output Con VīT model only reaches 81. 1.1% accuracy (vs. 82 22.%%) even on this longer training schedule.
Models
Params
Accuracy (%)
WideResNet-28
82.6
WedeResNNt MixMo"
83.3
ConViT D' Ascoli et (222))
12M
79.5
Vix Con
82.4
ConViT" D Ascoli et al. 2021 Mix Convil
81.4
12M
83.2 + 0.2
18

---

Under review as a conference paper at ICLR 2023
F FURTHER COMPARISONS ON TINYIMAGENET AND IMAGENET-100
We provide additional points of comparison for TinylmageNet and ImageNet- 100 as Tab. limited for these datasets.
1 iis fairly
Table 9: Comparison of MixViT against other transformers on TinylmageNet.
Tiny
Models
# Params
Accuracy (%)
Reported results (from [.])
ViT
Dosovitskiy al. (2021)
(Lee al.
2021))
3M
57.1
T2T
Yuan et al.
202T
[Lee et al.
2021
7M
60.6
PiT
Wang
2021
Lee al.
2021
7M
60.3
Swin
Liu et al.
2021b
(Lee et al.
2021
7M
60.9
Touvron et [202
(Lee et al.
2021
9M
64.4
SL-CaiT-XXXS-24 (Lee et al.
2021
9M
67.1
Our experiments
MiX
12M
70.2 + 0.2
Table 10: Comparison of Mix ViT against other transformers on ImageNet-106
ImageNet-100
Models
Params
Accuracy (%)
Reported results (from [.])
Dytox Douillard et al.
(022).
joint training)
10M
79.1
TiT-Vi1-14 Yuan et al.
2021
(Liu al.
202
22M
82.7
T2T-ViT-L 14 L Loc
Liu et (202
22M
83.7
Swin-1 Liu et al.
2021b)
(Liu et al. (202
29M
82.7
Swin-T LorLoc
(Liu al. (2021a
29M
84.0
CvT-13 Wu et al.
2021) (Liu et al.
2021a)
20M
85.6
CvT-13 LooLo
LLIU al. 2021a))
20M
86.1
Our experiments
MixVIT
26M
89.7
0.3
19

---

Under review as a conference paper at ICLR 2023
G INVENTORY OF LICENSED ASSETS USED IN THIS WORK
We built upon number of assets in this work, many of which are under open-source license. We provide below an inventory of the most important licensed assets we used:
We used the Python programming language under the PSF License agreement. We addition- ally made use of many python libraries, the most important of which is the pytorch library under BSD-3 Clause license.
We used a number of models and codebases, most notably the DeiT, Con ViT and timm codebases under Apache 2.0 license.
We used multiple datasets, among which are the CIFAR datasets under a MIT License and subset of the ImageNet dataset under a BSD-3 Clause license.
We used the strong data augmentation protocol AutoAugment under the MIT License.
H EXPERIMENTAL SETTING
We give here detailed inventory of hyperparameters and training procedure used in this paper. For contractual reasons, it is not possible to provide the codebase used in this paper, but Appendix ! contains pseudo-code and links to the public codebases we built upon.
All our experiments were run three times on three fixed seeds from the same codebase. Quantitative results are given in the form of mean std over the three runs.
FLOPS estimations We provide an estimate the relative number of FLOPS used by different methods in Tab. 6 and Tab. This number is traditionally computed for pytorch models on theoretical basis by standardized libraries like pt f flops or fvcore. These libraries however do not support custom attention layers and are therefore un un-usable for our purposes.
Fortunately, estimating the ratio of FLOPS is straightforward in Tab. 6 the additional overhead of MIMO and MixMo known be approximately 1x on CNNs (which also holds on transformers) and separate inference just leads to inferring twice, hence the 2x ratio. We use the upper bound derived in Appendix forr Mix and its variants that follow similar computational flows. For the last line of Tab. 7 we give rough approximation derived by considering the impact of the larger number tokens on the attention layers.
Corrected CutMix Scheme An issue that arises when considering a CutMix scheme to mix patch tokens that CutMix masks have binary values for each pixel value whereas a patch contains multiple patches. It stands to reason therefore that we will often have situations where the patch contains both masked pixels and unmasked pixels. While the natural solution simply to take a continuous mask value by averaging over the patch, we have seen in Sec. that subnetworks have trouble sharing tokens.
To address this, we correct the CutMix scheme by using majority vote within the patch: if most pixels are masked, then we mask the patch and vice-versa. Ties (where we have as many masked and unmasked pixels) are broken by choosing randomly for one set of inputs which input will be assigned the problematic patches. It worth noting that this scheme is for all intents and purposes the MixToken liang et al. (2021 1) scheme, with some practical differences in the sort of mask generated.
Architectural details We used a modified variant of the Con D' Ascoli et al. (202T) architecture for experiments hidden dimension 384, 12 attention heads, 5 GPSA blocks and 2 SA Blocks similarly to dytox Douillard et al. (2022). Patch size was taken to 4 x on CIFAR,
on TinyImageNe and 16 x 16 ImageNet- 100. It is worth noting GPSA blocks should in theory have a square number of heads to exactly emulate convolutional kernels at initialization, but we did not find this to be issue experimentally.
We took the CaiT models directly from the seminal paper Touvron et al. (2021b) and used the same patch size for our ViT model (which also coincides with CaiT's default settings).
20

---

Under review as a conference paper at ICLR 2023
Training details In general, train models over 150 epochs using the AdamW optimizer with a learning rate of 0,001, weight decay of 0.05 and linear warmup over 10% of training for most datasets. We follow step decay schedule with decay rate 0.1 at steps {125, 140} For the ImageNet dataset, we only trained for 133 epochs (CaiT-S-24, decay at steps (75,90) and 100 epochs (CaiTT-M-24, decay at steps (100, 120)) due to logistic constraints.
We used a base batch size B of 128 on CIFAR and TinylmageNet and 1048 on ImageNet. The effective batch size b used was dependant on memory constraints, and the learning rate was scaled
accordingly with ratio V / following heuristics derived from Granziol et al. (2020).
Following DeiT, we train with a lot of regularization. We use 0.05 weight decay. 0.05 stochastic depth for CaiT models, 0.1 label smoothing. Batch augmentations. We use the CIFAR configuration of Augment for CIFAR and Tiny ImageNet,, and the ImageNet configuration for ImageNet-1000 We also apply CutMix Yun et al. (2019) and MixUp Zhang et al. (2018) but find applying them simultaneously instead of alternatively yields better result and is more stable. Practically speaking, we draw CutMix mask with a 1.0 but instead of the mask having binary values (0 1), it has continuous values like a a 0.8 MixUp scheme.
Similarly MixMo, we progressively adapt the input mixing scheme reflect the mechanism used at inference. In the MixMo case, this means using summing scheme instead of CutMix with probability p that linearly grows over the last twelfth of training. the ViT case, we use masks that take all tokens from one of the inputs with probability that linearly grows over the last twelfth of training. We directly inherited the mechanism from MixMo and did not tune the point at which we start increasing the probability p linearly.
Choice of hyperparameters We tuned the parameter a E {0.25, 0.5, 1.0, used draw Mix ViT's input mixing CutMix mask on CIFAR-100 and kept it fixed for all other experiments. 0.5 was chosen as it was more stable and not necessarily for optimality purposes. We did not tune the root parameter used for loss balancing in MixMo and instead discarded it early on in our preliminary experiments (equivalent to choosing root 1).
While cosine decay is standard in the transformer literature, we failed to converge with it on CIFAR- 100 preliminary experiments. As we wished to ensure our single-input single-output backbone had converged for fair comparison with Mix we adopted step decay schedule and roughly set decay steps by considering the training loss on CIFAR- 100. Similarly, we train on 150 epochs instead of 100 because our models clearly failed to converge in 100 epochs on CIFAR-100.
We largely inherited all other parateters(patch sizes, batch sizes, 0.001 learning rate, 3 batch augmentations, 0.05 weight decay, AdamW optimizer, 0.05 stochastic depth dropout on CaiT, 0.1 label smoothing, CutMix a = 1.0, MixUp a = 0.8, RandErase and AutoAugment configurations) from the literature Touvron al. (222 Lee et al. (2021). We did double weight decay and stochastic depth on TinylmageNer and ImageNet- for MiXViT ViT to avoid overfitting and checked the performance of the backbone with those parameters, but did not look for an optimal value.
I ISEUDO-CODE AND CODEBASES
Our experiments were run through a codebase implemented using the PyTorch Paszke et al. (2019) library. Our codebase built upon the public official implementation of MixMo (atpps: //qithub com/ aalaaaleaaaeee / mi xmo-pytrrch) and adapted components from multiple public codebases:
We
took
Vin
building
blocks
from
https://githbb.co
We took CaiT building blocks from http::///gghhuu...cccaooooooo///
We
took
our AutoAugment implementation
from
https://github.com com/
We took a number of primitive functions from the timm library https://github. com/ rwi www.ghtmman/yytorchh www.:ehh-iiaaammmmmameeeee
21

---

Under review as a conference paper at ICLR 2023
The MixMo codebase in particular is well documented and should provide everything necessary to train multi-input multi-output networks. The main adaptation that must be made the addition the transformer networks and specific ViT routines in the transformer networks.
To facilitate this last point, we provide in Alg. T and Alg. 2 the training and inference pro- cedure for MixViT. PatchEmbed simply embeds the inputs into patch representations as standard in transformer frameworks, Mi Mixii thMakks simply mixes the inputs following the masks given, AddPosEmbed just adds the positional embeddings SourceAttributic implements Eq. 2 or Eq. 3 depending on which variant of MixViT is considered. SourceAtrr.butionni led is similar, except creates M sets of tokens, each equivalent to SourceAttrrbutionn one at position being a unit mask N.
the mixed tokens.
with T; being made up null masks On except for the
Algorithm 2: ViT forward inference.
Algorithm 1: MiX forward at training.
Input: Inputs
Input: Inputs zi. Masks M₁
Parameters:
Parameters:
Self-Attention blocks SAA.
Self-Attention blocks 5AAg blocks C'A.
Class-Attention blocks C.Ag
Classification token co
Classification token co
Dense classification layers d:
Dense classification layers dj
Result: Network predictions pi
Result: Network predictions Code:
Code:
Format the input;
# Format the input:
( } = {Pat ch Embed(x:)):
(rr) :
(Pat (.abed))mbed
No mixing needed
:= Mixwit t www.ithMaakeeo ((s)tttMaaas(()))aMMMM
X == Embe d(z;);
X = AddP os Embed
Standard forward evaluations;
# Standard forward evaluations;
for k=0,..,L--I do
for L= do
x: (A)
L x (x)(x);
Create sets of sourced tokens;
# Associate patches and subnetwork/input;
{ :=
X = SourceAtt r but ion
# Standard attentontnnn
10 # Standard class-attention:
11 x = Concat((cg;x)):
11 { := [Conca([[o;x::D]:
for k= do
12 for k=0, / do
13 x = (x2xx);;
13
(xx) == (C:)(z))):
14 Predict from features;
Predict from features;
15 co x(0]
15 (Co,s == [x;[0]))
16 [p.] = (co):
16 [p.] (d((co,)
Return (p;)
17 Return (p.]
22


Reviews:
[{"review_id": "fs2EPBucou", "reviewer": "Reviewer_sVL6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper considers the use of the multiple-input multiple-output (MIMO) paradigm for vision transformers. The authors propose to train a ViT by using two images A and B at a time – the two images are tokenized and the tokens are “mixed” by randomly selecting tokens from image A or B. The mixed tokens are then passed through the ViT, and a proposed “source attribution” layer at the top allows the network to predict the classes of both images A and B. Experimentally, the authors show that 1) their network beats previous small networks reported in the literature on small-scale datasets (table 1), 2) the proposed MIMO architecture gives improvements across datasets and architectures (table 2 and 3) and 3) that the proposed MIMO architecture beats previous MIMO architectures for ViTs (table 6).\n\n", "summary_of_the_review": "The authors present a natural extension of the MIMO framework for ViTs and show how this can improve performance for various vision datasets. While the empirical results are positive, it is not clear that all factors are controlled for and that the proposed method would actually improve SOTA models. Furthermore, the novelty is relatively low and the proposed method is not explained clearly.\n\n", "strength_and_weaknesses": "Strengths:\n\n* The proposed method generally gives solid gains in known benchmarks.\n* The authors give clear ablation experiments which show that their methods are better than other MIMO methods for VITs (table 6).\n\nWeaknesses:\n\n* The novelty is low.\n* Some results do not control for the training recipe, e.g. numbers in table 1 might have vastly different hyperparameters. \n* It is not clear if the method improves SOTA models.\n* The paper mostly focuses on small-scale datasets. It is known that VITs work best for large-scale datasets.\n* The clarity when explaining the proposed architecture could be improved.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fq4IMMTxJ3", "reviewer": "Reviewer_GaDC", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces a transformer-based MIMO framework, called MixViT. At training time, two subsets of tokens from two images are taken as input and the model is tasked to predict labels for the two input images. At test time, the same input is passed to the model multiple times with different sourced tokens. The features in the earlier layers are shared while two separate encoders, called source attribution, are introduced in the later layer for the input-dependent feature extraction. In this way, the model is expected to achieve feature sharing and separation for multi-prediction. Improved results are obtained on image classification tasks compared with the single input counterparts and baselines. ", "summary_of_the_review": "Overall, I think the proposed method sounds interesting and good performance is also achieved. However, the technical contribution is somewhat limited as discussed above, and the weaknesses overwhelm the strength. Hence, I lean to recommend rejection based on its current shape.", "strength_and_weaknesses": "Strength\n- Consistent improvements are achieved on the task of image classification.\n\nWeaknesses\n- The technical contribution seems limited. The input mixing for training is simply adopted from CutMix, and the source attribution appears to be quite straightforward. \n- The model needs to run multiple times for the inference, causing a significant increase in the computational cost.\n- The model is claimed to be multi-input and multi-output. However, it only supports two inputs and outputs, and the model needs to be changed for a different number of inputs, and it appears to be non-trivial to extend this model for more than two inputs and outputs. Hence, the model is not scalable and flexible for a different number of inputs and outputs.\n- As this is an architecture paper, it would be more convincing to also validate the effectiveness on other tasks like segmentation and detection, for which the model seems not directly applicable.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "X3bt8-IffuW", "reviewer": "Reviewer_6pnV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "Similar to MIMO in wireless communication, multi-input multi-output training improves network performance by optimizing multiple subnetworks simultaneously. Previous MIMO network architecture uses CNN. This paper proposes MixViT, the first MIMO framework for vision transformers which takes advantage of ViTs’ inherent mechanisms to share features between subnetworks. Unlike MIMO CNN, MixViT only separates subnetworks in the last layers leveraging a source attribution that ties tokens to specific subnetworks. The paper shows that MixViT can have significant gains across multiple architectures (ConViT, CaiT) and datasets (CIFAR, TinyImageNet, ImageNet-100, and ImageNet-1k).", "summary_of_the_review": "MIMO networks with vision transformer seems to novel. However, the paper does not provide deeper analysis on how it helps representation learning, besides the data augmentation benefits from input mixing, e.g. CutMix. The evaluation is also inadequate as it does not provide results on large datasets such as ImageNet.", "strength_and_weaknesses": "Strength\n1. This paper proposes MixViT, the first MIMO framework using vision transformers which takes advantage of ViTs’ inherent mechanisms to share features between subnetworks. \n\n2. The paper shows that MixViT can have significant gains across multiple architectures (ConViT, CaiT) and datasets (CIFAR, TinyImageNet, ImageNet-100, and ImageNet-1k).\n\nWeaknesses\n1. It is not clear how MIMO training helps representation learning. The paper lacks analysis on this. Does MIMO training improve hard example classification accuracy? Does it better cluster images into different clusters? Does most of the gains come from input mixing such as CutMix?\n\n2. The evaluation is on relatively small datasets. The paper should evaluate on larger datasets such as ImageNet.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "L4lnY5FLjp", "reviewer": "Reviewer_jCdb", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes to bring MIMO for CNNs to vision transformers. Specifically, it designs a source attributing module at last layer to separate the input source and perform classification for different input. Experiments are mainly conducted on small-scale datasets, CIFAR, Imagenet-100 to show its effectiveness. \n", "summary_of_the_review": "\nOverall the paper is a direct adaption of MIMO to vision transformers. However, the design space is not thoroughly studied, and current results are only on small-scale dataset and small models, raising the concern regarding if current design is optimal and can adapt to other large models. \n", "strength_and_weaknesses": " Strength:\n\t\tMany small-scale datasets are used for validating the results. And the paper proposes one working example of MIMO for vision transformers. \n\n\tWeaknesses:\n\t\tThe paper introduces one way of bring MIMO to vision transformers, which is quite straightforward and novelty-limited, \nSource attribution is a direct analogy of using different input projection in existing MIMO.  But there exist many other natural designs that can potentially work for vision transformers. The study space in the paper is not enough to show this proposed way is better.  i.e. important baselines are missing. For example, use different patch embedding for different inputs (similar to different project for inputs in MixMo paper). \n\t\tAgain, this specific design of late source embedding is not thoroughly studied. Early source embedding with carefully designs comparison is missing. Or middle source embedding, etc, this design should be carefully studied.\n\t\tExperiments are mainly conducted on small-scale datasets, ability to adapt to other large-scale datasets and large models remain unknown. Table 3 shows CaiT S-24 and M-24 results are 82.3 and 82.7. However, table 3 in CaiT paper shows s-24, m-24 results are 82.7 and 83.4 respectively. Can authors check the numbers here? And provides some large model (CaIT-M, or other large popular ViT variants) results would be better.  \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject