One query embedding,
many independent labels.

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":

• the class head (class_predictor) turns it into a primary label,
• the mask head (or a box head, in the detect family) turns it into geometry.

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

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.

The label space explodes

A normal detector has one flat list of classes. So when an object has a class and several orthogonal attributes, the only way to express it is the Cartesian product: class × attribute #1 × attribute #2 × …. 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 (class #1, class #2, …); push each orthogonal attribute into its own independent head that adds, not multiplies. Drag the sliders — this is the entire argument for the feature:

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.

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.

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)

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.

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": "attribute_1", "categories": [
     {"id": 1, "name": "value_a"}, {"id": 2, "name": "value_b"}, {"id": 3, "name": "value_c"}]},
  {"name": "attribute_2", "categories": [
     {"id": 0, "name": "value_x"}, {"id": 1, "name": "value_y"}, {"id": 2, "name": "value_z"}]}
]

2 · a per-annotation attributes map

{ "id": 1, "image_id": 42, "category_id": 1, "bbox": [...], "segmentation": [...],
  "attributes": {"attribute_1": 3, "attribute_2": 0} }

Raw ids are remapped to a contiguous 0..n-1 per head (so attribute_1's 1/2/3 → 0/1/2). 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.

The training recipe, and the knobs

Defaults follow the EoMT / Mask2Former fine-tuning recipe. Each piece is a keyword argument to train(), so legacy behaviour is one override away.

• Effective batch via gradient accumulation — LR/WD/clip are tuned for an effective batch of 16; EoMT is a ViT (LayerNorm, no BatchNorm), so accumulation ≈ a true large batch at a fraction of the memory.
• EMA weights are validated and saved as best.pt; last.pt holds live weights + optimizer + EMA state for exact resume.
• Large-Scale Jitter for training, letterbox for eval/inference (the mode is stored per-checkpoint so val/predict match training automatically).
• AdamW with no weight decay on norms/biases/embeddings, and layer-wise LR decay on the backbone.
• Masked-attention annealing 1 → 0 so the final stretch trains mask-free and matches efficient (mask-less) inference.

Compress to int8

compress() applies int8 weight-only quantization (via torchao) to the ViT transformer blocks — the bulk of the parameters — while keeping the prediction heads in full precision. It's data-free (no calibration) and shrinks the checkpoint ≈3.4× on l with no measurable mAP loss.

m = EoMT("runs/train/eomt-l")
m.compress("int8", data="coco", save="runs/train/eomt-int8/weights/best.pt")
# deltas returned + written to compression_metrics.json; omit `data` to skip validation

# reload like any run — the int8 layout is rebuilt from the checkpoint recipe
EoMT("runs/train/eomt-int8").predict("images/")

The compressed model is GPU-only (torchao's int8 kernels have no CPU path), and best.pt exported from an EMA run quantizes exactly those EMA weights.

Factorize the labels, tap the embedding, reuse the match

Keep the primary taxonomy small and well-populated. For every orthogonal axis — each an independent attribute — 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.