Composite framework API

Component interfaces

interfaces

Pluggable component interfaces for the composite WSI/VLM model.

Four building blocks compose into one model:

backbone -> SELECTORS (pipeline) -> PROMPTS (bank) -> AGGREGATORS (fused)

Selectors stack sequentially. Prompts fuse into a bank. Aggregators run in parallel and fuse via one of two modes:

• "logit_ensemble" (Level 1): each aggregator runs end-to-end and produces (C,) logits; logits combine at the very end. Most decoupled, most robust if any single aggregator is buggy.
• "vector_fusion" (Level 2): each aggregator returns a slide vector (D,); vectors are fused into one (D,) and a shared classifier head produces (C,). Tighter coupling, often higher capacity.

The user picks the mode in YAML.

PromptBank dataclass

PromptBank(text_features: Tensor, aux: Dict[str, Tensor])

Output of the prompt block.

text_features is the canonical [classes, dimension] tensor every aggregator can consume. aux is a mapping of optional tensors that specific aggregators may want (e.g. SLIP needs tissue prompts, MAPLE-graph optionally provides per-entity features).

PatchSelector

Bases: Module, ABC

Filters / re-orders patches before aggregation.

Implementations are functions (patches, text_features, coords) -> patches', where patches' is a (possibly smaller) subset of the input. Selectors stack: the output of one feeds the next.

PromptModule

Bases: Module, ABC

Produces text features.

Each prompt module returns a (C, D) tensor in text_features. Auxiliary tensors that don't fit that shape (entity attributes, tissue prompts, chain-of-diagnosis hierarchy) go in aux.

Many prompt modules are themselves composed of CoOp-style learnable context vectors; common code lives in common.models.coop.

Aggregator

Bases: Module, ABC

Pools patch features into a slide-level prediction.

Each aggregator implements BOTH return modes:

• forward_vector: returns (D,) slide vector (no classification)
• forward_logits: returns (C,) class logits (with internal head)

The composite model picks whichever mode matches the configured fusion strategy. Subclasses can implement only one of the two methods and inherit a default for the other (vector -> linear head; logits -> not vector-recoverable, raises).

forward_vector

forward_vector(patches: Tensor, bank: PromptBank) -> Tensor

Return a pre-classification slide vector when supported.

Raises:

Type	Description
NotImplementedError	If the aggregator exposes only logits.

Source code in common/composite/interfaces.py

def forward_vector(self,
                   patches: torch.Tensor,
                   bank: PromptBank
                   ) -> torch.Tensor:
    """Return a pre-classification slide vector when supported.

    Raises:
        NotImplementedError: If the aggregator exposes only logits.
    """
    raise NotImplementedError(
        f"{self.__class__.__name__} does not expose a slide-vector "
        f"output. Use fusion mode 'logit_ensemble' instead.")

Recipe

Bases: ABC

Bundles optimizer + scheduler + epoch count for one paper's recipe.

build_scheduler

build_scheduler(optimizer)

Build an optional scheduler; the default recipe uses none.

Source code in common/composite/interfaces.py

def build_scheduler(self, optimizer):
    """Build an optional scheduler; the default recipe uses none."""
    return None

Model

CompositeModel

CompositeModel(
    cfg: Dict[str, Any],
    backbone: Module,
    tokenizer: Any,
    info: Any,
    encoder_bundle: Optional[EncoderBundle] = None,
)

Bases: Module

Compose patch selectors, prompt modules, and slide aggregators.

Parameters:

Name	Type	Description	Default
cfg	Dict[str, Any]	Composite configuration containing class metadata and component registries.	required
backbone	Module	Native vision-language backbone used by prompt modules.	required
tokenizer	Any	Tokenizer paired with backbone.	required
info	Any	Legacy backbone dimension metadata.	required
encoder_bundle	Optional[EncoderBundle]	Validated capability-aware bundle for components that use black-box encoder operations.	None

The model accepts one slide at a time as [patches, dim] or [1, patches, dim] and returns class logits plus optional auxiliary loss inputs.

Source code in common/composite/model.py

def __init__(self, cfg: Dict[str, Any], backbone: nn.Module,
             tokenizer: Any, info: Any,
             encoder_bundle: Optional[EncoderBundle] = None):
    super().__init__()
    self.cfg = cfg
    self.backbone = backbone
    self.tokenizer = tokenizer
    self.info = info
    self.encoder_bundle = encoder_bundle
    self.n_classes = cfg["n_classes"]
    self.classnames = cfg["classnames"]

    if (encoder_bundle is not None and encoder_bundle.spec.shared_dim is not None
            and cfg.get("feature_dim", encoder_bundle.spec.shared_dim)
            != encoder_bundle.spec.shared_dim):
        raise ValueError(
            "Composite selectors compare patch and text vectors directly: "
            f"feature_dim={cfg.get('feature_dim')} but "
            f"{encoder_bundle.spec.name} shared_dim={encoder_bundle.spec.shared_dim}. "
            "Provide aligned features; no implicit projection is inserted.")

    # ---- 1. Selectors (stack) -------------------------------------
    self.selectors = nn.ModuleList(self._build_selectors(cfg))

    # ---- 2. Prompts (bank) ----------------------------------------
    self.prompt_fusion = self._build_prompt_bank(cfg, backbone, tokenizer)

    # ---- 3. Aggregators (fused) -----------------------------------
    in_dim = cfg.get("feature_dim", info.patch_dim)
    agg_cfg = cfg["aggregators"]
    self.fusion_mode = agg_cfg.get("fusion", "logit_ensemble")
    aggs = self._build_aggregators(agg_cfg, in_dim)
    if self.fusion_mode == "logit_ensemble":
        self.aggregator_block = LogitEnsemble(
            aggs, mode=agg_cfg.get("logit_mode", "mean"),
            n_classes=self.n_classes)
    elif self.fusion_mode == "vector_fusion":
        self.aggregator_block = VectorFusion(
            aggs, mode=agg_cfg.get("vector_mode", "concat"),
            n_classes=self.n_classes)
    else:
        raise KeyError(
            f"Unknown aggregator fusion '{self.fusion_mode}'. "
            f"Use 'logit_ensemble' or 'vector_fusion'.")

forward

forward(
    patches: Tensor,
    coords: Tensor = None,
    return_extras: bool = False,
) -> Tuple[Tensor, Dict[str, Any]]

Run the configured selector, prompt, and aggregation pipeline.

Parameters:

Name	Type	Description	Default
patches	Tensor	One variable-length patch bag.	required
coords	Tensor	Optional patch coordinates aligned with patches.	None
return_extras	bool	Include tensors used by optional composite losses.	False

Returns:

Type	Description
Tuple[Tensor, Dict[str, Any]]	A pair containing class logits and an auxiliary tensor mapping.

Source code in common/composite/model.py

def forward(self, patches: torch.Tensor,
            coords: torch.Tensor = None,
            return_extras: bool = False
            ) -> Tuple[torch.Tensor, Dict[str, Any]]:
    """Run the configured selector, prompt, and aggregation pipeline.

    Args:
        patches: One variable-length patch bag.
        coords: Optional patch coordinates aligned with ``patches``.
        return_extras: Include tensors used by optional composite losses.

    Returns:
        A pair containing class logits and an auxiliary tensor mapping.
    """
    # Strip batch dim if present (every method assumes one slide)
    if patches.dim() == 3:
        patches = patches.squeeze(0)
    if coords is not None and coords.dim() == 3:
        coords = coords.squeeze(0)

    # 1. prompt bank (don't depend on patches)
    bank: PromptBank = self.prompt_fusion()

    # 2. selectors (text-aware)
    patches = self._apply_selectors(patches, bank.text_features, coords)

    # 3. fused aggregators -> logits
    logits = self.aggregator_block(patches, bank)

    if not return_extras:
        return logits, {}

    # Extras for composite loss: SLIP cross-corr if available, MAPLE
    # attributes if available.
    extras: Dict[str, Any] = {}
    slip = next((a for a in self._iter_aggregators()
                 if a.__class__.__name__ == "SLIPRoutingAggregator"), None)
    if slip is not None and "tissue" in bank.aux:
        # cheap recomputation of cross_corr for the contrastive loss
        extras["slip_cross_corr"] = self._slip_cross_corr(slip, patches, bank)
        extras["slip_temperature"] = slip.temperature

    if "attributes" in bank.aux and "attribute_class_index" in bank.aux:
        # use the last vector aggregator's slide vector if available
        sv = self._slide_vector_for_aux(patches, bank)
        extras["maple_attributes"] = (
            sv, bank.aux["attributes"], bank.aux["attribute_class_index"])
    return logits, extras