Encoder-only Mask Transformer · annotated

One query embedding,
many independent labels.

attr-eomt is a standalone EoMT detector whose distinguishing trick is small: it reads extra per-instance attributes straight off the embedding the detector already computes — for free, without bloating the class list.

imagra93/attr-eomt · Apache-2.0 · this page focuses on the ⭐ auxiliary attribute heads
The base it builds on

EoMT in one breath

Before the attributes make sense, the host model: EoMT is a DINOv2-with-registers ViT whose last few transformer blocks are augmented with a fixed set of learnable queries (the Mask2Former idea). Each query is one "slot" that learns to latch onto one object instance.

After the encoder runs, every query produces a single vector — the per-query embedding, shape [B, Q, hidden] (batch × queries × width). The whole model is essentially "turn that embedding into predictions":

It's NMS-free: two overlapping cars stay two distinct queries instead of being merged by non-max suppression. That property is what lets attributes stay attached to the right instance.

sizebackbonehiddenqueries
sDINOv2-small384100
bDINOv2-base768200
lDINOv2-large1024200

Two families share all of the above and differ only in the geometry head: instance → masks, detect → boxes. Everything about attributes applies identically to both.

Why attributes exist

The label space explodes

A normal detector has one flat list of classes. So when an object has a type and a colour and a state, the only way to express it is the Cartesian product: type × colour × state. That blows up combinatorially, starves each leaf class of examples, and multiplies the targets the Hungarian matcher has to assign.

attr-eomt factorizes instead. Keep the primary taxonomy small and general (car, apple, person); push orthogonal distinctions into independent heads that add, not multiply. Drag the sliders — this is the entire argument for the feature:

Flat label space vs. factorized heads

4
6
3
Flat — one head
72
4 × 6 × 3 leaf classes
Factorized — 3 heads
13
4 + 6 + 3 outputs

The flat space is 5.5× larger — and each of those leaves has to be learned from its own scarce examples. Factorized heads each see every instance.

Because the heads are independent, the model can predict red on a class it never saw in red during training — it composes attribute × class combinations a flat label space cannot even represent.

How it's wired

The heads tap the same embedding

The key implementation idea: the attribute heads don't get a new feature extractor. They read the exact same per-query embedding that feeds class_predictor — captured non-invasively with a PyTorch forward hook. The detector doesn't even know they're there.

image 644×644 DINOv2 ViT last blocks + learnable queries NMS-free query_embed [B, Q, hidden] ▲ forward hook on class_predictor class_predictor primary class · nc+1 mask / box head geometry aux head · "ripeness" 3 classes aux head · "grade" 2 classes … as many as the data defines teal = the one shared embedding · every head is just a small classifier on top of it
model.py — a forward hook stashes the input to class_predictor; aux heads live in an nn.ModuleDict keyed by name.

In code it's about as small as it sounds. Each head is built by _build_aux_head — a bare linear probe, or a tiny MLP if you ask for more layers:

# model.py — one head per attribute spec, sharing one arch
self.aux_heads = nn.ModuleDict({
    s.name: _build_aux_head(hidden, s.num_classes, arch)  # Linear, or small MLP
    for s in self.aux_specs
})

# non-invasive: grab the embedding feeding the class head
self.eomt.class_predictor.register_forward_hook(self._capture_query_embed)

def _capture_query_embed(self, _m, inputs, _out):
    self._query_embed = inputs[0]   # [B, Q, hidden]

Default head arch is a 2-layer MLP (Linear → LayerNorm → GELU → Linear); layers=1 makes it a pure linear probe. The arch is stored in the checkpoint, so reloading rebuilds identical modules.

Training — aux_cls.py

It rides on the detector's own matching

The clever part of training is that attributes never run their own matcher. Detection already solves "which query is responsible for which ground-truth object" via the Hungarian matcher. Attributes simply reuse that same query→GT assignment, then read the answer off the matched queries.

  1. MatchReuse model.eomt.criterion.matcher on the model's geometry (masks for instance, boxes for detect) to get (query_idx, gt_idx) pairs — the same assignment the detection loss used.
  2. Gate (optional)The matcher pairs every GT, even barely-overlapping ones (common early on). gate_indices drops pairs whose predicted↔GT IoU is below a threshold, so attributes only learn from queries that actually localize the object.
  3. Gather_gather_matched pulls just the matched rows out of query_embed → a tight [N, hidden] tensor of "real" instances.
  4. LossEach head runs on that tensor; cross-entropy against the per-instance attribute labels, summed across heads, scaled by aux_w (default 1.0), and added to the detector loss. Missing labels use ignore_index and contribute nothing.
  5. Graph-safe zeroIf no query matched in a batch, it returns embed.sum()*0 — a real zero that keeps autograd happy instead of crashing.
# aux_cls.py — the whole supervision, condensed
indices  = match_queries(model, out, geom_labels, class_labels)  # reuse Hungarian
indices  = gate_indices(out, indices, ..., iou_thr=0.5)          # keep well-localized
matched  = gather_matched(out, indices)                          # [N, hidden]
for name, head in model.aux_heads.items():
    logits = head(matched)
    loss  += w * F.cross_entropy(logits, gt[name], ignore_index=-100)

The attribute "rides along for free": it reuses the backbone and the detector's matched queries, adding only a thin head and one CE term — not a second model or a second pass.

It never hijacks model selection

Per-head matched-query accuracy is printed live and written to metrics.csv, but it never drives which checkpoint becomes best.pt — that's still segm/mAP (or bbox/mAP). The attribute genuinely tags along without changing the detector's training contract.

Data & inference

Attributes live inside the COCO annotations

Because every COCO annotation is already a per-instance object, alignment is automatic — pycocotools still parses the file untouched. Two additions define the heads; no YAML changes, the heads (count, classes, names) are discovered straight from the JSON, just like nc.

1 · a top-level attributes list (the vocabulary)

"attributes": [
  {"name": "ripeness", "categories": [
     {"id": 0, "name": "unripe"}, {"id": 1, "name": "turning"}, {"id": 2, "name": "ripe"}]},
  {"name": "grade", "categories": [
     {"id": 10, "name": "A"}, {"id": 20, "name": "B"}]}
]

2 · a per-annotation attributes map

{ "id": 1, "image_id": 42, "category_id": 1, "bbox": [...], "segmentation": [...],
  "attributes": {"ripeness": 1, "grade": 20} }

Raw ids are remapped to a contiguous 0..n-1 per head (so grade's 10/200/1). A missing value defaults to 0; a JSON with no attributes is just detection-only, exactly as before.

At inference

Each kept detection gets an aux = {head: {"ids", "probs"}} entry, and predict(plot=True) renders every attribute next to the class label, using the names stored in the checkpoint — the primary class on the first row, each attribute + its confidence on the row beneath.

Roadmap idea worth noting: train an aux head with a contrastive objective and its per-instance embedding becomes a re-identification vector — matching the same object across frames and cameras, reusing the detector's matched queries instead of a separate tracker.

The whole idea, once more

Factorize the labels, tap the embedding, reuse the match

Keep the primary taxonomy small and well-populated. For everything orthogonal — colour, state, grade, laterality — attach an independent head that reads the per-query embedding the detector already produces, supervise it on the detector's own Hungarian match, and add a cross-entropy term. You get composable, generalizing per-instance attributes for almost no extra compute, and the primary detection metric stays exactly as it was.