---
title: Fast Universal Agglomerative Pooling (UniAP)
url: https://www.emergentmind.com/topics/fast-universal-agglomerative-pooling-uniap
type: topic
---

# Fast Universal Agglomerative Pooling (UniAP)

Fast Universal Agglomerative Pooling (UniAP) is a pseudo-mask generation algorithm introduced in "S2-UniSeg: Fast Universal Agglomerative Pooling for Scalable Segment Anything without Supervision" [2508.06995]. It is formulated as an efficient, nonparametric, heuristic pooling method that identifies groups of strongly-connected nodes in parallel via connected components, producing universal pseudo-masks at both instance level and semantic level, and at multiple granularities, within tens of milliseconds per image. In the S2-UniSeg framework, UniAP replaces the time-consuming offline pseudo-mask stages characteristic of multi-stage self-supervised segmentation pipelines, enabling single-stage continuous pretraining in a student–momentum teacher setup.

## 1. Motivation and problem setting

UniAP was proposed to address bottlenecks in prior self-supervised segmentation pipelines. Many earlier methods follow a multi-stage "cut-and-learn" schedule: pseudo-masks are first generated offline using graph partitioning methods such as TokenCut and MaskCut, sometimes in multiple rounds, and training alternates with pseudo-label regeneration as in UnSAM and U2Seg. According to the S2-UniSeg formulation, this regime is slow, can require seconds per image, is difficult to scale to large datasets such as SA-1B, and induces discontinuous optimization because training must stop for mask generation [2508.06995].

Within this context, UniAP is positioned as the mechanism that makes online pseudo-mask generation practical. It operates directly on self-supervised features, produces both instance-level and semantic-level masks, and supports hierarchical multi-granular outputs. In S2-UniSeg, the momentum teacher processes a global view once per iteration, runs UniAP to generate pseudo-masks and query features, and supervises the student on local views through Query-wise Self-Distillation. The stated consequence is continuous pretraining without offline cut stages or relabeling rounds, with scalability to millions of images.

The paper explicitly contrasts this design with normalized-cut-based methods. UniAP avoids eigenvalue problems and instead uses thresholded local similarities and connected components to cluster nodes in parallel, with a feature update designed to suppress boundaries. This places UniAP between classical agglomerative grouping and learned feature-based segmentation, while remaining nonparametric in the pooling stage.

## 2. Formal construction of the pooling graph

UniAP begins from the teacher encoder's largest-scale feature map at stride 4. After L2 normalization and flattening, the feature tensor is written as
$$
F \in \mathbb{R}^{HW \times d},
$$
where $H$ and $W$ are the spatial height and width at stride 4, and $d$ is the feature dimension [2508.06995].

From these features, the paper defines a global spatial affinity matrix
$$
A = F F^\top \in \mathbb{R}^{HW \times HW}.
$$
Because the features are L2-normalized, entries of $A$ lie in $[-1,1]$ and correspond to cosine similarity.

The initial graph for instance pooling is
$$
G^0 = \{V^0, M^0, E^0\}.
$$
Its components are specified as follows:

- **Nodes**: each pixel or token at stride-4 resolution is a node, so the initial number of nodes is $s^0 = HW$.
- **Adjacency**: $E^0 \in \{0,1\}^{s^0 \times s^0}$ is a 4-neighborhood graph, with edges only between horizontally or vertically adjacent nodes.
- **Node features**: $V^0 = F \in \mathbb{R}^{s^0 \times d}$.
- **Node masks**: $M^0 = I \in \{0,1\}^{HW \times HW}$, so each supernode initially covers exactly one token.

Similarity between adjacent nodes $i$ and $j$ at layer $t$ is a convex combination of feature similarity and spatial similarity. The feature term is
$$
S^f_{ij} = V^t_i (V^t_j)^\top \in [-1,1].
$$
The spatial term is defined as a voting-based measure leveraging the global affinity matrix $A$ and the current subtree masks:
$$
S^s_{ij} = 1 - (1/HW)\, |(\cdots)|\, 1^\top \in [-1,1].
$$
The paper states that the omitted term represents the difference between the affinity distributions of nodes $i$ and $j$ over the original $HW$ tokens as induced by $A$ and the current subtree masks $M^t_i, M^t_j$. Intuitively, $S^s_{ij}$ is a normalized $L_1$ distance between the two nodes' affinity profiles over the original tokens; each original token "votes" for how similar $i$ and $j$ are.

The final edge similarity is
$$
S_{ij} = \omega_f S^f_{ij} + \omega_s S^s_{ij}, \qquad \omega_f + \omega_s = 1.
$$
The default weights reported in ablations are $\omega_f = 0.6$ and $\omega_s = 0.4$.

## 3. Layerwise agglomeration and hierarchical mask formation

For a current graph $G^t$, UniAP computes $S_{ij}$ for all edges $(i,j)\in E^t$ in parallel, applies a layer-wise threshold $\tau_t$, and marks edges satisfying $S_{ij} \ge \tau_t$. A strongly/connected component routine, denoted SCC in the description, is then run on the thresholded edge set to find groups of nodes connected by high-similarity edges. Each group becomes a supernode [2508.06995].

The grouping is data-dependent. Larger $\tau_t$ yields more, smaller groups; smaller $\tau_t$ yields fewer, larger groups. If $\Omega \in \{0,1\}^{s^t \times s^{t+1}}$ is the assignment matrix from nodes to supernodes, the updates are written as
$$
E^{t+1} = \Omega^\top E^t (\Omega^\top E^t)^\top,
$$
and
$$
M^{t+1} = \Omega^\top M^t.
$$
Thus, supernode masks are unions of their children's masks, and adjacency preserves spatial sparsity.

The feature update is not simple mean pooling. For supernode $i$ at layer $t+1$, the paper uses an attention-like weighted voting over the original features:
$$
V^{t+1}_i = \mathrm{L2N}\left\{ \mathrm{softmax}\left(\frac{M^{t+1}_i A}{\sigma}\right) F \right\},
\qquad i = 1,\ldots,s^{t+1},
$$
with default softmax temperature $\sigma = 0.07$. This update is described as emphasizing main regions and suppressing boundaries. A mean-pooling alternative,
$$
V^{t+1} = \mathrm{L2N}\{\Omega^\top V^t\},
$$
was observed to produce premature merges of spatially adjacent but semantically different regions; the voting via $A$ is intended to alleviate this.

An optional area threshold is applied to instance masks: if $\mathrm{area}(M^{t+1}_i) \ge \phi$, with default $\phi = 5$ pixels at stride 4, the supernode is kept as an instance mask and its feature is appended as the corresponding instance query feature.

UniAP distinguishes two pooling modes per layer:

- **Instance pooling** uses the sparse 4-neighbor adjacency. Every resulting mask is a single connected region.
- **Semantic pooling** allows masks to be composed of multiple disjoint regions. To obtain semantic masks, a fully-connected graph $G^\*$ is constructed at the same layer from the current supernode set $\{V^{t+1}, M^{t+1}\}$, with all-to-all edges, and the Identify+Merge steps are repeated with the same $\tau_t$.

The method is hierarchical. A decreasing threshold schedule
$$
\tau = [0.8, 0.7, 0.6, 0.5, 0.4]
$$
is used over $\Gamma$ layers, with the default $\Gamma = 5$. Early layers merge very similar local regions, often background textures; later layers merge semantically coherent but more diverse regions, such as parts of a person. The result is a set of multi-granular masks from fine to coarse. This suggests that UniAP functions not only as a pseudo-mask extractor but also as a hierarchical organizer of local-to-global structure.

## 4. Role in S2-UniSeg and Query-wise Self-Distillation

In S2-UniSeg, UniAP is embedded in a teacher–student pretraining loop rather than used as a standalone post-processing module [2508.06995]. The teacher branch processes the original global view using a ViT-base/8 backbone augmented by ViT-Adapter multi-scale blocks. The largest-scale stride-4 features are sent to UniAP, which generates instance and semantic masks at multiple granularities together with their corresponding query features. These form the teacher queries.

The student branch processes $\delta$ local views in a multi-crop regime. A single Mask2Former decoder with parameter sharing across tasks is used, and object queries are partitioned into instance and semantic groups with two special learnable tokens, [INS] and [SEM]. The student predicts local-view universal masks and query features.

Teacher pseudo-masks are cropped to each local view, and teacher queries without overlap are dropped. The student's predicted masks are then matched to the cropped teacher masks using mask Dice similarity, with matching performed independently for instance and semantic levels. This matching stage connects UniAP's online pseudo-mask generation to the supervision signal used for learning.

The associated pretext task is Query-wise Self-Distillation (QuerySD). If matched teacher and student query embeddings are
$$
Q_t \in \mathbb{R}^{L \times d}, \qquad Q_s \in \mathbb{R}^{L \times d},
$$
then a 3-layer MLP head maps each query to logits over $K$ pseudo-classes, followed by softmax with temperature normalization:
$$
p_t^l = \mathrm{softmax}(h(Q_t^l)/T_t), \qquad
p_s^l = \mathrm{softmax}(h(Q_s^l)/T_s),
$$
where $l=1,\ldots,L$. The loss is
$$
L_{\mathrm{QuerySD}} =
-\sum_{s=1}^{\delta}\sum_{l=1}^{L}\sum_{k=1}^{K} p_t^{l,k}\log p_s^{l,k}.
$$
The paper states that this distills teacher query distributions into student query distributions across local crops, aligning local-to-global correspondences for both instance and semantic queries and learning universal segmentation features rather than global [CLS] tokens or pooled features.

Teacher weights are updated as an exponential moving average of the student, with a cosine momentum schedule from 0.996 to 1. Because UniAP runs on the teacher's current features every iteration, pseudo-masks improve continuously during training.

## 5. Computational profile, hyperparameters, and scaling behavior

UniAP is designed for parallel execution. On the sparse 4-neighbor graph, all edge similarities can be computed with batched dot-products and matrix-vector operations; thresholding is elementwise; and the connected-components stage can be implemented with parallel union-find or BFS/DFS across edges above threshold. The description notes that standard libraries such as SciPy's `connected_components` can be used efficiently because the graph is sparse and has small degree [2508.06995].

The principal dense object is the affinity matrix $A = F F^\top$, whose complexity is given as $O((HW)^2 d)$ and whose memory is $O((HW)^2)$. The paper reports that empirical runtime remains in tens of milliseconds per $512 \times 512$ image with ViT-B/8 features due to stride-4 downsampling and efficient GPU matrix multiplication. It also notes that $A$ can be computed in blocks or stored in half precision if memory becomes a concern.

The reported per-image generation times on COCOval2017 at $512 \times 512$ using DINO ViT-B/8 features are summarized below.

| Method | Time per image |
|---|---:|
| UniAP | 0.045 s |
| TokenCut | 2.25 s |
| MaskCut | 4.72 s |
| UnSAM divide-and-conquer | 5.27 s |

The paper states that UniAP is approximately $100\times$ faster under this comparison, enabling per-iteration online pseudo-mask generation during training.

The default UniAP hyperparameters are $\Gamma = 5$, $\tau = [0.8, 0.7, 0.6, 0.5, 0.4]$, $\sigma = 0.07$, $\phi = 5$, $\omega_f = 0.6$, and $\omega_s = 0.4$. Ablations report that starting with thresholds that are too low, such as $[0.5,\ldots,0.1]$, coarsens too aggressively and hurts performance, while a fine-grained schedule such as $[0.9,\ldots,0.1]$ can be slightly better but slower. Setting $\omega_s=0$ improves speed but degrades quality significantly, including a reported $-3.8$ AP on UVO.

The surrounding S2-UniSeg defaults include a DINO-pretrained ViT-base/8 backbone with hidden dimension 768 and 12 layers, ViT-Adapter multi-scale blocks inserted every 4 blocks using deformable attention for linear complexity, instance queries = 150, semantic queries = 50, projection head hidden dimension 2048, default $K = 512$, and $\delta = 2$ local views for speed. Optimization uses AdamW, batch size 16, learning-rate warmup to $6.25\mathrm{e}{-4}$ for 10k steps, and total 160k steps.

## 6. Empirical performance, comparative position, and limitations

Under the same setting, S2-UniSeg is reported to outperform the state-of-the-art UnSAM model, with notable improvements of AP +6.9 on COCO, AR +11.1 on UVO, PixelAcc +4.5 on COCOStuff-27, and RQ +8.0 on Cityscapes [2508.06995]. In the ImageNet pretraining setting without SA-1B, the detailed results are given as follows: COCO class-agnostic instance segmentation AP 14.2 and AR 34.3; UVO unsupervised instance segmentation AP 16.2 and AR 32.1; COCOStuff-27 unsupervised semantic segmentation PixelAcc 68.4 and mIoU 36.1; and Cityscapes unsupervised panoptic segmentation RQ 29.7, SQ 60.1, and PQ 20.5. In the same section, these gains are presented against U2Seg.

The paper further emphasizes scaling enabled by UniAP's online design. On a 0.4M subset of SA-1B, it reports COCO class-agnostic AP 33.6 and AR 43.9, with gains over UnSAM of +2.2 AP and +1.9 AR. After scaling to a 2.0M-image subset, the reported metrics are COCO class-agnostic AP 36.4 and AR 46.3; COCO instance AP 30.7 and AR 40.3; UVO instance AP 24.3 and AR 36.4; COCOStuff-27 PixelAcc 76.8 and mIoU 39.7; COCO panoptic PQ 29.6, SQ 90.3, and RQ 30.2; and Cityscapes panoptic PQ 25.4, SQ 66.7, and RQ 35.0.

Ablations on UVO instance segmentation attribute part of this behavior directly to UniAP design choices. The default similarity weights $(\omega_s,\omega_f)=(0.4,0.6)$ yield AP 16.2 and AR 32.1, whereas $(0,1)$ degrades to AP 12.4 and AR 25.4. For the projection head, $K=128$ yields AP 13.6, $K=512$ yields AP 16.2, and $K=1024$ yields AP 16.9. Local crop counts $\delta \in \{2,4,6\}$ are reported as robust, with $\delta=2$ chosen to reduce matching overhead.

In comparative positioning, UniAP is distinguished from UnSAM by eliminating lengthy offline pseudo-mask generation and replacing discontinuous multi-stage training with online continuous optimization. It is distinguished from TokenCut and MaskCut by avoiding normalized cuts and eigenvector computation over dense affinity graphs. It is also contrasted with classical oversegmentation methods such as Felzenszwalb and SLIC, and with hierarchical agglomerative clustering, on the grounds that those methods typically operate on color gradients or superpixels rather than learned self-supervised features and do not yield unified multi-granular universal semantics.

The paper also identifies limitations and failure modes. Low-texture homogeneous regions can merge aggressively in early layers and may absorb thin structures nearby if the feature update is naive; the voting-based update helps but does not fully eliminate over-merge. Very fine boundaries or small objects may be suppressed by softmax weighting or filtered by the area threshold $\phi$. Extreme occlusions or overlapping instances across disjoint parts can challenge semantic pooling when spatial similarity dominates. Hyperparameter sensitivity is acknowledged for thresholds, weights, and $K$, although the reported defaults worked well across datasets. A practical implication is that UniAP's speed and online operation do not remove the need for careful control of granularity, memory footprint for $A$, and the balance between feature and spatial similarity.

Source: https://www.emergentmind.com/topics/fast-universal-agglomerative-pooling-uniap