TaxoMIL: Taxonomy-Constrained Learning for Hierarchical Whole Slide Image Analysis¶

Conference: ECCV 2026
arXiv: 2606.31100
Code: https://github.com/QuIIL/TaxoMIL
Area: Medical Imaging
Keywords: Whole Slide Image Analysis, Multiple Instance Learning, Hierarchical Classification, Text Generation, Computational Pathology

TL;DR¶

TaxoMIL reformulates whole slide pathological image diagnosis as a hierarchical text generation task constrained by clinical taxonomies. Using a dual-head decoder, it simultaneously generates coarse- and fine-grained diagnostic texts. It enforces alignment of visual and label representation spaces with the medical taxonomy tree via hierarchical and image-text alignment losses, comprehensively outperforming existing methods on three WSI datasets: gastric, breast, and prostate.

Background & Motivation¶

Whole slide image (WSI) analysis is a core task in computational pathology. However, due to the gigapixel scale of WSIs, end-to-end optimization is computationally infeasible. Consequently, Multiple Instance Learning (MIL) has become the standard pipeline: splitting slides into local patches, encoding them into feature embeddings, and aggregating them into a global slide-level representation for classification prediction. However, traditional MIL methods invariably model diagnosis as flat, mutually exclusive classification—treating individual labels as independent and unrelated index values. This is severely decoupled from the inherent hierarchical structure of clinical diagnosis.

Pathological diagnosis is naturally a hierarchical, coarse-to-fine cognitive process. Pathologists typically determine broad categories first (e.g., "malignant") before refining them into specific subtypes or grades (e.g., "ductal carcinoma in situ" or "invasive carcinoma"). When labels are treated as mutually exclusive classification indices, standard optimization objectives do not spontaneously maintain coarse-to-fine consistency, nor can they encode semantic proximity between close diagnoses. Recent works have attempted to fill this structural gap but have followed two flawed paths: hierarchical classifiers (such as HMIL) introduce classification trees through multi-branch architectures, yet they still rely on discrete indices and fail to capture the semantic continuity between diagnostic labels; VLM methods (such as WSI-VQA, SlideChat, CAMP) can output clinically readable natural language text, but they lack explicit taxonomic constraints, rendering them prone to "taxonomic hallucinations" in generation—where contradictions between parent and child labels frequently occur. Neither paradigm can simultaneously achieve structured reasoning, taxonomy tree fidelity, and interpretable text generation.

The key insight of this paper is that since WSI diagnosis fundamentally performs "coarse category first, subtype second" structured reasoning on an image, prior knowledge of the clinical taxonomy tree should not be discarded. Core Idea: Redefine WSI diagnosis as a hierarchical text generation task constrained by clinical taxonomies—using a dual-head decoder to simultaneously output both coarse- and fine-grained diagnostic text, while imposing explicit hierarchical alignment and sibling margin constraints on the label embedding space to anchor visual representations into a semantic space conforming to the taxonomy tree.

Method¶

Overall Architecture¶

TaxoMIL consists of three primary components: a standard MIL backbone (feature extractor \(\mathcal{E}\) + MIL aggregator), a dual-branch multimodal conditional encoding module, and a dual-head text decoder. The input WSI is encoded into patch embeddings using the UNI foundation model, and the MIL aggregator outputs a global slide-level representation \(\mathbf{z}\). This representation is split along the feature dimension into coarse and fine halves, which are projected into the text decoder embedding space via linear adapters. Each branch further concatenates the visual embedding, a learnable image token, a label context token, and a text prompt to construct the decoding prefix. Both prefixes are fed into a shared GPT-2 backbone paired with two independent generation heads, autoregressively decoding the coarse- and fine-grained diagnostic text, respectively. Four loss functions are jointly optimized, accompanied by a cyclic loss scheduler to prevent training instability among multiple objectives.

flowchart TD
    A["Input WSI"] --> B["UNI Encoder<br/>Extract Patch Features"]
    B --> C["MIL Aggregator<br/>(Transformer/Attention)"]
    C --> D["Slide-level representation z"]
    D --> E["Split along feature dim"]
    E --> F1["Coarse-grained z_c<br/>→ Adapter g_c"]
    E --> F2["Fine-grained z_f<br/>→ Adapter g_f"]
    F1 --> G1["Construct coarse prefix<br/>v_c + &lt;img&gt; + ctx_c + Prompt"]
    F2 --> G2["Construct fine prefix<br/>v_f + &lt;img&gt; + ctx_f + Prompt"]
    G1 --> H["Shared GPT-2 Backbone"]
    G2 --> H
    H --> I1["Coarse Generation Head<br/>→ Coarse Diagnosis"]
    H --> I2["Fine Generation Head<br/>→ Fine Diagnosis"]
    I1 & I2 --> J["Joint Optimization of 4 Losses<br/>GEN + HIER + ITAL + CL"]

Key Designs¶

1. Dual-branch Multimodal Conditional Encoding: Constructing Decoder Prefixes for Different Granularities

The slide-level representation \(\mathbf{z}\) output by the MIL aggregator is a single vector. How can it "simultaneously carry coarse- and fine-grained information" to the two decoding heads? TaxoMIL's approach is to split \(\mathbf{z}\) along the feature dimension into two equal-length sub-vectors \(\mathbf{z}_c\) and \(\mathbf{z}_f\), representing the embeddings of the coarse and fine granularity perspectives, respectively. Each sub-vector is mapped to the text decoder embedding space via a lightweight adapter (linear projection + LayerNorm + GELU activation), yielding visual conditional embeddings \(\mathbf{v}^{(c)}\) and \(\mathbf{v}^{(f)}\).

However, visual conditions alone are insufficient to provide the decoder with "decision granularity" information. TaxoMIL constructs a decoding prefix for each branch by concatenating four parts: the visual conditional embedding \(\mathbf{v}\), a learnable image token <|img|> (acting as a fixed visual reference anchor), a label context token \(\mathbf{e}_{\text{ctx}}\) (computed as the mean of all label text embeddings at that granularity, providing a global semantic prior), and a natural language prompt template (e.g., "The condition shown in this gastroscopy image is:"). The design of the label context token is particularly ingenious: the coarse version \(\mathbf{e}_{\text{ctx}}^{(c)}\) is the average of all broad category name embeddings, while the fine version \(\mathbf{e}_{\text{ctx}}^{(f)}\) is the average of all subtype name embeddings. This allows the decoder to be aware of "which granularity is currently being decided" from the very first token. During training, the prompt template is randomly sampled from 10 semantically equivalent templates with different syntactic structures to enhance robustness, while a fixed template is used during inference.

2. Dual-Head Shared Decoder: One GPT-2 Backbone with Two Independent Generation Heads

To avoid training two independent language models for the two granularities (which would double the parameters and fail to share the coarse-to-fine reasoning knowledge), TaxoMIL employs a "shared backbone + independent generation heads" architecture. The decoding prefixes of both branches are sequentially fed into a single, small GPT-2 model (12 layers, 124M parameters). After obtaining the corresponding hidden state sequences for each branch, they are each connected to a lightweight Transformer decoder block and independent unembedding projection matrices, outputting vocabulary logits to autoregressively decode the complete diagnostic text.

The shared backbone implies that coarse-grained and fine-grained generation share the same language understanding capabilities—knowledge from the coarse-grained "broad category determination" can naturally assist in fine-grained "subtype differentiation," and vice versa. This aligns with the reasoning patterns of pathologists: knowing that "this is benign" helps narrow down the candidates for "which specific type of benign polyp." Meanwhile, the independent generation heads ensure flexibility in the output space's granularity (up to 5 tokens for coarse-grained, and up to 15 tokens for fine-grained). Greedy decoding is adopted during inference, requiring approximately 114ms per WSI, which is moderate among VLM-based methods.

3. Taxonomy-Guided Label-Visual Alignment: Four Losses with Cyclic Scheduling

Generation loss alone is insufficient to make visual representations inherently conform to the taxonomy tree structure. TaxoMIL imposes constraints from two levels: in the label embedding space, a hierarchical alignment loss (\(\mathcal{L}_{HIER}\)) pulls each coarse label embedding closer to its child fine label, pushes it away from other coarse labels, and separates fine labels under the same parent node; in the cross-modal space, an image-text alignment loss (\(\mathcal{L}_{ITAL}\)) pulls visual embeddings toward the semantic anchors of their corresponding label text embeddings, prompting the visual representation to inherit the hierarchical structure of the label space.

A fundamental contradiction exists when multiple logistics coexist: \(\mathcal{L}_{HIER}\) and contrastive loss \(\mathcal{L}_{CL}\) pull the structure within modalities, whereas \(\mathcal{L}_{ITAL}\) aligns across modalities; simultaneous optimization would cause objective conflicts and training instability. TaxoMIL proposes a Cyclic Loss Scheduler to resolve this contradiction: the weights of \(\mathcal{L}_{HIER}\) and \(\mathcal{L}_{CL}\) oscillate according to \(\cos^2(\pi \cdot n_{\text{cycles}} \cdot u)\) (where \(u\) is the training progress, and \(n_{\text{cycles}}=3\)), while the weight of \(\mathcal{L}_{ITAL}\) oscillates in anti-phase according to \(1-\cos^2(\pi \cdot n_{\text{cycles}} \cdot u)\). This establishes an alternating "pull structure \(\rightarrow\) align modalities \(\rightarrow\) pull structure" rhythm, avoiding gradient interference when optimizing all objectives simultaneously.

Loss & Training¶

The total loss is the sum of four terms:

\[ \mathcal{L} = w_{gen}\mathcal{L}_{GEN} + w_{hier}\mathcal{L}_{HIER} + w_{ital}\mathcal{L}_{ITAL} + w_{cl}\mathcal{L}_{CL} \]

\(\mathcal{L}_{GEN}\): Token-level cross-entropy for both coarse and fine branches, serving as the main loss (\(w_{gen}=1.0\) fixed).
\(\mathcal{L}_{HIER}\): Composed of hierarchical alignment loss \(\mathcal{L}_{HAL}\) and sibling margin loss \(\mathcal{L}_{SML}\). \(\mathcal{L}_{HAL}\) constructs a contrastive ratio of same-class similarity (coarse-to-child-fine) over different-class similarity (coarse-to-coarse, sibling-fine-to-fine), forcing the label embedding space to conform to the parent-child topology of the taxonomy tree. \(\mathcal{L}_{SML}\) introduces a hard Euclidean distance margin \(m=1.5\) on top of \(\mathcal{L}_{HAL}\) to prevent fine labels from collapsing completely toward coarse labels under coarse-fine alignment pressure, thereby preserving discrimination capability.
\(\mathcal{L}_{ITAL}\): Cross-modal contrastive learning that pulls visual embeddings \(\mathbf{v}^{(c)}/\mathbf{v}^{(f)}\) toward their corresponding label text embeddings, enabling visual representations to inherit the taxonomy structure of the label space.
\(\mathcal{L}_{CL}\): Standard supervised contrastive loss between visual samples of the same granularity, with temperature \(\tau=0.07\).
In the cyclic scheduler, \(w_{hier}^{\max}=0.3\), \(w_{ital}^{\max}=0.3\), \(w_{cl}^{\max}=0.1\). The model is trained for a maximum of 100 epochs (including early stopping) with a batch size of 64, using the AdamW optimizer (weight decay 0.01) and a cosine annealing learning rate scheduler.

Key Experimental Results¶

Main Results¶

Evaluated on three WSI datasets: gastric (GastWSI, 7,228 slides, 4 coarse and 23 fine classes), breast (BRACS, 545 slides, 3 coarse and 7 fine classes), and prostate (PANDA, 10,614 slides, 3 coarse and 6 fine classes). Compared against 12 baselines (5 single-label MILs, 4 hierarchical MILs, 3 VLMs) across three evaluation granularities: Holistic (both coarse and fine must be correct), Coarse-level, and Fine-level.

Holistic Evaluation (Table 1):

Dataset	Metric	TaxoMIL	Best Baseline	Gain
GastWSI	ACC	64.55	63.26 (CAMP)	+1.29
GastWSI	W-F1	0.6183	0.6044 (CAMP)	+0.0139
BRACS	ACC	75.78	68.52 (ABMIL)	+7.26
BRACS	W-F1	0.7662	0.7024 (ABMIL)	+0.0638
PANDA	ACC	54.60	50.28 (WSI-VQA)	+4.32
PANDA	W-F1	0.5347	0.4928 (WSI-VQA)	+0.0419

TaxoMIL achieves the best Holistic performance across all three datasets, particularly outperforming the strongest baseline by over 7 percentage points on BRACS. None of the baselines managed to maintain competitive performance across all three datasets simultaneously.

TaxoMIL also leads comprehensively at the Coarse/Fine levels (Table 2). The magnitude of the improvement at the fine-grained level (BRACS +7.41% to 25.93% ACC) is significantly larger than at the coarse-grained level (+0.93% to 5.12%), demonstrating the core value of taxonomic constraints for fine-grained discrimination—where the boundaries between classes are fuzzier, the constraint of prior knowledge from the taxonomy tree becomes more pronounced.

Ablation Study¶

Training Strategy Ablation (Table 3, Holistic ACC):

Configuration	GastWSI	BRACS	PANDA
Only \(\mathcal{L}_{GEN}\)	61.58	66.67	49.53
+ \(\mathcal{L}_{CL}\)	62.60	62.96	51.97
+ \(\mathcal{L}_{HIER}+\mathcal{L}_{ITAL}\)	63.07	72.22	51.41
+ Cyclic (Full Model)	64.55	75.78	54.60

Adding \(\mathcal{L}_{CL}\) alone on BRACS actually decreased the Holistic ACC from 66.67 to 62.96, indicating that contrastive learning without taxonomic constraints may lead to over-clustering of broad categories, thereby hindering fine-grained distinction. HIER+ITAL brings the largest magnitude of improvement (approx. +10% fine-grained on BRACS), and cyclic scheduling further stabilizes training.

Decoding Prefix Ablation (Table 4, Holistic ACC): Removing the label context token \(\mathbf{e}_{\text{ctx}}\) decreases ACC by 2.04% to 5.41%; removing the image token <|img|> decreases it by 1.50% to 5.41%. Both prefix components have distinct positive contributions across different datasets, proving both are indispensable.

Key Findings¶

The taxonomic constraint loss (HIER+ITAL) is the most significant contributor; using contrastive learning alone can be detrimental on certain datasets, indicating that intra-modality contrast must be coordinated with taxonomic constraints.
The cyclic scheduler significantly improves final performance by alternating between "pulling structure" and "aligning modalities," representing an effective, lightweight solution for multi-objective optimization.
Generative decoding comprehensively outshines classification head variants (TaxoMIL-Cls), showing that the semantic continuity of text forms benefits hierarchical differentiation.
The Parent-Child Violation Rate (PCVR) averages 4.09%, maintaining highly competitive hierarchical consistency while achieving peak ACC.
Embedding space visualization (MDS projection) shows that after training, both label and image embeddings exhibit a clear taxonomy tree topology.

Highlights & Insights¶

Directly Encoding the Taxonomy Tree Into the Loss Function: Prior hierarchical MIL methods either relied on multi-branch architectures with hard constraints (discrete index issues) or relied on the textual flexibility of VLMs (hallucination issues). TaxoMIL imposes explicit hierarchical alignment + sibling margin constraints in the label embedding space, which is the most direct and effective way to inject "taxonomy tree knowledge" into learning.
Lightweight and Elegant Cyclic Loss Scheduling: The conflict between intra-modality structuring (HIER+CL) and cross-modal alignment (ITAL) when optimized simultaneously is resolved via a simple cosine-squared oscillation, eliminating the need for complex gradient adjustment or multi-stage training pipelines.
Highly Efficient Label Context Token Design: Packaging the mean of all label text embeddings as part of the decoding prefix adds negligible overhead (only one embedding vector) but allows the decoder to understand "what the entire diagnostic space looks like" from the very first token, serving as a brilliant design choice.
Dual-Head Shared Decoder Aligns with Pathological Cognitive Models: Coarse-grained category judgment naturally assists fine-grained subtype distinction, while reducing parameters to almost half of those needed for independent dual decoders.
The structural design is generalizable beyond pathology—any medical diagnosis possessing hierarchical categorization (e.g., multi-level radiology reports, dermatoscopic typing) can adopt it directly.

Limitations & Future Work¶

TaxoMIL assumes a fixed taxonomy tree (composed of predefined label text embeddings and parent-child relations), making it unable to handle dynamically evolving classification systems. Future work needs to introduce online taxonomy tree construction or adaptive extension mechanisms.
The Parent-Child Violation Rate (PCVR), though lower than most baselines, is not completely eliminated (averaging 4.09%). This requires further reduction for high-safety clinical scenarios, potentially by introducing explicit contradiction detection post-processing.
The autoregressive decoder (small GPT-2 model) incurs additional overhead (218M trainable parameters, 114ms/WSI, 897MB VRAM), which is orders of magnitude higher than standard MIL aggregators and may be sensitive to resource-constrained deployments.
There is still substantial room for improvement in fine-grained ACC on the PANDA dataset (55.07%), partly because ISUP grading itself has fuzzy boundaries between classes and all ambiguous cases were retained; finer feature representation is needed to tackle low-discrimination scenarios.
The authors noted they did not fully explore deeper taxonomy trees (currently only two layers); the feasibility of deeper hierarchies (e.g., subtype \(\rightarrow\) molecular marker) remains to be validated.

vs. Traditional MIL (ABMIL/CLAM/TransMIL/S4MIL/MambaMIL): Traditional MILs are flat, mutually exclusive classification frameworks that cannot encode hierarchical relationships or semantic continuity between diagnostic labels. By switching to hierarchical text generation, TaxoMIL significantly leads by 10 to 26 percentage points in fine-grained accuracy.
vs. Hierarchical MIL (HMIL/Chang et al./ViLa-MIL/HiClass): These methods retain multi-branch classification head architectures with discrete indices, lacking semantic continuity. TaxoMIL naturally encodes semantic similarity using textual label embeddings, while maintaining stronger taxonomic constraints.
vs. VLM methods (SlideChat/WSI-VQA/CAMP): VLM methods output flexible natural language but lack explicit taxonomy tree constraints, rendering them prone to parent-child contradictions. TaxoMIL's hierarchical alignment loss fundamentally anchors the topology of the generation space, yielding a substantial lead in Holistic evaluation.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ Reformulates WSI diagnosis from flat classification to taxonomy-constrained hierarchical text generation, and uses multiple hierarchical constraints with cyclic scheduling to ensure hierarchical consistency. The approach is clean and effective.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Features three datasets, 12 baselines, three evaluation granularities, systematic ablations, embedding space visualization, and PCVR analysis. The experiments are highly comprehensive and robust.
Writing Quality: ⭐⭐⭐⭐ The methodology is described with ample detail and clear logical structure. However, the ablation tables are somewhat dense, and some design choices (e.g., why 3 cycles were chosen, and why label context uses a simple average instead of attention weighting) deserve deeper discussion.
Value: ⭐⭐⭐⭐⭐ Directly addresses the fundamental structural deficiency in computational pathology when moving from flat classification to hierarchical diagnosis. The framework possesses great generalizability, can be transferred to any hierarchical classification task, and holds high practical value.