SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shuyang Sun $^{1*}$  Weijun Wang $^{2}$  Qihang Yu $^{\dagger}$  Andrew Howard $^{2}$  Philip Torr $^{1}$  Liang-Chieh Chen $^{\dagger}$ $^{1}$ University of Oxford  $^{2}$ Google Research
This paper presents a new mechanism to facilitate the training of mask transformers for efficient panoptic segmentation, democratizing its deployment. We observe that due to the high complexity in the training objective of panoptic segmentation, it will inevitably lead to much higher penalization on false positive. Such unbalanced loss makes the training process of the end-to-end mask-transformer based architectures difficult, especially for efficient models. In this paper, we present ReMaX that adds relaxation to mask predictions and class predictions during the training phase for panoptic segmentation. We demonstrate that via these simple relaxation techniques during training, our model can be consistently improved by a clear margin without any extra computational cost on inference. By combining our method with efficient backbones like MobileNetV3-Small, our method achieves new state-of-the-art results for efficient panoptic segmentation on COCO, ADE20K and Cityscapes. Code and pre-trained checkpoints will be available at https://github.com/google-research/deeplab2.
Panoptic segmentation [36] aims to provide a holistic scene understanding [63] by unifying instance segmentation [21] and semantic segmentation [24]. The comprehensive understanding of the scene is obtained by assigning each pixel a label, encoding both semantic class and instance identity. Prior works adopt separate segmentation modules, specific to instance and semantic segmentation, followed by another fusion module to resolve the discrepancy [71, 12, 35, 70, 53, 42]. More recently, thanks to the transformer architecture [64, 4], mask transformers [66, 13, 76, 43, 72, 14, 73] are proposed for end-to-end panoptic segmentation by directly predicting class-labeled masks.
Although the definition of panoptic segmentation only permits each pixel to be associated with just one mask entity, some recent mask transformer-based methods [13, 76, 14, 40] apply sigmoid cross-entropy loss (i.e., not enforcing a single prediction via softmax cross-entropy loss) for mask supervision. This allows each pixel to be associated with multiple mask predictions, leading to an extremely unbalanced loss during training. As shown in Figure 1, when using the sigmoid cross-entropy loss to supervise the mask branch, the false-positive (FP) loss can be even  $10^{3} \times$  larger than the false-negative (FN) loss. Surprisingly, such unbalanced loss leads to better results than using softmax cross-entropy, which indicates that the gradients produced by the FP loss are still helpful for better performance.
However, the radical imbalance in the losses makes it difficult for the network to produce confident predictions, especially for efficient backbones [28, 57, 27], as they tend to make more mistakes given the smaller model size. Meanwhile, the training process will also become unstable due to the large scale loss fluctuation. To address this issue, recent approaches [4, 13, 14, 40] need to carefully clip
*Work done during internship at Google Research. Correspondence to: kevinsun@robots.ox.ac.uk
$\dagger$  Work done while at Google.
Figure 1: The histogram shows the ratio of false positives to false negatives for the cross-entropy loss, on a logarithmic scale. When using sigmoid as the activation function, the false positive loss is always over  $100 \times$  greater than the false negative, making the total loss to be extremely unbalanced.
Figure 2: The overall pipeline for mask-transformers.  $\otimes$  represents the matrix multiplication. Here  $\mathbf{x}_{\mathrm{pan}}^i$  and  $\mathbf{x}_{\mathrm{sem}}^i$  represent the input features of the  $i^{th}$  stage for panoptic and semantic heads respectively.
the gradients during training to a very small value like 0.01; otherwise, the loss would explode and the training would collapse. In this way, the convergence of the network will also be slower. A natural question thus emerges: Is there a way to keep those positive gradients, while better stabilizing the training of the network?
To deal with the aforementioned conflicts in the learning objectives, one naive solution is to apply weighted sigmoid cross entropy loss during training. However, simply applying the hand-crafted weights would equivalently scale the losses for all data points, which means those positive and helpful gradients will be also scaled down. Therefore, in this paper, we present a way that can adaptively adjust the loss weights by only adding training-time relaxation to mask-transformers [73, 66, 13, 14, 43, 76]. In particular, we propose two types of relaxation: Relaxation on Masks (ReMask) and Relaxation on Classes (ReClass).
The proposed ReMask is motivated by the observation that semantic segmentation is a relatively easier task than panoptic segmentation, where only the predicted semantic class is required for each pixel without distinguishing between multiple instances of the same class. As a result, semantic segmentation prediction could serve as a coarse-grained task and guide the semantic learning of panoptic segmentation. Specifically, instead of directly learning to predict the panoptic masks, we add another auxiliary branch during training to predict the semantic segmentation outputs for the corresponding image. The panoptic prediction is then calibrated by the semantic segmentation outputs to avoid producing too many false positive predictions. In this way, the network can be penalized less by false positive losses.
The proposed ReClass is motivated by the observation that each predicted mask may potentially contain regions involving multiple classes, especially during the early training stage, although each ground-truth mask and final predicted mask should only contain one target in the mask transformer framework [66]. To account for this discrepancy, we replace the original one-hot class label for each mask with a softened label, allowing the ground-truth labels to have multiple classes. The weights of each class is determined by the overlap of each predicted mask with all ground-truth masks.
By applying such simple techniques for relaxation to the state-of-the-art kMaX-DeepLab [73], our method, called ReMaX, can train the network stably without any gradient-clipping operation with a over  $10\times$  greater learning rate than the baseline. Experimental results have shown that our method not only speeds up the training by  $3\times$ , but also leads to much better results for panoptic segmentation. Overall, ReMaX sets a new state-of-the-art record for efficient panoptic segmentation. Notably, for efficient backbones like MobileNetV3-Small and MobileNetV3-Large [27], our method can outperform the strong baseline by 4.9 and 5.2 in PQ on COCO panoptic for short schedule training; while achieves 2.9 and 2.1 improvement in PQ for the final results (i.e., long schedules). Meanwhile,
our model with a Axial-ResNet50 (MaX-S) [65] backbone outperforms all state-of-the-art methods with  $3 \times$  larger backbones like ConvNeXt-L [47] on Cityscapes [17]. Our model can also achieve the state-of-the-art performance when compared with the other state-of-the-art efficient panoptic segmentation architectures like YOSO [29] and MaskConver [55] on COCO [44], ADE20K [77] and Cityscapes [17] for efficient panoptic segmentation.
Mask Transformers for image segmentation. Recent advancements in image segmentation has proven that Mask Transformers [66], which predict class-labeled object masks through the Hungarian matching of predicted and ground truth masks using Transformers as task decoders [64, 4], outperform box-based methods [35, 70, 54] that decompose panoptic segmentation into multiple surrogate tasks, such as predicting masks for detected object bounding boxes [23] and fusing instance and semantic segmentation [48, 10] with merging modules [42, 53, 45, 71, 12, 41]. The Mask Transformer based methods rely on converting object queries to mask embedding vectors [32, 62, 67], which are then multiplied with pixel features to generate predicted masks. Other approaches such as Segmenter [59] and MaskFormer [14] have also used mask transformers for semantic segmentation. K-Net [76] proposes dynamic kernels for generating masks. CMT-DeepLab [72] suggests an additional clustering update term to improve transformer's cross-attention. Panoptic Segformer [43] enhances mask transformers with deformable attention [79]. Mask2Former [14] adopts masked-attention, along with other technical improvements such as cascaded transformer decoders [4], deformable attention [79], and uncertainty-based point supervision [37], while kMaX-DeepLab [73] employs k-means cross-attention. OneFormer [31] extends Mask2Former with a multi-task train-once design. Our work builds on top of the modern mask transformer, kMaX-DeepLab [73], and adopts novel relaxation methods to improve model capacity.
The proposed Relaxation on Masks (ReMask) is similar to the masked-attention in Mask2Former [14] and the k-means attention in kMaX-DeepLab [73] in the sense that we also apply pixel-filtering operations to the predicted masks. However, our ReMask operation is fundamentally distinct from theirs in several ways: (1) we learn the threshold used to filter pixels in panoptic mask predictions through a semantic head during training, while both masked-attention [14] and k-means attention [73] use either hard thresholding or argmax operation on pixel-wise confidence for filtering; (2) our approach relaxes the training objective by applying a pixel-wise semantic loss on the semantic mask for ReMask, while they do not have explicit supervision for that purpose; and (3) we demonstrate that ReMask can complement k-means attention in Section 4.
Acceleration for Mask Transformers for efficient panoptic segmentation. DETR [4] successfully proves that Transformer-based approaches can be used as decoders for panoptic segmentation, however, it still suffers from the slow training problem which requires over 300 epochs for just one go. Recent works [14, 73, 79, 50] have found that applying locality-enhanced attention mechanism can help to boost the speed of training for instance and panoptic segmentation. Meanwhile, some other works [76, 43, 33] found that by removing the bi-partite matching for stuff classes and applying a separate group of mask queries for stuff classes can also help to speed up the convergence. Unlike them, which apply architectural level changes to the network, our method only applies training-time relaxation to the framework, which do not introduce any extra cost during testing. Apart from the training acceleration, recent works [26, 29, 12, 55, 51] focus on how to make the system for panoptic segmentation more efficient. However, all these works focus on the modulated architecultural design while our approach focuses on the training pipeline, which should be two orthogonal directions.
Coarse-to-fine refinement for image segmentation. In the field of computer vision, it is a common practice to learn representations from coarse to fine, particularly in image segmentation. For instance, DeepLab [6, 8] proposes a graph-based approach [38, 7] that gradually refines segmentation results. Recently, transformer-based methods for image and video segmentation such as [66, 14, 76, 69, 43, 20, 74, 78] have also adopted a multi-stage strategy to iteratively improve predicted segmentation outcomes in transformer decoders. The concept of using coarse-grained features (e.g., semantic segmentation) to adjust fine-grained predictions (e.g., instance segmentation) is present in certain existing works, including [11, 2, 3]. However, these approaches can lead to a substantial increase in model size and number of parameters during both training and inference.
Figure 3: ReMask Operation. Modules, representations and operations rendered in gray are not used in testing.  $\otimes$  and  $\odot$  represent the matrix multiplication and Hadamard multiplication and + means element-wise sum. The  $\times$  symbol and "stop grad" mean that no gradient is flown to  $\mathbf{m}_{\mathrm{sem}}$  from  $\mathcal{L}_{\mathrm{pan}}$  during training.
Figure 4: Demonstration on how ReClass works. We utilize the mask rendered in blue as an example. Our ReClass operation aims to soften the class-wise ground-truth by considering the degree of overlap between the prediction and ground-truth mask. The blue mask intersects with both masks of "baseball glove" and "person", so the final class weights contain both and the activation of "person" in the prediction will no longer be regarded as a false positive case during training.
By contrast, our ReMaX focuses solely on utilizing the coarse-fine hierarchy for relaxation without introducing any additional parameters or computational costs during inference.
Regularization and relaxation techniques. The proposed Relaxation on Classes (ReClass) involves adjusting label weights based on the prior knowledge of mask overlaps, which is analogous to the re-labeling strategy employed in CutMix-based methods such as [75, 5, 60], as well as label smoothing [61] used in image classification. However, the problem that we are tackling is substantially different from the above label smoothing related methods in image classification. In image classification, especially for large-scale single-class image recognition benchmarks like ImageNet [56], it is unavoidable for images to cover some of the content for other similar classes, and label smoothing is proposed to alleviate such labelling noise into the training process. However, since our approach is designed for Mask Transformers [66, 13, 14, 73, 72] for panoptic segmentation, each image is precisely labelled to pixel-level, there is no such label noise in our dataset. We observe that other than the class prediction, the Mask Transformer approaches also introduce a primary class identification task. The proposal of ReClass operation reduces the complexity for the classification task in Mask Transformers. Prior to the emergence of Mask Transformers, earlier approaches did not encounter this issue as they predicted class labels directly on pixels instead of masks.
Before delving into the details of our method, we briefly recap the framework of mask transformers [66] for end-to-end panoptic segmentation. Mask Transformers like [66, 14, 76, 69, 43] perform both semantic and instance segmentation on the entire image using a single Transformer-based model. These approaches basically divide the entire framework into 3 parts: a backbone for feature extraction, a pixel decoder with feature pyramid that fuses the feature generated by the backbone, and a transformer mask decoder that translates features from the pixel decoder into panoptic masks and their corresponding class categories.
In the transformer decoder, a set of mask queries is learnt to segment the image into a set of masks by a mask head and their corresponding categories by a classification head. These queries are updated within each transformer decoder (typically, there are at least 6 transformer decoders) by the cross-attention mechanism [64] so that the mask and class predictions are gradually refined. The set of predictions are matched with the ground truth via bipartite matching during training; while these queries will be filtered with different thresholds as post-processing during inference. We follow the same post-processing as kMaX-DeepLab [73].
3.1 Relaxation on Masks (ReMask)
The proposed Relaxation on Masks (ReMask) aims to ease the training of panoptic segmentation models. Panoptic segmentation is commonly viewed as a more intricate task than semantic segmentation, since it requires the model to undertake two types of segmentation (namely, instance segmentation and semantic segmentation). In semantic segmentation, all pixels in an image are labeled with their respective class, without distinguishing between multiple instances (things) of the same class. As a result, semantic segmentation is regarded as a more coarse-grained task when compared to panoptic segmentation. Current trend in panoptic segmentation is to model things and stuff in a unified framework and resorts to train both the coarse-grained segmentation task on stuff and the more fine-grained segmentation task on things together using a stricter composite objective on things, which makes the model training more difficult. We thus propose ReMask to exploit an auxiliary semantic segmentation branch to facilitate the training.
Definition. Here we first define  $H, W$  as the height and width of the feature,  $N_{Q}$  as the number of mask queries.  $N_{C}$  denotes the number of semantic classes for the target dataset,  $d_{q}$  is the number of channels for the query representation, and  $d_{\mathrm{sem}}$  is the number of channels for the input of semantic head. As shown in Figure 2 and 3, given a mask representation  $\mathbf{x}_{\mathrm{pan}} \in \mathbb{R}^{HW \times N_{Q}}$ , we apply a panoptic mask head to generate panoptic mask logs  $\mathbf{m}_{\mathrm{pan}} \in \mathbb{R}^{HW \times N_{Q}}$ . A mask classification head to generate the corresponding classification result  $\mathbf{p} \in \mathbb{R}^{N_{Q} \times N_{C}}$  is applied for each query representation  $\mathbf{q} \in \mathbb{R}^{N_{Q} \times d_{q}}$ . A semantic head is applied after the semantic feature  $\mathbf{x}_{\mathrm{sem}} \in \mathbb{R}^{HW \times d_{\mathrm{sem}}}$  from the pixel decoder to produce a pixel-wise semantic segmentation map  $\mathbf{m}_{\mathrm{sem}} \in \mathbb{R}^{HW \times N_{C}}$  assigning a class label to each pixel. As for the structure for semantic head, we apply an ASPP module [9] and a  $1 \times 1$  convolution layer afterwards to transform  $d_{\mathrm{sem}}$  channels into  $N_{C}$  channels as the semantic prediction. Note that the whole auxiliary semantic branch will be skipped during inference as shown in Figure 3. Since the channel dimensionality between  $\mathbf{m}_{\mathrm{sem}}$  and  $\mathbf{m}_{\mathrm{pan}}$  is different, we map the semantic masks into the panoptic space by:
$$
\widehat {\mathbf {m}} _ {\text {s e m}} = \sigma \left(\mathbf {m} _ {\text {s e m}}\right) \sigma \left(\mathbf {p} ^ {\top}\right), \tag {1}
$$
where  $\sigma (\cdot)$  function represents the sigmoid function that normalizes the logits into interval [0, 1]. Then we can generate the relaxed panoptic outputs  $\hat{\mathbf{m}}_{\mathrm{pan}}$  in the semantic masking process as follows:
$$
\widehat {\mathbf {m}} _ {\text {p a n}} = \mathbf {m} _ {\text {p a n}} + \left(\widehat {\mathbf {m}} _ {\text {s e m}} \odot \mathbf {m} _ {\text {p a n}}\right), \tag {2}
$$
where the  $\odot$  represents the Hadamard product operation. Through the ReMask operation, the false positive predictions in  $\mathbf{m}_{\mathrm{pan}}$  can be suppressed by  $\widehat{\mathbf{m}}_{\mathrm{sem}}$ , so that during training each relaxed mask query can quickly focus on areas of their corresponding classes. Here we apply identity mapping to keep the original magnitude of  $\mathbf{m}_{\mathrm{pan}}$  so that we can remove the semantic branch during testing. This makes ReMask as a complete relaxation technique that does not incur any overhead cost during testing. The re-scaled panoptic outputs  $\widehat{\mathbf{m}}_{\mathrm{pan}}$  will be supervised by the losses  $\mathcal{L}_{\mathrm{pan}}$ .
Stop gradient for a simpler objective to  $\widehat{\mathbf{m}}_{\mathrm{sem}}$ . In order to prevent the losses designed for panoptic segmentation from affecting the parameters in the semantic head, we halt the gradient flow to  $\mathbf{m}_{\mathrm{sem}}$  as illustrated in Figure 3. This means that the semantic head is solely supervised by a semantic loss  $\mathcal{L}_{\mathrm{sem}}$ , so that it can focus on the objective of semantic segmentation, which is a less complex task.
How does ReMask work? As defined above, there are two factors that ReMask operation helps training, (1) the Hadamard product operation between the semantic outputs and the panoptic outputs that helps to suppress the false positive loss; and (2) the relaxation on training objectives that trains the entire network simultaneously with consistent (coarse-grained) semantic predictions. Since the semantic masking can also enhance the locality of the transformer decoder like [14, 73], we conducted experiments by replacing  $\mathbf{m}_{\mathrm{sem}}$  with ground truth semantic masks to determine whether it is the training relaxation or the local enhancement that improves the training. When  $\mathbf{m}_{\mathrm{sem}}$  is assigned with ground truth, there will be no  $\mathcal{L}_{\mathrm{sem}}$  applied to each stage, so that  $\mathbf{m}_{\mathrm{pan}}$  is applied with the most accurate local enhancement. In this way, there are large amounts of false positive predictions masked by the ground truth semantic masks, so that the false positive gradient will be greatly reduced. The results will be reported in Section 4. The semantic masking can be viewed as local enhancement as it would suppress the extreme false-positive predictions via a simple masking operation.
Figure 5: Performance on COCO val compared to the baseline kMaX-DeepLab [73]. ReMaX can lead to  $3 \times$  faster convergence compared to the baseline, and can improve the baselines by a clear margin. The performance of ResNet-50 can be further improved to 54.2 PQ when the model is trained for 200K iterations.
Method	Backbone	Resolution	FPS	PQ
Panoptic-DeepLab [12]	MNV3-L [27]	641×641	26.3	30.0
Panoptic-DeepLab [12]	R50 [22]	641×641	20.0	35.1
Real-time [26]	R50 [22]	800×1333	15.9	37.1
MaskConver [55]	MN-MH [16]	640×640	40.2	37.2
MaskFormer [14]	R50 [22]	800×1333	17.6	46.5
YOso [29]	R50 [22]	800×1333	23.6	48.4
YOso [29]	R50 [22]	512×800	45.6	46.4
kMaX-DeepLab [73]	R50 [22]	1281×1281	16.3	53.0
ReMaX-T†	MNV3-S [27]	641×641	108.7	40.4
ReMaX-S†	MNV3-L [27]	641×641	80.9	44.6
ReMaX-M‡	R50 [22]	641×641	51.9	49.1
ReMaX-B	R50 [22]	1281×1281	16.3	54.2
Table 1: Comparison with other state-of-the-art efficient models (≥ 15 FPS) on COCO val set. The Pareto curve is shown in Figure 6 (b). The FPS of all models are evaluated on a NVIDIA V100 GPU with batch size 1.  $\dagger \ddagger$  represent the application of efficient pixel and transformer decoders. Please check the appendix for details.
3.2 Relaxation on Classes (ReClass)
Mask Transformers [66, 14, 73, 43] operate under the assumption that each mask prediction corresponds to a single class, and therefore, the ground truth for the classification head are one-hot vectors. However, in practice, each imperfect mask predicted by the model during the training process may intersect with multiple ground truth masks, especially during the early stage of training. As shown in Figure 4, the blue mask, which is the mask prediction, actually covers two classes ("baseball glove" and "person") defined in the ground truth. If the class-wise ground truth only contains the class "baseball glove", the prediction for "person" will be regarded as a false positive case. However, the existence of features of other entities would bring over-penalization that makes the network predictions to be under-confident.
To resolve the above problem, we introduce another relaxation strategy on class logits, namely Class-wise Relaxation (ReClass), that re-assigns the class confidence for the label of each predicted mask according to the overlap between the predicted and ground truth semantic masks. We denote the one-hot class labels as  $\mathbf{y}$ , the ground truth binary semantic masks as  $\mathcal{S} = [\mathbf{s}_0,\dots,\mathbf{s}_{HW}]\in \{0,1\}^{HW\times N_C}$ , the supplement class weights is calculated by:
$$
\mathbf {y} _ {m} = \frac {\sigma \left(\mathbf {m} _ {\text {p a n}}\right) ^ {\intercal} \mathcal {S}}{\sum_ {i} ^ {H W} \mathbf {s} _ {i}}, \tag {3}
$$
where  $\mathbf{y}_m$  denotes the label weighted by the normalized intersections between the predicted and the ground truth masks. With  $\mathbf{y}_m$ , we further define the final class weight  $\widehat{\mathbf{y}} \in [0,1]^{N_C}$  as follows:
$$
\widehat {\mathbf {y}} = \eta \mathbf {y} _ {m} + (1 - \eta \mathbf {y} _ {m}) \mathbf {y}, \tag {4}
$$