GrooMeD-NMS: Differentiable 3D NMS

• The paper introduces GrooMeD-NMS, a differentiable closed-form approximation of greedy NMS that integrates suppression into training for monocular 3D detection.
• It reformulates NMS as a matrix operation with unsupervised grouping and masking, aligning training and inference by penalizing ranking and suppression errors.
• Experimental results on KITTI benchmarks show improved 3D and BEV average precision, evidencing its practical benefits in optimizing detection scores and geometry.

Searching arXiv for the primary paper and closely related works mentioned in the provided data. GrooMeD-NMS is a differentiable, closed-form approximation to greedy non-maximum suppression (NMS) that is inserted into training for monocular 3D object detection, so the detector is optimized not only on raw proposals before NMS but also on the post-NMS predictions that matter at test time. Introduced in "GrooMeD-NMS: Grouped Mathematically Differentiable NMS for Monocular 3D Object Detection" (Kumar et al., 2021), the method reformulates NMS as a matrix operation and then uses unsupervised grouping and masking to obtain a simple closed-form expression. Its stated purpose is to address the train/inference mismatch created when detectors are trained on dense proposals before suppression but evaluated only on boxes surviving NMS.

1. Problem formulation and motivation

Modern detectors, including monocular 3D detectors, typically predict many overlapping boxes per object and apply NMS at inference to prune duplicates and keep one box per object. During training, however, the standard pipeline computes losses on the dense set of proposals before NMS. The method is explicitly motivated by this discrepancy: training optimizes scores and regression on all boxes before suppression, whereas inference evaluates only boxes surviving NMS (Kumar et al., 2021).

The paper argues that this mismatch is particularly severe for monocular 3D detection because the correlation between classification confidence and 3D localization quality is weaker than for 2D localization. Its oracle experiment is described as showing that if one replaces NMS scores with oracle 3D overlap quality, AP improves dramatically, much more than analogous 2D oracle scoring. This motivates a training objective that does not merely regress boxes and classify proposals, but also encourages the network to assign higher scores to the best-localized 3D box for each object.

Within that framing, GrooMeD-NMS is not presented as a generic differentiable relaxation of suppression in isolation. It is designed to include NMS inside the training graph, let gradients flow through it, and impose a loss after NMS that directly penalizes poor ranking and suppression decisions. A plausible implication is that the method treats ranking quality under suppression as a first-class training target rather than as an emergent byproduct of proposal scoring.

2. Mathematical formulation of differentiable suppression

Classical NMS is described as a greedy set procedure: choose the highest-scoring box via $\arg\max$, add it to the keep set, suppress overlapping boxes according to a hard threshold, and repeat. In the paper’s notation, the inputs are scores $S=\{s_i\}_{i=1}^n$, an overlap matrix $O=[o_{ij}]$, threshold $N_t$, pruning function $p$, and temperature $\tau$ for soft pruning (Kumar et al., 2021). The difficulty of incorporating standard NMS into end-to-end training is attributed to four properties: discrete $\arg\max$ selection, hard threshold suppression, iterative greedy set-valued logic, and lack of natural parallelizability.

The paper develops a progression from greedy suppression to a differentiable approximation. It first writes a classical greedy-like rescoring form and then generalizes it to account for the influence of all higher-ranked boxes: $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$ Because this expression is recursive and multiplicative, the authors expand the product, drop higher-order terms, and obtain a first-order additive approximation: $r_i \approx s_i - \sum_{j=1}^{i-1} o_{ij}r_j.$ To prevent negative scores, they clamp from below: $$r_i \approx \max\left(s_i-\sum_{j=1}^{i-1} o_{ij}r_j,\;0\right).$$ The paper also motivates the sum by relaxing the “suppressed if it overlaps with any selected higher box” logic from a logical OR into summation (Kumar et al., 2021).

After sorting boxes by decreasing score, the recursive update is written in matrix form. Let $S=\{s_i\}_{i=1}^n$0 denote the lower-triangular part of the overlap matrix without diagonal, and let the prune matrix be $S=\{s_i\}_{i=1}^n$1, with $S=\{s_i\}_{i=1}^n$2 applied elementwise. Then the recursive form becomes

$S=\{s_i\}_{i=1}^n$3

Ignoring the max yields a linear system with closed-form solution

$S=\{s_i\}_{i=1}^n$4

followed by clipping to $S=\{s_i\}_{i=1}^n$5: $S=\{s_i\}_{i=1}^n$6 The stated reason invertibility holds is that $S=\{s_i\}_{i=1}^n$7 is lower triangular with ones on the diagonal, hence always full rank and invertible. This gives a mathematically differentiable mapping from scores and overlaps to post-NMS scores (Kumar et al., 2021).

3. Grouping, masking, and the final GrooMeD-NMS expression

The defining feature of GrooMeD-NMS is the transition from the full matrix inverse to a grouped and masked approximation. The method partitions boxes into groups $S=\{s_i\}_{i=1}^n$8 based on 2D overlaps, where each group is intended to correspond roughly to one object hypothesis neighborhood. Suppression is then restricted to within-group interactions: $S=\{s_i\}_{i=1}^n$9 Grouping is performed after sorting by score. The top remaining box acts as a seed, boxes whose overlap with that top box exceeds the threshold $O=[o_{ij}]$0 are collected into the next group, the group is truncated to maximum size $O=[o_{ij}]$1, and grouped boxes are removed before the process repeats. This grouping is unsupervised because it uses only predicted box overlaps and scores, not ground-truth object identities (Kumar et al., 2021).

The masking step is the crucial simplification. Within a group, classical NMS compares all other boxes against the current top box only. To mimic this, the method introduces a binary mask $O=[o_{ij}]$2 that keeps only the column of the group’s top-ranked box and zeros the rest. Thus only one column of $O=[o_{ij}]$3 is non-zero. The paper states that this makes $O=[o_{ij}]$4 a Frobenius matrix / Gaussian transformation, whose inverse is simply

$O=[o_{ij}]$5

Substituting this identity yields the final grouped and masked formula: $O=[o_{ij}]$6 or equivalently

$O=[o_{ij}]$7

This is the final GrooMeD-NMS expression reported in the paper (Kumar et al., 2021).

The paper characterizes this expression as differentiable with respect to the scores and overlap entries, except for the hard sorting and group membership decisions. That permits gradients through both the score head and the box geometry via the overlaps. This suggests that the method is designed not only to recalibrate confidence under suppression but also to shape predicted geometry indirectly through the overlap matrix.

4. Detector integration and learning objective

GrooMeD-NMS is integrated into a monocular 3D detector based on the PyTorch implementation of Kinematic-3D, using its best image model as baseline. The reported base configuration consists of a DenseNet-121 backbone, the M3D-RPN detector structure, and the binning and self-balancing confidence components from Kinematic-3D (Kumar et al., 2021).

The detector predicts class scores, 2D box outputs, 3D box outputs, and self-balancing confidence. The self-balancing confidence is used as the NMS score $O=[o_{ij}]$8 during training, and this confidence is described as intended to reflect localization quality better than plain class confidence. GrooMeD-NMS takes those scores together with an overlap matrix $O=[o_{ij}]$9 built from predicted 2D boxes. The pipeline is summarized in the source as

$N_t$0

The post-NMS supervision uses a “best” box definition for each ground-truth $N_t$1: $N_t$2 Targets are then assigned by selecting the maximizer of this quality measure subject to threshold $N_t$3: $N_t$4 The paper describes this as extending prior “best 2D box” logic to 3D (Kumar et al., 2021).

Because positive post-NMS “best boxes” are extremely sparse, the post-NMS objective is not BCE but a modified AP loss called Imagewise AP: $N_t$5 The final training loss is

$N_t$6

with

$N_t$7

and $N_t$8. The pre-NMS loss is given in the source in expanded form and combines classification, 2D box regression, confidence-weighted 3D losses, orientation terms, and a confidence regularizer. The central conceptual point is that the network receives supervision based on which boxes remain after suppression, not only on raw proposal quality (Kumar et al., 2021).

5. Geometry, pruning functions, and implementation assumptions

For suppression and grouping, GrooMeD-NMS uses 2D overlap: $N_t$9 The source states that 2D overlap is easy to compute and provides meaningful clustering, while the post-NMS target assignment is based on the combined 2D–3D quality

$p$0

To obtain nonzero signal when regular 3D IoU is often zero, especially early in training or on difficult examples, the paper uses a generalized 3D IoU-like quantity: $p$1 It states that the hull volume is built from the BEV hull and the $p$2-dimension hull (Kumar et al., 2021).

Three pruning functions $p$3 are considered: $p$4

$p$5

and

$p$6

The paper reports that linear pruning works best. It speculates that the exponential and sigmoidal forms can suffer vanishing gradients near overlap $p$7 or $p$8 (Kumar et al., 2021).

The method includes several explicit approximations and assumptions. Hard sorting remains, because soft sort methods were found to be $p$9, too temperature-sensitive, and to produce unreliable gradients in the reported setup. The derivation also uses first-order truncation by dropping $\tau$0 terms, replaces $\tau$1 with $\tau$2 in the supplement to obtain a cleaner form, and substitutes the recursive max-form with a linear solve plus clipping. Grouping depends on threshold $\tau$3 and maximum group size $\tau$4, with default values $\tau$5 and $\tau$6; the valid-score threshold after rescoring is $\tau$7 by default. The paper states that the layer provides gradient through both scores $\tau$8 and overlap matrix $\tau$9, and characterizes GrooMeD-NMS as a 1-layer, matrix-based method, parallelizable up to groups, with complexity $\arg\max$0 in the number of groups as a sequential factor (Kumar et al., 2021).

6. Experimental evaluation on KITTI

The main benchmark is KITTI monocular 3D car detection. Evaluation is reported on the KITTI test split, Val 1 split with 3712 train / 3769 val, and Val 2 split with 3682 train / 3799 val. The training setup listed in the source includes warmup plus full training stages, Adam, batch size 2, random horizontal flip, $\arg\max$1, $\arg\max$2, $\arg\max$3, $\arg\max$4, and $\arg\max$5; the inference score is class score multiplied by predicted confidence (Kumar et al., 2021).

On KITTI test cars at $\arg\max$6, GrooMeD-NMS reports 3D AP of 18.10 / 12.32 / 9.65 and BEV AP of 26.19 / 18.27 / 14.05 for Easy/Moderate/Hard. The source compares this with Kinematic (Video), reported as 19.07 / 12.72 / 9.17 3D AP, and notes that GrooMeD-NMS outperforms it on Hard despite being image-only (Kumar et al., 2021).

On KITTI Val 1 at $\arg\max$7 for Cars, GrooMeD-NMS reports 3D AP of 19.67 / 14.32 / 11.27 and BEV AP of 27.38 / 19.75 / 15.92, compared with Kinematic (Image) at 18.28 / 13.55 / 10.13 3D AP and 25.72 / 18.82 / 14.48 BEV AP. At $\arg\max$8, GrooMeD-NMS reports 3D AP of 55.62 / 41.07 / 32.89 and BEV AP of 61.83 / 44.98 / 36.29. The source emphasizes that the largest gains come on Moderate and Hard cases (Kumar et al., 2021).

On KITTI Val 2, the paper states that GrooMeD-NMS is best in all reported categories, with 3D AP @0.7 of 14.72 / 10.87 / 7.67 and BEV AP @0.7 of 22.03 / 16.05 / 11.93 (Kumar et al., 2021).

A central experimental claim concerns training versus inference. When trained without GrooMeD and only swapping inference NMS, classical NMS, Soft-NMS, Distance-NMS, and GrooMeD at test time give nearly identical results. When trained with GrooMeD-NMS, performance is significantly higher, and the choice of inference NMS barely matters. The paper presents this as evidence that the gain comes from training through NMS rather than from replacing classical NMS at inference (Kumar et al., 2021).

Additional findings reported in the source include a score–3D-IoU correlation after NMS of 0.345, higher than the M3D-RPN and Kinematic baselines; consistent gains across object depths and IoU thresholds; and only a slight increase in full training time, from 8h to 8.5h after warmup. Inference time is reported as 0.12 ms/image for classical NMS and 0.15 ms/image for GrooMeD-NMS, with the practical observation that models trained with GrooMeD-NMS can use classical NMS at test time with almost no change in accuracy (Kumar et al., 2021).

7. Relation to prior NMS variants, limitations, and scope

The paper positions GrooMeD-NMS against several classes of prior methods. Soft-NMS changes hard suppression to soft decay but remains greedy and inference-only, and thus does not address the end-to-end train/inference mismatch. Neural or learned NMS approaches, including those attributed in the source to Hosang et al. and Prokudin et al., use separate networks, multiple layers, or message passing to approximate NMS behavior rather than yielding a clean mathematical closed form. Structured or optimization-based formulations such as QUBO, point-process, MAP, and structured-SVM methods are described as often not integrated as end-to-end differentiable modules in modern detectors. Adaptive-NMS is described as predicting thresholds adaptively while still relying on standard NMS logic, and Distance-NMS as a 3D-specific inference refinement for the $\arg\max$9-coordinate that remains inference-time and hard-suppression-based (Kumar et al., 2021).

Within that comparison, the paper claims five specific novelties: first closed-form mathematically differentiable NMS integrated into end-to-end training for object detection; a matrix-based formulation of NMS; unsupervised grouping and masking to simplify the inverse into a closed-form subtraction; a single-layer formulation rather than a separate multi-layer neural approximator; and successful application to monocular 3D detection, where the mismatch is especially harmful (Kumar et al., 2021). The emphasis is therefore not merely on “differentiable suppression,” but on a closed-form matrix approximation of NMS itself.

The method also has explicit limitations and failure modes in the source. Hard sorting remains non-differentiable because soft sorting was found unstable in this setting. GrooMeD-NMS is an approximation to classical NMS rather than an exact differentiable rendering, since it relies on OR-to-sum relaxation, first-order truncation, linearization, clipping, and grouping/masking simplifications. Grouping is heuristic, built via 2D overlap with the current top box and capped at size $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$0. Suppression uses 2D overlap rather than 3D overlap, which is pragmatic for monocular setups but may limit fidelity when 2D overlap does not reflect 3D duplication well. Best-box positives are extremely sparse, requiring a specialized post-NMS ranking loss, and the method depends on several hyperparameters, including $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$1, $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$2, $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$3, $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$4, and $$r_i = s_i \prod_{j=1}^{i-1}\left(1-o_{ij}r_j\right).$$5 (Kumar et al., 2021).

The paper suggests extending GrooMeD-NMS beyond monocular 3D detection to LiDAR-based 3D detection, pedestrian detection, and more general object detection tasks. A plausible implication is that the grouped matrix formulation could serve as a template for other post-processing layers where train-time and test-time operators differ, provided that the relevant ranking and suppression structure can be expressed through overlaps and groupwise masking.