TreemapMix: Dirichlet-Controlled Multi-Image Augmentation for Probability and Ordinal Supervision

AI summary

Overview

Research area: Computer vision — data augmentation and label supervision for image classification.

Technical level: Intermediate. Readers should be comfortable with soft-label training, cross-entropy, and the idea of listwise/ranking losses; the geometric construction itself is easy to follow.

Scope: The paper introduces and evaluates TreemapMix, a multi-image augmentation that samples region sizes from a Dirichlet distribution, lays them out as an area-proportional treemap, and supervises the result either as soft class probabilities (TreemapMix-SCE) or as an ordinal ranking (TreemapMix-PL).

What This Paper Is About

Most mixed-sample augmentation methods either blend just two images (Mixup, CutMix) or use fixed small layouts with a set number of sources (RICAP, Mosaic), which limits how much control a practitioner has over how many images appear, how much area each occupies, and what kind of label signal results. TreemapMix addresses this by composing K single-label images into one canvas whose region sizes are drawn from a Dirichlet distribution and realized as a squarified treemap, so region area becomes an explicit, tunable quantity. The paper then asks what that area information should supervise: exact probability matching, or a coarser ordinal statement that bigger regions should score higher.

Key Contributions

1. The TreemapMix augmentation itself. A Dirichlet-controlled K-way image-mixing method that assigns sampled weights to non-overlapping, area-proportional treemap regions, supports a variable number of sources, and preserves the training sample count by emitting K cyclic views per group of K images.

2. Two supervision objectives from one composition. TreemapMix-SCE treats region areas as soft class-probability targets under a KL/soft cross-entropy objective; TreemapMix-PL converts the same areas into an area-induced ordering of target classes and optimizes a Plackett-Luce listwise objective.

3. An analysis of the Dirichlet concentration parameter α. The paper characterizes how α controls target variability, probability-matching supervision, and ranking emphasis, including a result that as α increases the sorted weights concentrate near 1/K, so the distinct-label PL loss approaches 1/K times an unweighted top-K ListMLE-style objective.

4. An empirical study on ImageNet-1K and beyond. Accuracy, calibration (ECE over 15 bins), area-logit alignment diagnostics, COCO detection and instance-segmentation transfer, plus CIFAR-100, additional architectures, long-tailed recognition, 320-epoch training, and loader throughput/memory measurements.

Main Findings

• TreemapMix-PL has the best Top-1 accuracy on every evaluated ImageNet-1K backbone. Under the matched 110-epoch protocol it reaches 78.01 ± 0.13 (ViT-Wee), 80.12 ± 0.04 (ViT-Little), 81.05 ± 0.14 (ViT-Medium), and 81.69 ± 0.16 (ViT-Betwixt), ahead of all 13 baseline configurations.

• TreemapMix-SCE has the lowest calibration error. ECE is 1.04 ± 0.35, 1.89 ± 0.23, 1.95 ± 0.25, and 2.31 ± 0.41 percentage points across the same four backbones, substantially below every evaluated mixed-sample baseline (Mixup+CutMix, for instance, ranges from 7.53 to 9.05).

• The two objectives trade off rather than combine cleanly. The convex combination (TreemapMix-mixed, β = 0.5) lands between the two on both metrics — for example 77.65 accuracy / 2.63 ECE on ViT-Wee — and the paper states it does not consistently improve on both objectives.

• No single baseline is runner-up everywhere. The strongest baseline varies by backbone: single-image augmentation on ViT-Wee, SaliencyMix on ViT-Little, and ResizeMix on ViT-Medium and ViT-Betwixt.

• PL also produces better area-logit alignment. On 500 held-out mixtures per K ∈ {3, 4, 5, 6} with ViT-Betwixt checkpoints, PL has the highest pairwise rank accuracy, Spearman correlation, and largest-area top-class accuracy. RICAP is competitive and beats TreemapMix-SCE on largest-area top-class accuracy at every evaluated K despite also using area-weighted soft targets.

• Pretrained features transfer to COCO. After ViTDet-style Mask R-CNN fine-tuning for 30 epochs at 256×256, TreemapMix-PL is best: 24.32 box AP / 20.71 mask AP on ViT-Wee and 29.71 / 25.75 on ViT-Betwixt, versus 23.66 / 20.34 and 29.04 / 25.25 for Mixup+CutMix and 23.08 / 19.82 and 27.86 / 24.21 for RICAP.

• The accuracy/calibration separation holds on other architectures. On ConvNeXtV2-Nano, DeiT-Tiny, and Swin-Tiny under the same 110-epoch protocol, SCE has the lowest ECE (2.74, 1.32, 2.52) and PL the highest accuracy (80.45, 71.81, 80.57); the paper explicitly does not claim the small PL–SCE accuracy gaps on DeiT/Swin as statistically significant.

• On ImageNet-LT the calibration advantage disappears. PL beats Mixup+CutMix in accuracy by 1.40, 1.02, 1.31, and 1.82 points across the four backbones, but SCE's ECE advantage does not persist, and Mixup+CutMix has the lowest ECE on all four.

• Longer training keeps the pattern but narrows the accuracy gap. At 320 epochs (20 warmup + 300 regular) on ViT-Wee and ViT-Little, PL still has the highest accuracy (80.64, 81.95) and SCE the lowest ECE (3.23, 2.77), but PL's margin over Mixup+CutMix shrinks to 0.91 and 0.21 points.

• Loader cost is moderate relative to RICAP. Versus single-image augmentation at 1282.8 images/s, TreemapMix-SCE (K = 4) loses 19.0% throughput and PL (K = 6) loses 25.6%, while RICAP (K = 4) loses 71.6%. Peak RSS rises by +2.47 GiB (SCE) and +4.04 GiB (PL), against +8.44 GiB for RICAP. The authors note these are CPU-side measurements that exclude model computation.

Methodology in Plain English

The method processes images in groups of K. For each group, it draws K weights from a Dirichlet distribution with concentration α — small α produces one dominant region plus small slivers, large α produces near-equal regions. Those weights are sorted ascending and handed to a squarified treemap algorithm, which carves the image canvas into K non-overlapping rectangles whose areas approximate the weights while staying as compact as possible. Each source image is resized into its rectangle. If integer pixel conversion collapses a box to zero area, the code falls back to a stripe layout, keeping the target weights unchanged.

To stop the model from learning a shortcut like "the biggest region is always class 7," the method emits K cyclic views per group: the source images are rotated through the slots so every image occupies every region size exactly once. Random area-preserving symmetries (the eight D4 transforms for square inputs; identity, 180° rotation, and flips for rectangular ones) add spatial variety without altering the targets.

Supervision then splits two ways. SCE builds a target distribution q over classes from region areas and minimizes the cross-entropy against the softmax output, whose gradient is the familiar prediction-minus-target form. PL instead sorts the present classes by area and asks that larger-area classes outrank smaller ones, normalizing each choice over the remaining classes plus the total probability mass of absent classes; each ranking term is weighted by its target area.

Training uses class-balanced rows split into groups, so complete groups contain distinct labels by construction. The main ImageNet-1K study uses four compact ViT/EVA02-family backbones from timm, two NVIDIA H20 GPUs, global batch size 1024, a 110-epoch recipe, seeds 42/43/44, K = 4 for SCE and K = 6 for PL, and a cosine α schedule from 0.1 to 0.5.

Why This Matters

For research: The paper isolates a design question that mixed-sample literature usually leaves implicit — whether region area should be treated as a probability target or as an ordering signal — and shows the two produce measurably different accuracy/calibration profiles from the identical augmented image. It also gives a concrete, reproducible recipe (Dirichlet weights plus treemap layout plus cyclic views) for getting controllable many-source mixtures without enumerating K! assignments.

Real-world applications:

• Calibration-sensitive deployment, such as medical triage or autonomous-driving perception, where a model's confidence score is used downstream and the low ECE of TreemapMix-SCE (around 1–2 percentage points on ImageNet-1K) matters more than a fraction of a point of accuracy.
• Detection and segmentation pipelines, where the COCO transfer results suggest TreemapMix-pretrained backbones are a reasonable initialization; the augmentation also resembles the multi-image mosaics already common in detection training.
• Long-tailed and skewed class distributions, since the paper evaluates on ImageNet-LT (1,000 classes, 5–1,280 images per class) and reports PL accuracy gains over Mixup+CutMix under that skew.
• Throughput-constrained training infrastructure, where the reported trade-off (19.0% / 25.6% loader slowdown versus 71.6% for RICAP) is decision-relevant for teams already bottlenecked on data loading.

Industry relevance: The augmentation is a drop-in that keeps the standard single-input model interface and the same number of training samples, so it does not require architecture changes or extra epochs of data. The published loader benchmark and the code release at github.com/Ejafa/treemapmix lower the cost of trying it.

Future Directions

• Tuning TreemapMix's own hyperparameters properly. The authors state that the best attainable TreemapMix performance under separately optimized hyperparameters is unexplored, and that short-schedule settings were reused for the 320-epoch runs.

• Separating the ingredients. The conclusion calls for isolating the contributions of cyclic views and class-balanced grouping, and for studying semantic content alongside region area — the paper's own diagnostics show RICAP beating SCE on largest-area top-class accuracy, hinting that area is a noisy proxy for evidence.

• Adaptive and end-to-end settings. Suggested extensions include adaptive grouping and schedules for K and α, calibration under skewed class priors, and measuring true end-to-end training efficiency rather than loader-side cost.

• Beyond supervised classification. The authors propose extensions to contrastive or self-supervised learning, where the composed multi-source image could supply a natural positive/negative structure.

Target Audience

Researchers and practitioners working on data augmentation, label assignment, and calibration for image classification who already know Mixup and CutMix and want a controllable many-source alternative. It is also useful for engineers deciding whether a heavier augmentation is worth its loader cost, and for anyone interested in trading exact probability matching against ordinal/ranking supervision. Readers without background in soft-label training or listwise ranking losses will find the loss derivations in Section 3.3 heavy going, though the augmentation construction in Section 3.1 is accessible on its own.

Authors’ abstract

Data augmentation improves visual recognition by exposing models to synthetic training examples that encourage generalization beyond the observed data. Mixed-sample methods extend this idea by combining multiple images. However, many common approaches either mix only pairs of images or rely on fixed small layouts, limiting control over source count, visible area proportions, and induced supervision. We propose TreemapMix, a distribution-controlled multi-image augmentation that samples source weights from a Dirichlet distribution and assigns them to area-proportional regions using a treemap partition. The resulting known region areas support two complementary objectives. TreemapMix with soft cross-entropy (SCE) treats areas as soft class-probability targets, while TreemapMix with a Plackett-Luce (PL) objective converts areas into ordinal supervision. On ImageNet-1K under matched training recipes, TreemapMix-SCE achieves substantially lower calibration error than evaluated mixed-sample baselines, while TreemapMix-PL achieves the strongest Top-1 accuracy. Transfer experiments on object detection and instance segmentation suggest that TreemapMix-pretrained features remain useful beyond classification. These results show that controlled multi-image composition can expose the same region-area information as either probability-matching or rank-matching supervision. Code is available at https://github.com/Ejafa/treemapmix.