Perceptual Group Tokenizer: Building Perception with Iterative Grouping
Authors: Zhiwei Deng
Ting Chen
Yang Li
Venue: ICLR 2024 poster (2024)
Keywords: Self supervised learning
representation learning
image encoder
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Human visual recognition system shows astonishing capability of compressing visual information into a set of tokens containing rich representations without label supervision. One critical driving principle behind it is perceptual grouping. Despite being widely used in computer vision in the early 2010s, it remains a mystery whether perceptual grouping can be leveraged to derive a neural visual recognition backbone that generates as powerful representations. In this paper, we propose the Perceptual Group Tokenizer, a model that entirely relies on grouping operations to extract visual features and perform self-supervised representation learning, where a series of grouping operations are used to iteratively hypothesize the context for pixels or superpixels to refine feature representations. We show that the proposed model can achieve competitive performance compared to state-of-the-art vision architectures, and inherits desirable properties including adaptive computation without re-training, and interpretability. Specifically, Perceptual Group Tokenizer achieves 79.7% on ImageNet-1K self-supervised learning benchmark with linear probe evaluation, marking a new progress under this paradigm.

Published as a conference paper at ICLR 2024
PERCEPTUAL GROUP TOKENIZER: BUILDING PERCEPTION WITH ITERATIVE GROUPING
Zhiwei Deng, Ting Chen, and Yang Li Google Research and Deepmind
ABSTRACT
Human visual recognition system shows astonishing capability of compressing visual information into set of tokens containing rich representations without la- bel supervision. One critical driving principle behind it perceptual grouping (Palmer, 2002; Wagemans et al., 201 Herzog, 2018). Despite being widely used in computer vision in the early 2010s, it remains a mystery whether perceptual grouping can be leveraged to derive neural visual recognition backbone that generates as powerful representations. In this paper, we propose the Perceptual Group Tokenizer, model that entirely relies grouping operations extract vi- sual features and perform self-supervised representation learning, where a series of grouping operations are used iteratively hypothesize the context for pix- els or superpixels to refine feature representations. We show that the proposed model can achieve competitive performance compared to state-of-the-art vision architectures. and inherits desirable properties including adaptive computation without re-training, and interpretability. Specifically, Perceptual Group Tokenizer achieves 80.3% on ImageNet- K self-supervised learning benchmark with linear probe evaluation, establishing a new milestone for this paradigm.
INTRODUCTION
Visual recognition mechanisms matter. The pursuit of advanced vision algorithms that encode an image to meaningful representations dates back to late 80s, with two paradigms marking the progress over the past 40 years: feature detection (LeCun et al., 1998; Lowe, 2004; He al., 2016; Liu al., 2022b) and perceptual grouping (Shi & Malik, 2000; Uijlings et al., 2013; Arbeláez et al., 2014), where feature detection focuses on specific distinetive patterns, while perceptual grouping considers similarities among all pixels to produce compact set of tokens as proxies for image representation. Ever since the surge of deep learning, feature detection has predominated the vision field and become the main principle behind representation learning backbone designs and made impressive progress (Simonyan & Zisserman, 2014; Szegedy et al., 2015; He et al., 2016; Chen et al., 2017; Tan & Le, 2019: Qi et al., 2020; Liu et al., 2022b). The success of the former paradigm is, although striking, raising the question of whether perceptual grouping can also be used as the driving principle to construct visual recognition model.
Different from detecting and selecting distinctive features, perceptual grouping emphasizes on learn- ing feature space where similarity all pixels can be effectively measured (Uijlings et al., 2013; Arbeláez et al., 2014). With such a feature space, semantically meaningful objects and regions can be easily discovered with simple grouping algorithm and used compact set to represent an image (Uijlings et al., 2013; Arbeláez et al., 2014; Locatello et al., 2020). This indicates that image understanding essentially "pixel space tokenization", and being able produce generalizable fea- ture representations is tightly connected to whether the correct contextual pixels are binded together (Hinton, 2022; Culp et al., 2022).
The intriguing properties of perceptual grouping, including natural object discovery, deep connec- tions with information theory and compression (Ma et al., 2007), and association with biological vision system (Herzog, 2018) 8) cognitive science explanations (Palmer, 2002). have led to strong revival recently under deep learning frameworks (Locatello et al., 2020; Elsayed et al., 2022; Xu et al., 2022; Wu et al., 2022; Biza et al., 2023). However, these methods are either still focusing on small or toy datasets (Locatello et al., 2020; Chang et al., 2022; Biza et al., 2023), or used as auxilliary component (Xu et al., 2022; Ke & Yu, 2022; Seitzer et al., 2022) to strengthen exist-

---

Published as a conference paper at ICLR 2024
Image
Group tokens
Image
Group tokens
T
5d
TT.
Figure 1: Perceptual Group Tokenizer is entirely driven by grouping operations to perform rep- resentation learning. Group tokens (discovered objects) are shown above. See more in the appendix.
ing vision architectures for increased interpretability. Whether perceptual grouping can be used to build models and learn representations that are as informative and expressive as those learned by state-of-the-art vision architectures remains an open question.
In this paper, we propose Perceptual Group Tokenizer, a model trained under a self-supervised learn- ing framework, which builds visual representation entirely based on perceptual grouping operations. Given an image, the core of our model is to understand each pixel or patch through hypothesizing its contexts with grouping operations. Starting from given input patches, the grouping operation performs an iterative binding process onto a set of randomly sampled group tokens to determine the affinity groups based on similarities. The group tokens are then used as hypothesized contexts to refine the feature representation for the image. We show that applying this simple principle can al- ready produce expressive representations and works well with sulf-super pretraining on large vision dataset.
The grouping operation is also closely related to self-attention, a highly popular method commonly used in modern vision backbones. We build connection between the proposed grouping operation and self-attention and show that, if group tokens are treated as communication channels, self at- tention can potentially automatically emerge during learning processes special case, while the grouping operation can produce even richer interactions among tokens. Under this viewpoint, ViT (Dosovitskiy et al., 2020) can be considered as a grouping backbone, with a fixed number of group- ing slots equal to the number of input tokens, and the binding is achieved through stacking more than one layer with non-shared weights. This provides one explanation on why grouping mechanism can be effective on visual representation learning and has the potential to be promising competitive paradigm for vision architecture designs.
The primary contribution of this work is proposing a new architecture derived purely by perceptual grouping that achieves competitive performance compared to other state-of-the-art architectures on self-supervised learning benchmarks, contributing to new paradigm of developing vision architec- tures. The model has several key differences and advantages over ViT, including (1) explicit separat- ing out the "group token" concept to allow for automatic image parsing and flexible customization on the number groups without being binded to the number of patches; (2) much less peak mem- ory usage during inference time given the same number of input tokens; (3) adaptive computation without re-training the model, leading to flexible usage according to domains and computes.
RELATED WORKS
Vision architectures. There are two main frameworks for vision backbones. The first framework is Convolutioinal neural networks, which rely on local filters, sliding windows and translational equivariance to perform representation learning. Since the introduction of ConvNets in 1980s, Con- vNet was repopularized by AlexNet (Krizhevsky et al., 2012). The line of ConvNet is a classical inheritance from traditional feature detection methods (Lowe, 2004; Dalal & Triggs, 2005; Rosten et al., 2008), where instead of hand crafting features, an overcomplete set of filters are automati- cally learned to obtain high-response regions. The object understanding is built along the depth axis (Simonyan & Zisserman, 2014; Szegedy et al., 2015; He et al., 2016), with early layers capturing low-level parts and higher-level layers producing object structure representations (Zeiler & Fergus, 2014; Zhou al., 2014; Yosinski et al., 2015; Bau et al., 2017). In the feature detection framework,
2

---

Published as a conference paper at ICLR 2024
not every pixel is worth being used depending on particular tasks, leading to difficulty in obtaining representation for each pixel.
Recently, Vision Transformer (ViT) (Dosovitskiy et al., 2020), a second vision backbone framework, shows impressive performance and has surpassed ConvNet on visual recognition. The core of ViT is the iterative applying of self-attention operations (Vaswani et al., 2017; Dosovitskiy al., 2020). A direct usage of ViT on small patches (thus a high-resolution grid) is extremely computationally expensive due to its associated quadratic cost. Therefore, common practice is often partitioning the image into large non-overlapping patches (Dosovitskiy et al., 2020; Touvron et al., 2021), or constrain the operation to local regions (Liu et al.. 2021).
Self-supervise learning. The field of representation learning has seen significant interest in self- supervised learning during the past few years. The main evaluation results using linear probe on ImageNet benchmarks is approaching the results obtained by supervised learning (Oquab et al., 2023). Contrastive representation learning the early method that shows promising results (Oord et al., 2018; Chen et al., 2020a; Tian et al., 2020). BYOL (Grill et al., 2020) and DINO (Caron et al., 2021) propose to use a moving average target of an online network to perform self representation matching. Masked image modeling also shows to be effective on representation learning, where the masking is either at the pixel level (He et al., 2022) or the learned codebook level (Bao et al., 2021).
Object discovery. The perceptual grouping is essentially performing "object and stuff" discovery in the pixel space. It has broad connections with the early works in computer vision (Shi & Ma- lik, 2000; Uijlings et al., 2013; Levinshtein et al., 2013; Arbeláez et al., 2014; Pont-Tuset et al., 2016), the recent progress on object-centric representation (Burgess et al., 2019; Locatello et al., 2020; Chang et al., 2022; Hinton, 2022; Hénaff al., 2022; Culp al., 2022; Elsayed al., 2022), and biological or neural mechanisms on perceptual grouping (Palmer, 2002; Wagemans et al., 2012: Herzog, 2018: Kim et al., 2019). Despite the early popularity of perceptual grouping methods on various computer vision tasks (Shi & Malik, 2000; Uijlings et al., 2013; Levinshtein et al., 2013; Krähenbühl & Koltun, 2011), it has not attracted significant attention until several recent works that apply it as a side component on top of another main backbone (Seitzer et al., 2022; Liu et al., 2022a; Xu al., 2022; Ke & Yu, 2022). Some relevant works demonstrate alternative possibilities in archi- tecture design, but only uses cross attention without refining the patch feature space (Jaegle et al., 2021), or apply it on diffusion tasks (Jabri et al., 2022). Other methods also attempt to use ad-hoc sparsification methods top of ViT (Rao et al., 2021; Yin et al., 2022; Bolya al., 2023) for effi- ciency and are orthogonal to our work. A most related work (Ma et al., 2023) focuses on supervised learning and relies fixed-center pooling and less standard operations. In our proposed model, we adopt a design as VIT except for self attention, and highlight several key technical contributions, including multi-grouping with multi-seeding, adaptive computation without re-training, and other design choices for self-supervised representation learning.
MODELS
In this section, we introduce Perceptual Group Tokenizer (PGT), a visual recognition architecture entirely driven by perceptual grouping principles. We discuss the core operations for grouping in section 3.1, the building blocks and network architectures section 3.2, the loss function used for self-supervised learning in section 3.3, and the connections with other models in section 3.4.
3.1 PERCEPTUAL GROUPING
We start with introducing notations for method. Given image RxxxxC first reshape it as a sequence of small patches'. Each patch &p E wxxxxe has spatial shape h x w, where h < H and w < W, leading to N = fow number of patches per image. To represent patch, we embed it into a high-dimensional vector h IR.. The set of embedded tokens {h,) is referred to as input tokens in later parts, and used as inputs for the following grouping blocks.
HW
Feature refinement through hypothesizing contexts. Individual pixels do not have meanings with- out putting into contexts. At high level, image understanding or feature learning is equivalent to binding the correct contextual pixels at all locations. The core idea of our model is to generate many (e.g. over-complete w.r. number of objects in the image) hypothesized contexts and use the
We use 4 patches inputs in this work. Note that our method is generalizable to either pure pixels other forms of superpixels given proper patch-to-vector embedding layer.
3

---

Published as a conference paper at ICLR 2024
Sampling init.
GRU
Model outputs:
Linear
Attention
MLP
Embed
GRU
:
Concat
Grouping head
Patches
Input tokens
Group takens
Figure 2: Perceptual Group Tokenizer takes in sequence of patches (or pixels), generates high- dimensional embedding vectors for all patches, then them passes through series of grouping layers to refine the embedding vectors as feature representations. Each grouping layer performs K rounds of binding from input tokens to group tokens. To consider various grouping possibilities, multi- ple grouping heads adopted. Each group token provides useful context for input tokens for feature refinement. The final output of the model contains refined input token, group tokens, and assignments between input tokens and groups tokens.
hypothesized contexts cues to refine the feature representation of each patch. This process is achieved through a grouping module. Given input tokens }h,}%, the grouping module starts from a set of random samples (referred as group tokens) from random distribution, then performs binding process to aggregate information from input tokens to the group tokens, and ends up with set of group tokens c* {c² ) representing hypothesized contexts among input tokens. The relation between h; and Cj is soft assigment, indicating how likely an input token belongs to that context. Note that there are often various ways of generating groupings for an image, e.g. different seman- ties, colors, textures, etc., we propose the "multi-groupin operation" to hypothesize rich contexts for tokens. The overall model shown in figure 2.
Multi-grouping operation. The building block of our model is the multi- grouping operation G. which contains multiple heads to perform the binding process in parallel. This design encourages the model to consider multiple ways of generating groups under different projection spaces. Each head owns a separate Gaussian distribution with learnable means and variance, similar to (Kingma & Welling, 2013; Locatello et al., 2020). Starting from set of randomly sampled initial group tokens (0)
CHEAD Pini(-), the grouping operation uses doubly normalized attention weights aggregate information from h, and the produced group tokens CREAD are used for the next round binding. The attention normalization and feature projection are performed all heads separately.
CHEAD
((CHEAD, h; 0)
(1)
CHEAD
CHEAD
G(c
(K-1) h;; 0)
(2)
where after K steps the final group tokens c* (K)) is obtained, and 0 is learnable parameters in G. The grouping operator is summarized in algorithm
The sampling distribution PiNi(.) for initializing group tokens CHEAD needs to be lightweight. We explore two variations: (1) Gaussian distribution P(HULADS THEAD) with Iearnable means and vari- ance, and a one-step normalizing flow module that transforms a unit Gaussian noise to sample that follows more complex distributions. More details can be found in the appendix in section A.
(0)
Implicit differentiation. The iterative grouping process unrolls K steps per operation and leads to heavy burden in the training computation graph. Instead of explicitly backpropagating through the unrolled graph, we follow (Chang et al., 2022) and treat the multi- grouping process as a fixed point iteration per head. The gradient the backpropagation is approximated using first-order Neumann series, which can be simply achieved by detaching the output before the final iteration.
3.2 NETWORK ARCHITECTURE
Similar to standard ViT, our model refines the hidden representation h using L model layers. We use h to denote the representation after each layer, and explain the design in this section.
4

---

Published as a conference paper at ICLR 2024
Algorithm Multi-grouping operation using G.
der multi__grouge (h_y,y, h_valué, ateps, num_tokens, nar_heada)
... Inppu tenaors:
hukey and h_value aze projected malti head tensors with shape [aun_heads x B d]
Thitial group tokens.
group_toxens aampling distgggdddiiiiiiiiiiii ion (nasmples-num tokens, choice-' Gausslan") or "Flow!
group_tokens group_tckens..ret reshape heads, num_tckens, d) 4 [num_ heads x M.
Binding proceas.
for step it cange(st)ps);
implicis differentiation.
it step steps 1:
group_ toup_toees--- stop_grad (group_tokena) ean The follewing Attertion operation for group assi anseet.
cne-step grouping operatien. u..
attn aetn_natrix = attentianggroup.o wtt.agggoppooooaaaa.ooooooo h h_key) # ||num_
K M)
attn_matrix /- astn_matrix.aum u_updates ("asuum"aa,iidd-93555 group_tckens gru mon_celll [h_i(hh_paates grosp_tokana)
keeo_dim-Tzue)
atta_matrik, h_value) (nun X
Grouped mip/layernory pertorms independent m.p/ layernorm for each nead. group_tokens
group_tckene grouped
return group_sokems
Grouping layer. Each grouping layer takes h'-1 as input, and uses the grouping operation in equation to generate group tokens CREAD {Cj, J To the group tokens to provide context for each h! we perform another attention operation to obtain the attention matrix (only normalized over group token axis) A R X representing the assignment from input tokens to group tokens, and aggregate the feature back to the input token space:
h Alci C2,HAADi .... C.SHHEDD)
(3)
h' Linear Linear(hh (HAD) :... ....hH
(4)
h' = h'-1 + MLP(LN(N'))
(5)
This layer definition follows the standard ViT layer close possible, where features from each head are aggregated through concatenation and a linear layer transformation. Each token h is further refined using follow up multi-layer perceptron.
Grouping blocks. Similar to previous architecture designs (He et al., 2016; Liu et al., 2021). we define blocks for the model. One block contains multiple grouping layers that share the same hy- perparameters setups, i.e. the number of group tokens, and group token dimensions. The full model contains three grouping blocks. This increases the flexibility when exploring model design spaces.
3.3 SELF-SUPERVISION LOSS
We strictly follow the student-teacher self-supervision loss (Caron et al., 2021; Oquab al., 2023), and use moving average of online network (student model) as the teacher model to perform repre- sentation learning. To summarize group tokens outputed from the final layer, we use one multi-heal attention layer with learnable token to attend to all group tokens. The produced single vector is treated as the feature representation for the image and input to the loss function.
3.4 DISCUSSION
Our proposed model, perceptual group tokenizer, does contain self-attention operations and purely relies on group- ing operations. In this section, we link the grouping process to several techniques and discuss the rationale on why this model can be effective on representation learning.
Communication channels
Bind
Group tokens as "communication channels" The core of feature representation learning is how information ex- changed among pixels. In perceptual grouping backbones, we can consider the set of group tokens as communication channels, where information from different input tokens aggregated in various ways. Each group token represents a high-order channel that links input tokens with high affinity under certain pro- jected space exchange information among them. As thought experiment, if each input token is
Painwise attenton
Grouping operation
Figure Operation comparison.

---

Published as a conference paper at ICLR 2024
solely assigned to different group token (given enough group tokens), then the perceptual grouping layer is equivalent to one self attention layer (up to some engineering design difference). While self attention layers mainly rely pairwise communications, grouping operation, hypothetically. automatically learn and emerge both pairwise and higher-order information exchange through the group token communication channels. This can also be linked traditional factor graphs in proba- bilistic graphical models. Through the lens of that, grouping is forming factor nodes automatically through the learning processes. With a properly designed loss and grouping operation, it has the potential to be more effective if adopting per-layer comparison with self self-attention operations.
Efficiency. Due to the flexibility in customizing number of group tokens (controlled by initial num- ber of samples), grouping operation does not require a strict ( operation and is O(NM) on complexity. Furthermore, we show that inference time, number of group tokens can even be adaptively customized, given an already trained model.
4 EXPERIMENTS
We evaluate the representation learned by our model standard benchmarks based on the ImageNet-IK: K dataset. We also explore and analyze the design space of perceptual group tokenizer in section 4.2, investigate its adaptive computation ability section 4.3, demonstrate its generalization ability on semantic segmentation in section 4.4, and visualize learned attentions in section 4.5.
4.1 MAIN RESULTS
Setup. The widely-adopted standard benchmark for evaluating self-supervised learning methods ImageNet ILSV ILSVRC-2012 (ImageNet-IK) (Russakovsky al., 2015). Performance of models are measured by top-1 classification accuracy. The pre-trained backbones are frozen, with linear classifier trained on top. For fair comparison, we follow the standard data augmentation used in (Caron et al., 2021), with the same number of global views and local views. The model is optimized using AdamW (Loshchilov & Hutter, 2018) with learning rate 0.0005 and 1024 batch size for 600 epochs, trained with TPUv5 for 2 core hrs (512 cores for 41 hrs). We use 4x4 patches as image tokens, which keeps as much details as possible while maintaining reasonable computation costs.
Method
Arch
Param.
Linear probe (top-1 acc)
(Other backbones with different losses within the same batch of DINO for reference)
SCLR (Chen et al., 2020a)
RN50W4
375
76.8
SwAV (Caron et al., 2020)
RN50W2
93
77.3
BYOL (Caron et al., 2020)
RN50W2
93
77.4
SwAV (Caron et al., 2020)
RN5OW5
586
78.5
BYOL (Caron et al., 2020)
RN50W4
375
78.6
iBOT (Zhou et al., 2021)
ViT-B/16
85
79.5
BYOL (Caron et al., 2020)
RN200W2
250
79.6
SCLRv2 (Chen et al., 2020b)
RNI52W3+SK
794
79.8
BE (Peng et al., 2022)
ViT-B/16
85
80.1
(Fair comparison under the DINO loss and framework)
DINO (Caron al., 2021)
ViT-S/8
21
79.7
Ours (PGTG-S-1024
PGT-S
34
79.8
DINO (Caron et al., 2021)
ViT-B/16
85
78.2
DINO (Caron et al., 2021)
ViT-B/8
85
80.1
Ours BB25c)
PGT-B
70
79.7
Ours BB5T2)
PGT-B
70
79.9
Ours (PGTc-B-1024)
PGT-B
70
80.1
Ours (PGTr-B-256)
PGT-B
115
80.0
Ours (PGT,-B-512)
PGT-B
115
80.1
Ours (PGT,-B-1024)
PGT-B
115
80.3
Table 1: Comparison with strong baselines on ImageNet-11 1K under linear probe evaluation protocal. PGT -B-X X represents X number group tokens per grouping layer in inference (same trained model with 256 tokens is used). DIST: the distribution choice for group token initialization. G and represent Gaussian and Flow, respectively. Our model achieves 80.3%, competitive with of the-art vision backbones.

---

Published as a conference paper at ICLR 2024
Descend
Flat
Ascend
Token size Accuracy
576, 384, 192]
(384, 384, 384] 63.1
[192, 384, 576]
62.0
63.4
Token shape Accuracy
[192, 128, 64]
[128, 128, 128)
[64, 128, 192]
63.6
63.7
63.1
Table 2: Exploring the design choices for PGT. Token size: dimensions for group tokens in three grouping blocks. Token shape: number of tokens for group tokens in three grouping blocks. Ac- curacy measured ImageNet-1K K under linear probe protocal. Results indicate progressively large group token dimensions with flat or descend number of tokens arrangements work the best.
Architecture details. In the experiments, we mainly evaluate two variants of PGT: the main model and a tiny version for exploring design choices. On the ImageNet- IK benchmark, we report the performance metries of our main model. Three grouping blocks are used, with 10 grouping layers in each block. The dimension for input token is 384, with 256 group tokens per layer. The dimensions for group tokens are 98, 192, and 288 for the three blocks, respectively. There are 6 grouping heads used. For number of grouping iterations, we observe three rounds are sufficient to achieve good performance. The MLP hidden size for each layer is 384 as well, i... the MLP multiplication factor is 1. The final multihead attention layer uses learnable token with 2048 dimensions to summarize all group tokens outputs from the model.
The main results are summarized in table 1. We mainly compare with ResNet and ViT backbones, the two main stream vision architectures to show that perceptual grouping architecture can also achieve competitive results the challenging Image benchmark. Although our model is trained with 256 group tokens, the model can use different numbers of group tokens in inference (more experiments in section 4.2). We evalaute PGT with 256, 512, and 1024 number of group tokens and observe that the model can achieve 80.3% top-1 accuracy, showing the self-supervised learned feature of PGT is as good as the ones learned by VIT architectures.
4.2 ABLATIONS
To explore design choices of PGT, use tiny version of PGT with 3 blocks, 2 layer in each block (6 layers in total), 256 hidden size for input tokens, and number of grouping iterations. The learnable token in MAP head has 512 dimensions. There 100 parameters in this PGT-tiny.
Group token layouts. Given a fixed number of budget on group tokens, we explore three choices on how they should be arranged across grouping blocks and layers: descend, flat and ascend. Intuitively, more group tokens will have higher capacity of capturing smaller parts and detailed visual features, while less group tokens more prone carry global information. As shown in table 2 bottom row, flat or descend number of group tokens performs the best. In practice, we find that using flat (same number of group tokens in three grouping blocks) version achieves better training stability.
Group token dimension shapes. Similar to token number arrangements, we explore how group token dimensions should be set. Under three choices, progressively increasing the dimension size in the later layers performs the best, shown in first row of table 2. This also aligns with the intuition that later layers contain more information and requires higher capacity to represent groups.
Multi-grouping vs single grouping. We further test whether multi-head grouping helps improve performance. As a fair comparison, we use heads and 128 group tokens per head for a multi- grouping model, and head with 6x 128 group tokens for single grouping model. We find that adopting multi-head design can improve the performance from 62.2% to 66.3%, 4.1% accuracy boosts, showing that having multiple heads indeed helps with representation learning.
Grouping distribution entropy. Will grouping process collapse to some specific group token dur- ing training? We visualize the entropy of marginal distribution over tokens p(c) and conditional distribution p(c/2) ) in figure 4. Interestingly, we observe that conditional probability, i.c. the assign- ment to group tokens, tends to become more certain during training, while the marginal distribution remains having descend entropy, indicating collapses not happening in training.
Peak memory usage. As discusssed in section 3.4, given the same number of tokens, the grouping operation uses less memory than the self-attention operation. We show the percentage of peak memory usage in PGT compared to VIT-B with the same patch size (4x4) in table 3. The
7

---

Published as a conference paper at ICLR 2024
5.2
Layer
5.2
layer
layer
X2
4.8
Max
8 5.0
Max entropy
4.4.9
4.6
Max entropy
8 3.5
Layer 16
8.6.
4.6
Entropy
B 3.0
entropy
4.4
2.5
4.2
4.2
Epochs (600)
Epocht (600)
Couchs (600)
Epochs (600)
Figure 4: The entropy curves of grouping distributions p(c) and p(c]æz) across different layers.
usage obtained from the forward inference graph, as in practice the underlying complex hardware optimizer is a less accurate measurement and varies across infrastructures.
#group tokens
16
32
64
128
256
384
512
768
1024
ViT-B
Peak memory(%)
4.6
4.6
4.6
4.6
4.6
6.1
8.2
12.2
16.3
100
Table 3: Peak memory usage of PGT-B compared to the baseline model VTT-B with x patch size.
4.3
OUT-OF-DISTRIBUTTON ADAPTIVE COMPUTATION
One surprising and powerful ability of PGT is adaptive computation. For example, given a model trained using Mi group tokens per layer, one can choose to M2 group tokens in inference, where M2 * M1. This is because that the initial seeding group tokens are drawn from a proba- bilistic distribution, and the number of samples can be eustomized. This property leads to highly customizable inference without re-training the model. When M, * M2, the model copes with an -ittr--distiiuution (OOD) problem where test time setting is different from training. We observe surprisingly strong generalization with our model. Specifically, with more tokens M2 > Mi in in- ference, the performance can actually outperform the setting (M2 = M1) used training, even if is OOD for the model.
The results for OOD adaptive computation are summarized in table We mainly test PGTg-Tiny with grid evaluation that varies the number of group tokens in training M and the number of group tokens in inference N, and also show the main model's 's results in the last row. When using the main model PGTG-B perform adaptive inference, with only 12.5% of the number of group tokens compared to training, the performance can still be maintained at 72.1% % with only a ~88 drop on top-1 accuracy. The adaptive computation ability is important for both general image understanding where images have varying number of objects and need different numbers of groups, and scenarios where test-tim cest-time computational resource is constrained. This flexibility is an important advantage that grouping backbones hold.
tr/inf
16
32
64
128
256
384
PGT6-Ti-116
57.4 (x1)
58.3 (x2)
58.5 (x4)
58.5 (x8)
58.5(x16)
58.4 (x24)
PGTG-Ti-32
57.3 (xx)
59.9 (x1)
60.8 (x2)
61.0 (x4)
61.0 (xs)
60.9 (x12)
PGTG-Ti-64
53.0 (xi)
59.2 (x+)
61.7 (x1)
62.6 (x2)
62.9 (x4)
62.9 (x6)
PGTg-Ti-128
44.9 (x²)
56.6
61.8 (x))
63.9 (x1)
64.7 (x2)
64.8 (x3)
PGT2-Ti-256
27.2 (x)
47.4 (x
(x
58.8 (x4)
63.3 (x4)
65.1 (x1)
65.5 (x
PG-G-Ti-384
26.1 (x+)
43.0 (x++)
(x
61.7 (x4)
(x4)
65.5 (x1)
PGTG-B-256
60.4 (x+
72.1 (x
77.1 (x4)
78.9 78.9 (x
79.7 (x1)
79.9 (x3)
Table 4: Out-of-distribution adaptive computation by selecting different numbers of initially sam- pled tokens. Row: number of tokens used for training. Column: number of tokens used for infer- ence. Top-1 accuracy is reported under linear evaluation protocol using ImageNet-IK K. The reported performance of first six rows obtained using a tiny version of PGT, and last row the main model. Number of group tokens is the same for underlined numbers in training and inference. Bold numbers are the best results.
4.4 DOWNSTREAM TASK TRANSFER: SEMANTIC SEGMENTATION ON ADE20K
To evaluate the generalizability of pretrained feature produced by PGT, we test the transfer per- formance of semantic segmentation with ADE20k. Following the standard setup, finetune
8

---

Published as a conference paper at ICLR 2024
Grouping head
Grouping head 2
Grouping head
Layer
Layer
Input
Group taken
Group token
Group taken 3
Group takon
Layer
Layer
Input
Group token
Group token
Group token
Group token
Figure 5: Visualization of attention maps of each group tokens across layers and grouping head. indicates layer indices. Five group tokens for each grouping head. Smaller images are for early layers, arranged as five group tokens per grouping head. Large images are for the last layer.
model with the same data augmentation for 128 epoch. The baseline method uses DINO VIT- B/16 (Zheng et al., 2021). For our model, we add one linear classification layer after the pre-trained BG for fine-tuning. To adapt to more objects and complex scenes in the segmentation datasets, we use 1024 group tokens inference, benefiting from the adaptive computation ability of our model. We find that our model can obtain 45. % on mean IoU while the baseline achieves 44, % (Bao et al., 2021), leading to a 1.0% improvements.
4.5 GROUPING VISUALIZATION
We visualize the attention maps calculated between group tokens and input tokens in figure 4.5. We find that (1) using multiple grouping heads can capture different information within each head. For example, in layer O, the first head captures light and color, second head focuses only spatial locations, and the third head potentially relies on textures; (2) group tokens can capture different semantic parts, for example, in the first image, group tokens separate apple, jar, handle, and back- ground. In the second image, camel, legs, camel hump. and human are separately grouped. Com- pared standard ViT in DINO (Caron et al., 2021) where only single foreground can be extracted using [CLS] token, our model can flexibly group different parts given an image, leading to set of tokens that are potentially more meaningful and customizable. Note that the grouping results are still different from human's 's vision, and sometimes generates parts that seem to be "fragmented". This is possibly due to the "parts-to-whole with data augmentation" training loss. Human vision, in contrast, is sensitive to moving objects and trained within 4D space. Nevertheless, we believe with a similar dataset, environment and loss design, our grouping model can potentially produce groupings more coherent and sensitive to boundaries and moving objects.
CONCLUSION
In this paper, we propose Perceptual Group Tokenizer (PGT), new visual recognition architecture entirely built through perceptual grouping principles. The proposed model shows strong perfor- mance self-supervised learning benchmark ImageNet-IK with linear probe evaluation, and has desirable properties such as adaptive computation and high model interpretability each operation. This work can enable a paradigm for designing visual recognition backbones, and we hope to inspire more research progress along this direction. One limitation of the proposed model is its relatively expensive computation cost due the iterative grouping processes. This be poten- tially addressed by other grouping operations, such as those grouping operations with closed-form solutions, which promising direction for the future work.

---

Published as a conference paper at ICLR 2024
REFERENCES
Alexander A Alemi, lan Fischer, Joshua Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint wrXiv:662..00410. 2016.
Pablo Arbeláez, Jordi Pont-Tuset, Jonathan T Barron, Ferran Marques, and Jitendra Malik. Mul- tiscale combinatorial grouping. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 328-335, 2014.
Hangbo Bao, Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training image transformers. arXiv preprini arXiv:2206.08220 2021.
David Bau, Bolei Zhou, Aditya Khosla, Aude Oliva, and Antonio Torralba. Network dissection: Quantifying interpretability of deep visual representations. Proceedings of IEEE conference computer vision and pattern recognition, pp. 6541 16549, 2017.
Ondrej Biza, Sjoerd van Steenkiste, Mehdi SM Sajjadi, Gamaleldin F Elsayed, Aravindh Mahen- dran, and Thomas Kipf. Invariant slot attention: Object discovery with slot-centric reference frames. arXiv preprint wrXiv:230..04973. 2023.
Daniel Bolya, Sheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https:/ wwwpr/veeneeeeeeeeeeeeeeeeeeewe. net/ forum?id-
JroZRaRw 7Eu.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Trina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint 1901 2019.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 33:9912-9999 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transformers. Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660, 2021.
Michael Chang. Tom Griffiths, and Sergey Levine. Object representations as fixed points: Training
iterative refinement algorithms with implicit differentiation. Advances in Neural Information
Processing Systems, 35:32694-32708, 2022.
Liang-Chich Chen, George Papandreou, lasonas Kokkinos, Kevin Murphy, and Alan Yuille. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. IEEE transactions on pattern analysis and machine intelligence, 40(4): 834 848, 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning visual representations. International conference on machine learning, pp. 1597-1607. PMLR, 2020a.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton. Big self-supervised models are strong semi-supervised learners. Advances in neural information pro- cessing systems, 33:22243 00.2223--22255. 2020b.
Laura Culp, Sara Sabour, and Geoffrey Hinton. Testing glom's ability to infer wholes from ambiguous parts. arXiv preprint (.di:22/1.16666 2022.
Navneet Dalal and Bill Triggs. Histograms of oriented gradients for human detection. In 2005 IEEE computer society conference on computer vision and pattern recognition CCVPR'05), volume 1, pp. 886-893. Icee, 2005.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 35:16344-16359 2022.
10

---

Published as a conference paper at ICLR 2024
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 2016.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. image is worth 16x words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11922 2020.
Gamaleldin Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C Mozer, and Thomas Kipf. Savi++: Towards end-to-end object-centric learning from real-world videos. Advances in Neural Information Processing Systems, 35:28940-28954 2022.
Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural 2020.
information processing systems,
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders scalable vision learners. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 16000-16009, 2022.
Olivier J Hénaff, Skanda Koppula, Evan Shelhamer, Daniel Zoran, Andrew Jaegle, Andrew Zisser- man, João Carreira, and Relja Arandjelovié. Object discovery and representation networks. In European Conference Computer Vision, pp. 123-143. Springer, 2022.
Michael H Herzog. Perceptual grouping. Current Biology, 28(12) 2018.
Geoffrey Hinton. How to represent part-whole hierarchies in a neural network. Neural Computation, pp. 1-40, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation.
arXiv preprint axXiv:2212.11972, 2022.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint arXiv:2107.147999 2021.
Tsung-Wer Ke and Stella X Yu. Cast: Concurrent recognition and segmentation with adaptive segment tokens. arXiv preprint arXivv2210.003314, 2022.
Junkyung Kim, Drew Linsley, Kalpit Thakkar, and Thomas Serre. Disentangling neural mechanisms for perceptual grouping. arXiv preprint arXiv:1906.01558. 2019.
Diederík Kingma and Max Welling.
Auto-encoding variational bayes. arXiv preprint
11X 2013.
Diederik Kingma and Max Welling. Auto-encoding variational bayes, 2022.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018.
Philipp Krähenbühl and Vladlen Koltun. Efficient inference in fully connected crfs with gaussian edge potentials. Advances neural information processing systems, 24, 2011.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.

---

Published as a conference paper at ICLR 2024
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-25 2324, 1998.
Alex Levinshtein, Cristian Sminchisescu, and Sven Dickinson. Multiscale symmetric part detection and grouping. International journal of computer vision, (04:117-134) 2013.
Kai Liu, Tianyi Wu, Cong Liu, and Guodong Guo. Dynamic group transformer: A general vision transformer backbone with dynamic group attention. arXiv preprint arXiv:2203.03993 2022a.
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 Proceedings of the IEEE/CVF international conference computer vision, pp- 10012-10022, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer. Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1997-11986, 2022b.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. Advances in Neural Information Processing Systems, 33: 1523-11538, 2020.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. 2018.
David G Lowe. Distinctive image features from scale-invariant keypoints. International journal of computer vision, 10:91-110, 2004.
Xu Ma, Yuqian Zhou, Huan Wang. Can Qin, Bin Sun, Chang Liu, and Yun Fu. Image as set of points. In The Eleventh International Conference Learning Representations, 2023. URL
httt:s:///penreview wtw.eeiiieeeeeeeeeeeeeeee wwt.:/eeee//tneeeeeeeeeeeeeeeeee
Yi Ma, Harm Derksen, Wei Hong, and John Wright. Segmentation of multivariate mixed data via lossy data coding and compression. IEEE transactions on pattern analysis and machine intelligence, 29(9):1546
2007.
Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Con- ference on Machine Learning, pp. 3403-3412. PMLR, 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint 1rX 2018
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint urXiv:2304.07199, 2023.
Stephen E Palmer. Perceptual grouping: It's later than you think. Current Directions in Psycholog- ical Science, 11(3):101-166, 2002.
Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. Beit v2: Masked image modeling with vector-quantized visual tokenizers. arXiv preprint arXiv:2208.06:: 2022.
Jordi Pont-Tuset, Pablo Arbelae7, Jonathan T Barron, Ferran Marques, and Jitendra Malik. Mul- tiscale combinatorial grouping for image segmentation and object proposal generation. JEEE transactions on pattern analysis and machine intelligence, 39(1): 128-::24-440, 2016.
Haozhi Qi, Chong You, Xiaolong Wang, Yi Ma, and Jitendra Malik. Deep isometric learning for visual recognition. In International conference on machine learning, pp. 7824-7835. PMLR, 2020.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 14317997-,3949, 2021.
Sid Reddy, Anca Dragan, and Sergey Levine. Pragmatic image compression for human -n--iii-ttttt -hop-looo decision-making. Advances in Neural Information Processing Systems, 34:26499-2651 10, 2021.
12

---

Published as a conference paper at ICLR 2024
Edward Rosten, Reid Porter, and Tom Drummond. Faster and better: A machine learning approach to corner detection. IEEE transactions on pattern analysis and machine intelligence. 32(1):105- 119, 2008.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 11::211-252 2015.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Simon-Gabriel, Tong He, Zheng Zhang, Bernhard Schölkopf, Thomas Brox, al. Bridging the gap to real-world object-centric learning. In The Eleventh International Conference on Learning Representations, 2022.
Jianbo Shi and Jitendra Malik. Normalized cuts and image segmentation. IEEE Transactions on pattern analysis and machine intelligence, 22(8):888-99)) 100-90055220 2000.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1509, 2014.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International conference on machine learning, pp. 4105-6114. PMLR, 2019.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding, In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part 16, pp. 776-794. Springer, 2020.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Jasper RR Uijlings, Koen EA Van De Sande, Theo Gevers, and Aroold WM Smeulders. Selective search for object recognition. International journal of computer vision, 104:154-171, 2013.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Johan Wagemans, James H Elder, Michael Kubovy, Stephen E Palmer, Mary A Peterson, Manish Singh, and Rüdiger Von der Heydt. A century of gestalt psychology in visual perception: I. perceptual grouping and figure-ground organization. Psychological bulletin, 138(6):112,, 2012.
Ziyi Wu, Nikita Dvomik, Klaus Greff, Thomas Kipf, and Animesh Garg. Slotformer: Unsupervised visual dynamics simulation with object-centric models. In The Eleventh International Conference on Learning Representations, 2022.
Jiarui Xu, Shalini De Mello, Sifei Liu, Wonmin Byeon, Thomas Breuel, Jan Kautz, and Xiaolong
Wang. Groupvit: Semantic segmentation emerges from text supervision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18134-18144 2022.
Hongxu Yin, Arash Vahdat, Jose M Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10809-10818. 2022.
Jason Yosinski, Jeff Clune, Anh Nguyen, Thomas Fuchs, and Hod Lipson. Understanding neural networks through deep visualization. arXiv preprint arXiv:1506.066 2015.
Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part / 13, pp. 818-88 133. Springer, 2014.
13

---

Published as a conference paper at ICLR 2024
Sixiao Zheng. Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Tom,, et al. Rethinking semantic segmentation from sequence- tosseueee- perspective with transformers. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 6881-6890, 2021.
Bolei Zhou, Aditya Khosla, Agata Lapedriza. Aude Oliva, and Antonio Torralba. Object detectors emerge in deep scene enns. arXiv preprint 44Xiv:1412.68566 2014.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pre-training with online tokenizer. arXiv preprint arXiv:2111.07832, 2021.
14

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 LEARNABLE SAMPLING DISTRIBUTIONS
Our proposed Perceptual Group Tokenizer (PGT) model initializes a set of group tokens through sampling from distribution. This set of group tokens then serve as the initial "seeding" for the grouping process. We explore two methods to serve as the initial distribution: learnable Gaussian distribution and Normalizing Flows. We would like the extra cost of the grouping process to be minimal, therefore, use two leigh-te versions.
A.1.1 GAUSSIAN
Similar to the standard usage learnable Gaussians in generative model literature (Kingma & Welling, 2022; Ho al., 2020), the reparameterization to perform learnable sampling process: e = μ + a * e, where e is drawn from a unit Gaussian (0,,
A.1.2 FLOw
As Gaussian distribution might have limitations in covering complex distribution shapes, especially in the high-dimentional space, we also explore version with one step of affine coupling flow trans- formation (Dinh et al., 2016). Since we only require the differentiable sampling procedure and do not need to compute the determinant of Jacobian matrix, we directly apply the transformation without splitting the dimensions by half:
c a a+e+b
(6)
(log t) MLP(c)
8 = exp(logs)
(8)
c s*c+t
where is drawn from unit Gaussian NN 1). This transformation is simply a re-scaling and translation (similar to Gaussian) but conditioned on per sample c. More details are in (Dinh et al., 2016; Kingma & Dhariwal, 2018). We only apply one step of this transformation, leading to minimal parameter increase and negligible inference time difference.
A.2
MODEL ANALYSIS
In this section, we add more analysis on our model's performance and computational costs.
A.2.1 GROUPING ENTROPY
Grouping distribution entropy. The main paper has discussed and shown the grouping distribution entropy curves several layers. In the appendix, demonstrate curves from more layers in figure 6 and figure? 7, where the first one is marginal distribution and second one is conditional distribution.
5.0
Layer
5.2
Layer
Entropy at d
5.2
Layer
4.5
Max entrupy
8.5 5.0
5.8
Layer
5.0
Max entropy
Max entropy
4.0
Max entropy
4.6
3.5
4.4
3.0
Epochs (600)
4.2
4.0
Epochs (600)
4.2
Epochs (600)
Epochs (600)
5.0
5.2
Max entropy
Layer 16
5.00
4.75
5.25
5.2
5.0
Layer 24
entropy
0.4.5
05.00
5.25
Leyer Max entropy
4.2
4.8
4.8
4.4
4
4.50
M
Layer
Max entropy
3.8
4.0
22 3.755
24225
8..00
MM
Epochs (600)
Epochs (600)
Epochs (600)
(600)
Figure 6: The entropy curves of marginal distribution p(c) grouping across different layers.
15

---

Published as a conference paper at ICLR 2024
Entrapy plena
5.0
Layer 0
entropy
2.5.2
1.5.0
X 5.2
5.00
Layer 8
entropy
Entropy pfcix)
4.5
Layer 12
entropy
3 4.6
4.0
1.4
4.4
4.0
4.2
3.8
Layer entropy
3.8
3.6
4.0
2.5
3.0
Epochs (600)
Epochs (600)
Epochs (600)
Epochs (600)
Enirreay ofcc)
4.55
5.0
5.5
Layer
Max entropy
Ertropy of picK
5.0
4.5
3.5
4.0
3.0
Layer Max entropy
ENTIONN of 4.5
4.0
5.0
5.5
3.5
Layer 24
entropy
Entropy
4.5
5.0
3.3
Max entropy
Layer
mww
2.5
2.0
2.0
2.5
2-5
3.0
2.5
Epochs (600)
1.5
Epochs (600)
2.0
Epochs (600)
Epochs (800)
Figure 7: The entropy curves of conditional distribution p(clæz) grouping across different layers.
A.2.2 GROUPING ITERATIONS
In our backbone, we find that more grouping iterations will lead to better performance. We explore the number of grouping iterations on the PGT-Tiny model and PGTG-B-256. On the tiny version, we find the model achieves 61.4, 63.8, and 65.1 on the linear probe evaluation with number of interations is 2, and For the main model, the performances are 79.3, 79.6, and 79.7 respectively. The model's 's increased depth potentially helps with the lack of grouping iterations in the deep model. But in general, having the grouping process still important obtaining higher performance.
A.2.3 INFERENCE TIME
We also profile our model's inference time, compared with ViT-B with 4x4 patches (the same amount tokens) for ablation study on the grouping operation. Note that model and framework are built upon a complex infrastructure that uses XLA and other hardware accelerator to optimize speed. We find varying number of group tokens only lead to small influence. PGT-B-256 has 640 im/sec/core and ViT-B/4 has 680 im/sec/core. Using smaller number grouping iterations can speed up the inference 710 im/sec/core (2 iter2) and 820 im/sec/core iter).
Note that this is only due to the specialty of the underlying infrastructure. In general, having less number of group tokens should still increase the inference speed, since the attention operation is key computation bottleneck for vision models.
A.2.4 GELOPS
In table A.2.4, we show the gflops for our model under various inference budgets. Note that, as pointed in other works (Dao et al., 2022), gflops often do not fully reflect the model's computation performance. Due to that our model needs iterative grouping process, it'll increase the gflops count. But as shown in peak memory usage and inference time, the model's ss computation costs are either similar are much less.
#group tokens
16
32
64
128
256
384
512
768
1024
ViT-B
Gflops
99.2
131.1
194.9
322.6
577.8
833.1
1088.3
1599.0
2109.3
451.6
Table 5: Gflops of PGT-B compared to the baseline model ViT-B with the same patch size.
A.2.5 PROBABILISTIC PERSPECTIVE OF GROUPING OPERATIONS
Due to the probabilistic nature, our model is also quite compatible with full "treatment" with the variational inference framework, which can provide certain backup for our grouping operations already in the current model. We can treat the group token embeddings where the grouping process uses iterative amortized inference (Marino et al., 2018) to refine the la- tent variable. The grouping modules, including GRU, MIP, attention, and other layers are designed
the latent variables,
16

---

Published as a conference paper at ICLR 2024
to better infer the embeddings (latent variables). The training signal is pragmatic loss (instead of reconstruction loss), which has been demonstrated in (Reddy et al., 2021; Alemi et al., 2016). The key differences are: (1) there is sampling in each inference step; (2) the regularization from unit Gaussian distribution is set to zero. We do believe a full probabilistic treatment of the perceptual grouping architecture can be very interesting next step.
A.2.6 MORE VISUALIZATIONS
In this section, we show more visualizations of the attention maps for generated group tokens by Perceptual Group Tokenizers in figure 8, and 10.
Grouping iterations
Grouping iterations
Group
Group
Cke
Input
Group
Group
Group
Group
Input
Group
Group
Group token
Group
Input
Group
Group
Group token
Group
Input
Group
Group Token
Group token
Input
Group taken
Group Token
Figure 8: Visualization for attention maps of group token samples during the grouping process. PGT
uses 256 group tokens in inference time.
17

---

Published as a conference paper at ICLR 2024
Input
Group taken
Group token 2
Group taken
Group token 4
Group taken 5
Group token
Group taken
Group taken
B 0ee!!
5d
Figure 9: Grouping results from the 2 layer, using 8 group tokens in inference time.
18

---

Published as a conference paper at ICLR 2024
Input
Group token
Group token
Group token 3
Group token 4
Group token
Group token
Figure 10: Grouping results from the last layer, using 256 group tokens in inference time.
19


Reviews:
[{"review_id": "JwgIMRkI9Z", "reviewer": "Reviewer_8m3b", "summary": "This paper introduces the Perceptual Group Tokenizer (PGT), a novel vision model that relies on iterative grouping to extract visual features and learn representations in a self-supervised manner. PGT proposes to use grouping operations instead of self-attention layers in ViT, which yields a self-attention-free visual backbone. It demonstrates competitive performance on the ImageNet-1K benchmark. The model also shows properties like adaptive computation and high interpretability. The paper provides comprehensive analysis, ablation studies, and visualizations that underscore the model's capability and potential as a new paradigm in visual backbone architecture design.", "questions": "1. The model uses patch size 4x4, which means the visual backbone only downsamples the image by 4. Intuitively, this model should be good at dense prediction tasks that require high feature resolution, e.g. semantic segmentation and object detection. The authors reported the results of semantic segmentation on ADE20K in Section 4.4, but it only outperforms the ViT-B by a smaller margin. I understand that the segmentation architecture is different. So it would be interesting to compare ViT-B vs PGT-B with the same segmentation architecture (linear classification layer). \n2. Since PGT is self-attention-free, the computation cost is not quadratically increasing with the input resolution. But it is still reasonable to compare to ViT under the same resolution, for example, 16x16 patch size and 8x8 patch size. I am wondering whether authors have done this ablation. \n3. As mentioned in the weakness section, it would be interesting to have the computation cost of PGT. As Table 1 shows, as we increase the number of group token from 256 to 768, the linear probe accuracy increases from 79.3 to 79.7 How about less than 256 tokens and more than 768 tokens? It would be insightful to have a graph of number of tokens vs accuracy and inference speed. \n4. In Mask Autoencoder paper, researchers find that higher linear probing accuracy may be not necessarily stand for better representation. It would be also interesting to compare with MAE-ViT against PGT under the fine-tuning setting.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. I like the idea of grouping operations only, without self-attention. Using perceptual grouping is innovative and theoretically sound, offering a fresh perspective on representation learning and architecture design. \n2. The adaptability of computation in inference mode is also interesting. Without re-training, the model could inference with different number of group tokens. Table 1 also shows the accuracy will increase as there are more group tokens\n3. The visualization of the attention map is very interesting. It not only shows more iterations yield clearer grouping, but also shows different grouping heads kind of learning disjoint visual representations.", "weaknesses": "1. Only ViT-B level(70-80M parameter) model is reported. It would justify the effectiveness if the proposed architecture works when scaling the model size up. \n2. The computation cost, peak memory usage, and inference speed comparison with ViT-B are not reported. It would be informative for readers how fast the PGT is since PGT doesn't have memory/computation demanding self-attention operations.", "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": "", "extra_scores": {}}, {"review_id": "kFuEEeuDa2", "reviewer": "Reviewer_qQDA", "summary": "This paper proposes a perceptual group tokenizer, which just uses the grouping operations to extract visual features and perform self-supervised learning. The authors also explain the connection between the proposed perceptual group tokenizer and the self-attention. The experimental results show the performance is competitive with some state-of-the-art self-supervised methods.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The concept of perceptual group tokenizer is novel, and seems to enable the networks to have more good properties including interpretability and so on.\n\n2. Discussion between the perceptual group tokenizer and the self-attention is interesting, and can provide a new vision for vision transformer design.", "weaknesses": "1. The motivation is not clear. Since we have powerful vision transformers already, what is the advantages of the proposed perceptual group tokenizer. The authors claim that the perceptual group tokenizer have good properties such as adaptive computation without re-training and interpretability. However, I don't see the experimental results or the visualization that can provides proof of this claim.\n\n2. The performance is still concerned. Because we should also focus on the accuracy despite some good properties, the results shown in Table 1 demonstrate that the method cannot beat the baselines. Moreover, the state-of-the-art methods proposed in 2023 are not compared.\n\n3. The explanations to the grouping operation should give more details. Since grouping seems to be a explicit operation, the implicit operations for grouping such as MLP should show the correlation of the term \"grouping\".", "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": "", "extra_scores": {}}, {"review_id": "AY5Qg4MfL1", "reviewer": "Reviewer_P2FR", "summary": "This paper proposes the Perceptual Group Tokenizer (PGT), a ViT-like architecture that clusters tokens to implement the principle of perceptual grouping for visual recognition. Specifically, PGT combines a slot attention-like token grouping module within ViT blocks.", "questions": "There are many similar works in recent years. What are the substantial differences or unique advantages of this work?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is generally well-written.\n- Prior work on feature detection and perceptual grouping is well-discussed, although recent works on token grouping are missing.", "weaknesses": "**Limited technical novelty**\n\nThe proposed method is essentially a combination of ViT and slot attention (or its token grouping variants).\nThe concept of token grouping has been extensively studied in prior work. Please refer to the related work section of ToMe [1] and CoCs [2] for examples.\nOn the other hand, DINOSAUR [3] combined a ResNet/ViT encoder with slot attention to scale up object-centric learning for real-world images, which is also relevant to this work.\n\n[1] Token merging: Your vit but faster. ICLR'23.\\\n[2] Image as Set of Points. ICLR'23.\\\n[3] Bridging the Gap to Real-World Object-Centric Learning. ICLR'23.\n\n---\n**Unclear empirical benefits**\n\nIn addition to the technical novelty, the empirical merit of the proposed method is unclear.\n1. Performance: It is not better than prior works, such as DINO.\n2. Efficiency: Several efficient ViT-based works exist, such as ToMe, DynamicViT [4], A-ViT [5], etc.\n3. Adaptive computation w/o retraining: ToMe claims to offer the same benefit.\n4. Grouping visualization: Other grouping methods also offer similar advantages (Fig. 4 of ToMe). CAST [6] is even better in this regard.\n\n[4] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification. NeurIPS'21.\\\n[5] A-ViT: Adaptive Tokens for Efficient Vision Transformer. CVPR'22.\\\n[6] CAST: Concurrent Recognition and Segmentation with Adaptive Segment Tokens. arXiv'22.", "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": "", "extra_scores": {}}, {"review_id": "bKyKLgjGiG", "reviewer": "Reviewer_TtGg", "summary": "The authors propose a self-supervised learning method that relies only on grouping operations. They call their model the Perceptual Group Tokenizer (PGT). The model's linear probe performance is on par with other state-of-the-art models such as ViT. It also produces a highly interpretable representation and shows some interesting properties like being somewhat adaptable at inference time. Overall, the PGT demonstrates that grouping operations alone can produce a rich visual representation.", "questions": "As I mention in the weaknesses, could the authors touch on what are the benefits of the PGT over ViT given that they perform similarly on the ImageNet 1k linear probe?\n\nIn a similar vein, I would also like some of the model interpretability results to be expanded on in the context of other state-of-the-art methods as well as human visual perception.\n\nFinally, I thought adaptive computation (Section 4.3) was a bit too brief for me to appreciate. Could the authors explain what they mean by adaptive computation and why it is a benefit of PGT?\n\nGenerally, I enjoyed this paper. I am open to increasing my score if these concerns are addressed,", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The writing of this work is very clear. The authors do a great job of motivating the problem and contextualizing its significance. As the authors note, perceptual grouping has historically been an important concept in computer vision, but it had not been proven to be as powerful as feature detection. This work proposes a novel way to use grouping principles for self-supervised representation learning on large-scale natural image data. \n\nThe approach also appears to be well done. The analyses of number and size of grouping are also insightful, and demonstrate a thorough evaluation of their model. Making the connection to self-attention in ViT is also an important and original contribution for understanding why both the PGT and ViT work.", "weaknesses": "One of the main weaknesses I see is that the PGT performs very similarly to the ViT. As the authors note, their method is in some sense a more general version of the self-attention approach so it is somewhat unsurprising that performance is on par with ViT. I think the paper could be strengthened by more discussion about why the perpetual grouping tokenizer might be better or more useful than other methods like ViT.\n\nAlong similar lines, I think the interpretability results in Figure 6 and section 4.5 could be better contextualized. It is not immediately obvious how the interpretability conveyed by the attention maps compares to other state-of-the-art models. In addition, this analysis could be strengthened by a discussion of how the attention maps at each stage do or do not agree with what we would expect of human vision.", "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": "", "extra_scores": {}}, {"review_id": "Fnlnh4Nfwf", "reviewer": "Reviewer_ABtJ", "summary": "This paper proposes the Perceptual Group Tokenizer model (PGT), which utilizes grouping operations for visual feature extraction. It demonstrates competitive performance in self-supervised learning while reduces the computation complexity into O(n*m) compared with the complexity O(n^2) of vision transformer. The PGT also offers adaptive computation without re-training via flexible number of grouping tokens, and offers interpretability in feature representations.  Quantitative experiments and visualization demonstrate the PGT's effectiveness.", "questions": "1. Do the group tokens in each block generate independently? If not, maybe the relavant description is ambiguous. If so, why don't the group tokens take its' output of the previous block as the next block's input?\n2.  Could the iterative grouping processes be considered multiple cross-attention between input tokens and group tokens? What's the difference or advantages of the iterative grouping processes compared with cross-attention? How much does the number of grouping iterations matter?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The idea of Perceptual Group Tokenized is intrersting. By alternatively refining the group tokens and image feature tokens, the computation complexity has been reduced comparing with self-attention. \n2. The adaptive computation ability of PGT is good, which could be used to meet the needs of different inference speeds. \n3. The interpretability of model is relatively good.", "weaknesses": "1.  The experiment is not sufficiently comprehensive and the results are weak. Since a new model architecture is proposed, the performance under both supervised learning and self-supervised learning should be presented to show its’ universal. In performance comparison, the listed comparable architectures should be compared with the same pre-training/supervised-training strategy, a similar amount of parameters. It will be better to do experiment with more model sizes(e.g, PGT-B, PGT-S,PGT-Ti or more ) and compare with counterparts under different sizes. The performance on downstream tasks are not clear, although segmentation performance is reported, more comparisons with relative methods are not sufficient.\n2.  Besides the number of parameters, for fair comparison, the computational cost(Gflops or MACs) should be shown. The inference time is missing, which is important for evaluating models’ efficiency.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)