Aggregators

Aggregators combine a slide’s tile features into a representation for multiple-instance learning (MIL). Select a preset as the mil learner’s aggregator leaf: learner.params.aggregator.name in YAML, {"aggregator": "abmil"} (or {"name": ..., "params": ...}) in the Python LearnerConfig. Patient-level pipelines use a frozen patient encoder instead.

A frozen tile encoder turns a slide into a bag of features; a trained MIL aggregator or a frozen slide-level FM pools it into one slide-level vector for the task head.

A trained MIL aggregator pools the bag of tile features into one slide-level vector.

Aggregator Zoo

Preset

Description

Notes

mean_pool

Mean over all tile features.

max_pool

Element-wise max over tile features.

abmil

Gated attention pooling for slide-level aggregation.

Ilse et al., 2018

clam_sb

Single-branch CLAM with instance-level supervision.

Lu et al., 2021

clam_mb

Multi-branch CLAM with one attention branch per class.

Lu et al., 2021

dsmil

Dual-stream MIL with a critical-instance query.

Li et al., 2021

dtfdmil

Two-tier MIL with pseudo-bag distillation.

Zhang et al., 2022

transmil

Transformer-based MIL with Nystrom attention.

Shao et al., 2021

hipt

Hierarchical image pyramid transformer.

Chen et al., 2022

Aggregator details

ABMIL

abmil returns tile-level attention scores for heatmap generation.

class soma.models.aggregators.mil.abmil.ABMIL(input_dim, hidden_dim=128, activation='tanh', gated=True, dropout=0.25)

Bases: Aggregator

Attention-Based MIL aggregator.

Uses gated attention pooling to aggregate tile features into a slide-level representation, with per-tile attention weights for interpretability and heatmap generation.

Parameters:
  • input_dim (int) – Feature dimension of input tiles.

  • hidden_dim (int) – Attention bottleneck dimension.

  • activation (str) – Activation function (‘tanh’, ‘relu’, ‘gelu’).

  • gated (bool) – If True, use gated attention.

  • dropout (float) – Dropout rate applied before attention.

CLAM-SB

clam_sb supports binary, multiclass, ordinal, and single-target regression tasks, and can mix bag-level and instance-level supervision.

class soma.models.aggregators.mil.clam.CLAM_SB(input_dim, hidden_dim=512, attn_dim=256, gated=True, dropout=0.0, k_sample=8, n_classes=2, inst_loss='ce', use_negative_class_instance_loss=False, bag_weight=0.7, instance_loss_mode=None, low_attention_weight=0.1, topk_target_weight=1.0)

Bases: _CLAMBase

Single-branch CLAM aggregator.

CLAM-MB

clam_mb requires multiclass_classification.

class soma.models.aggregators.mil.clam.CLAM_MB(input_dim, hidden_dim=512, attn_dim=256, gated=True, dropout=0.0, k_sample=8, n_classes=2, inst_loss='ce', use_negative_class_instance_loss=False, bag_weight=0.7)

Bases: _CLAMBase

Multi-branch CLAM aggregator.

DSMIL

dsmil requires binary_classification.

class soma.models.aggregators.mil.dsmil.DSMIL(input_dim, att_dim=128, nonlinear_q=False, nonlinear_v=False, dropout=0.0)

Bases: Aggregator

Dual-Stream MIL aggregator.

Uses an instance classifier to identify a critical instance, then computes attention via query-key matching with the critical instance.

Parameters:
  • input_dim (int) – Feature dimension of input tiles.

  • att_dim (int) – Attention/query dimension.

  • nonlinear_q (bool) – If True, use nonlinear query projection.

  • nonlinear_v (bool) – If True, use nonlinear value projection.

  • dropout (float) – Dropout rate for value projection (when nonlinear_v=True).

DTFDMIL

instances_per_group (default 1) sets how many instances each pseudo-bag contributes; the partition is random during training and contiguous in eval mode.

class soma.models.aggregators.mil.dtfdmil.DTFDMIL(input_dim, embed_dim=512, hidden_dim=128, n_groups=8, distill_mode='maxmin', instances_per_group=1, dropout=0.25)

Bases: Aggregator

Double-Tier Feature Distillation MIL aggregator.

Parameters:
  • input_dim (int) – Feature dimension of input tiles.

  • embed_dim (int) – Width of the learned tile embedding every later step works in (the reference’s DimReduction to mDim = 512). Without it, pseudo-bag pooling, CAM scoring and distillation all act on frozen encoder features, and the default max-min distillation fails to learn even a task where one marker tile decides the label.

  • hidden_dim (int) – Attention bottleneck dimension.

  • n_groups (int) – Number of pseudo-bags to partition into.

  • distill_mode (str) – Feature distillation mode (‘maxmin’, ‘max’, ‘afs’).

  • instances_per_group (int) – Instances distilled per pseudo-bag and per direction (‘max’ keeps the top-k CAM instances, ‘maxmin’ the top-k and bottom-k). The reference implementation uses total_instance // numGroup = 1; clamped to the pseudo-bag size.

  • dropout (float) – Dropout rate applied before tier-1 attention.

The pseudo-bag partition is a random permutation drawn from torch’s global RNG while the module is in training mode (so it follows the run seed and differs every step, as in the reference), and a deterministic contiguous split in eval mode so scoring a checkpoint is reproducible.

TransMIL

class soma.models.aggregators.mil.transmil.TransMIL(input_dim, att_dim=512, n_layers=2, n_heads=4, n_landmarks=None, pinv_iterations=6, dropout=0.0, use_mlp=False)

Bases: Aggregator

Transformer-based Correlated MIL aggregator.

Uses Nystromformer layers with PPEG positional encoding and a learnable class token to aggregate tile features.

Parameters:
  • input_dim (int) – Feature dimension of input tiles.

  • att_dim (int) – Transformer embedding dimension.

  • n_layers (int) – Number of Nystromformer layers (must be >= 2).

  • n_heads (int) – Number of attention heads.

  • n_landmarks (int | None) – Landmarks for Nystrom approximation (default: att_dim//2).

  • pinv_iterations (int) – Pseudo-inverse iterations.

  • dropout (float) – Dropout rate.

  • use_mlp (bool) – Whether to use MLP blocks in transformer layers.

HIPT

hipt assumes hierarchical tiling; its tile_multiple param (or features.extraction.region_tile_multiple) controls how many tiles fit along each side of a region.

class soma.models.aggregators.mil.hipt.HIPT(input_dim, region_size, patch_size, embed_dim_region=192, embed_dim_slide=192, num_heads=6, dropout=0.25)

Bases: Aggregator

Hierarchical Image Pyramid Transformer aggregator (Chen et al., 2022).

Two-level hierarchy: a region-level ViT aggregates P tile features per region, then a global transformer + gated attention pools M region embeddings into a slide-level representation.

Features are stored either flat as (B, N, D) where N = M × P, or natively hierarchical as (B, M, P, D) where P = (region_size / patch_size)². HIPT reshapes internally.

Parameters:
  • input_dim (int) – Feature dimension of input tiles (auto-resolved from CachedSetSource).

  • region_size (int) – Region pixel size (e.g. 4096). Required.

  • patch_size (int) – Subtile pixel size within region (e.g. 256). Required.

  • embed_dim_region (int) – Region ViT output dimension.

  • embed_dim_slide (int) – Global transformer / output dimension.

  • num_heads (int) – Attention heads in region ViT.

  • dropout (float) – Dropout rate.

Aggregator interface

The shared base classes are soma.models.aggregators.base.Aggregator and soma.models.aggregators.base.AggregatorOutput.

class soma.models.aggregators.base.Aggregator(*args, **kwargs)

Bases: ABC, Module

Abstract base class for MIL aggregators.

Consumes a bag of tile features (B, N, D) and produces a slide-level representation (B, D_out) wrapped in an AggregatorOutput.

class soma.models.aggregators.base.AggregatorOutput(bag_representation, tile_attention=None, auxiliary=None)

Bases: object

Structured output from an aggregator.

bag_representation

Slide-level representation, shape (B, D_out).

tile_attention

Per-tile attention weights, shape (B, N). None for non-attention methods (e.g. MeanPool, MaxPool).

auxiliary

Optional dict of auxiliary tensors for training-time losses (e.g. instance logits for DSMIL, embeddings for CLAM).

soma.components("aggregators") lists the registered aggregators with their params and defaults.