PaperlayerView exact revision on arXiv

Source-linked AI summary

TAP-Path: Task-Adaptive Structural and Token Pruning for Efficient and Trustworthy Pathology Foundation Models

Mehedi Hasan, Ashfak Yeafi, Md Khairul Islam

arXiv:2609.04071v1cs.CVcs.AI

TL;DR

Large pathology foundation models offer transferable representations but impose substantial inference costs and may lose useful evidence or reliability under aggressive compression. TAP-Path restructures a pretrained Virchow2 encoder through task-adaptive block and token pruning, recovering multi-depth features with a gated head. The resulting model reduces computational cost while retaining competitive internal performance, probabilistic behavior, and frozen external evaluation results.

  • Problem

    Large pathology foundation models provide transferable representations but use hundreds of millions of parameters and high inference cost, while aggressive pruning can remove sparse diagnostically useful evidence.

  • Method

    TAP-Path physically removes task-irrelevant Virchow2 blocks, adaptively prunes patch tokens, recovers multi-depth features, and integrates them with a learned gate.

  • Results

    TAP-Path reduced encoder parameters by 24.96% and analytical FLOPs by 35.20% while achieving 87.98±0.067% internal test accuracy and 82.38±0.48% macro-F1 across three runs.

  • Takeaways & Limitations

    Task-adaptive structural and token sparsification can remove task-specific redundancy while retaining competitive discrimination and informative uncertainty in the studied pathology setting.

  • Takeaways & Limitations

    Analytical FLOPs are an architecture-level estimate rather than a latency guarantee, and external evaluation currently covers only CCRCC and UCEC.

Abstract

from arXiv · show

Pathology foundation models improve transferable representation learning for histopathology, but recent gains often rely on encoders with hundreds of millions of parameters and high inference cost. We propose TAP-Path, a task-adaptive compression framework that directly restructures a pretrained Virchow2 encoder rather than distilling it into a separate student. TAP-Path combines validation-driven transformer-block selection, physical removal of redundant blocks, input-adaptive patch-token pruning, multi-depth feature recovery, and a lightweight gated task head. The final model retains 24 of 32 transformer blocks and 70% of patch tokens after pruning, reducing encoder parameters by 24.96% (631.24M to 473.70M) and analytical encoder compute by 35.20% (340.13G to 220.40G FLOPs). Across three task-head optimization seeds, TAP-Path achieved $87.98 \pm 0.067%$ test accuracy, $81.26 \pm 0.49%$ balanced accuracy, and $82.38 \pm 0.48%$ macro-F1 on a 32-class histopathology benchmark, compared with 86.89% for full Virchow2 and 87.67% for UNI2-h. TAP-Path achieved a Brier score of $0.1800 \pm 0.0005$ and failure-detection AUROC of $0.9047 \pm 0.0060$. A validation-only rare-aware objective improved rare-class balanced accuracy in a secondary operating analysis. Frozen external evaluation on 433 CPTAC samples yielded $91.22 \pm 0.83%$ accuracy and $91.10 \pm 0.81%$ balanced accuracy. These results show that task-adaptive structural and token sparsification can improve the accuracy-efficiency trade-off of large pathology foundation models while preserving reliability under internal and external evaluation.

1 Introduction

TAP-Path addresses the efficiency and trustworthiness challenges of large pathology foundation models by restructuring Virchow2 for task-specific deployment. It combines structural and token pruning with reliability-oriented evaluation while retaining competitive predictive performance.

  • Motivation and contribution: 35.20% per-patch compute reduction compounds across whole-slide images containing thousands to tens of thousands of candidate tiles.This motivates restructuring the patch encoder rather than limiting efficiency analysis to trainable parameter counts.
  • Motivation and contribution: TAP-Path targets a validated Pareto point because pathology evidence can be spatially sparse and aggressive pruning may erase diagnostically useful regions.The objective balances compression with task utility rather than maximizing compression alone.
  • Motivation and contribution: TAP-Path physically reconstructs a task-relevant, non-contiguous Virchow2 subnetwork instead of training or distilling a separate student.The framework identifies relevant transformer depth and removes unselected blocks from the deployed encoder.
  • Method: Input-adaptive token pruning and multi-depth recovery preserve complementary representations, which a learned gate integrates for downstream prediction.The method reduces redundant patch-token computation while retaining information from multiple compressed-hierarchy stages.
  • Motivation and contribution: 35.20% lower FLOPs and 24.96% fewer encoder parameters demonstrate substantial compression while preserving competitive performance on the 32-class task.The study evaluates deployed parameter count and analytical FLOPs alongside prediction.
  • Evaluation: The evaluation covers calibration, failure awareness, selective prediction, rare-class behavior, uncertainty, and locked external CPTAC testing.The external cohorts are evaluated without adaptation or threshold tuning.

2 Related Work

Related work establishes both the broad utility of pathology foundation models and the need for task- and domain-sensitive efficiency evaluation. TAP-Path extends pruning research by combining non-contiguous block selection with adaptive token pruning in a large pretrained pathology encoder.

  • Foundation models: Pathology foundation models have expanded from general representation learning to multimodal, whole-slide, multi-institutional, and broader clinical-task settings.Examples include UNI, CONCH, Virchow, Prov-GigaPath, CHIEF, Hibou, BEPH, GPFM, and PathOrchestra.
  • Evaluation and robustness: Independent evaluations show that model rankings and preferred adaptation procedures vary with downstream tasks, data regimes, and external cohorts.This motivates evaluating compression beyond a single fixed benchmark setting.
  • Evaluation and robustness: Domain robustness matters because pathology models can encode scanner and laboratory signatures that compromise generalization.Related work emphasizes external evaluation, fairness, and transparent reporting of data sources.
  • Efficient adaptation and compression: Parameter-efficient adaptation reduces optimization memory but leaves a frozen or adapter-tuned 600M-parameter encoder executing most of its original forward pass.Prompt tuning and adapters address fine-tuning cost rather than necessarily reducing inference computation.
  • Efficient adaptation and compression: Distillation creates a smaller student, whereas TAP-Path restructures the pretrained encoder in place while preserving its representation family.The two approaches answer different efficiency questions.
  • Transformer sparsification: Pathology-specific pruning work shows that redundancy can be removed selectively, while TAP-Path operates on non-contiguous blocks and then patch tokens.The distinction is between attention-head pruning in a distilled ViT and hierarchical block-plus-token pruning in a pretrained foundation model.
  • Transformer sparsification: Because discriminative pathology regions may be spatially sparse, TAP-Path conditions token removal on feature content and validates the pruning ratio rather than assuming it.Multi-depth recovery preserves information from distinct stages of the compressed hierarchy.

3 Materials and Methods

The study uses a locked, deterministic evaluation pipeline with fixed multi-scale patch bags, balanced assessment, and comparisons across pathology foundation encoders. TAP-Path then selects and physically removes task-redundant blocks using validation-only profiling, retaining anchored early and late depth while enforcing a parameter-reduction constraint.

  • Dataset and evaluation: 25,495 image-level records spanning 32 cancer classes are partitioned into 17,769 training, 3,867 validation, and 3,859 test images.Each image uses the final 12-patch protocol.
  • Dataset and evaluation: 433 CPTAC whole-slide images from CCRCC and UCEC are evaluated with the locked model without external fitting or threshold optimization.The external cohorts cover two of the 32 internal cancer classes.
  • Preprocessing and patch selection: The deterministic 12-patch manifest combines multi-scale candidate regions, confidence, diversity, and entropy to preserve representative and difficult tissue regions.One high-confidence patch per available scale and two high-entropy patches are retained before utility-based filling.
  • Preprocessing and patch selection: Validation, internal test, and external evaluation use deterministic unaugmented patches, while training applies reflections, 90° rotations, and mild contrast or brightness perturbations.Aggressive color remapping is avoided because stain-independent appearance may carry diagnostic information.
  • TAP-Path architecture: TAP-Path retains 24 of 32 transformer blocks and 70% of tokens after pruning, using a retained block set and token-retention ratio.The method is initialized from Virchow2 and combines structural selection, physical block removal, token pruning, and multi-depth feature recovery.
  • Task-adaptive block selection: Block selection profiles normalized residual changes on development data, anchors the first and final four blocks, and restores selected middle blocks to their original order.A parameter-reduction constraint with validation tolerance avoids choosing the smallest network after accuracy collapses.

3.9 Optimization objectives

TAP-Path separates validation-only architecture and token selection from final task-head optimization and locked evaluation. Its objective combines mild class weighting, calibration, gated multi-depth fusion, and comprehensive predictive, efficiency, and reliability assessment.

  • Task-head objectives: The primary head uses mild class weighting with label smoothing ϵ_ls = 0.01 and a gate-entropy term to discourage degenerate single-tap fusion.Four hierarchy taps are recovered and adaptively fused before optimizing the 32-class head.
  • Rare-class analysis: Rare-class objectives are treated as a post-locking ablation, comparing mild weighted cross-entropy, Balanced Softmax, two logit-adjustment settings, and class-balanced focal loss.The ablation remains separate from the primary head rather than combining metrics into an artificial best-of-all-metrics model.
  • Architecture and token selection: Validation-only structural screening selects compressed candidates using validation utility and compression, without using internal test data for retained-block selection.Contiguous-depth controls and non-contiguous task-sparse candidates are evaluated with a lightweight validation probe.
  • Architecture and token selection: ρ = 0.70 is locked after validation comparison of token-retention rates 1.00, 0.85, and 0.70, yielding the TaskSparse24 configuration.The final executable compressed encoder is rebuilt only after this selection.
  • Task-head objectives: Three independent seeds train the single gated head, while temperature scaling is fitted on validation logits before locked internal-test and frozen CPTAC evaluation.The protocol keeps external data from modifying architecture, threshold, temperature, or other configuration decisions.
  • Reliability objectives: Calibration and failure awareness are assessed with ECE, NLL, Brier score, uncertainty-based failure-detection AUROC, and risk–coverage analysis.Risk–coverage supports a selective-prediction interpretation but is not presented as proof of clinical safety.

4 Results

TAP-Path improves the accuracy–efficiency trade-off of Virchow2 through structural and token compression while maintaining informative confidence estimates and transfer performance. Its results also reveal operating trade-offs for rare-class sensitivity and limits on external generalization.

  • Foundation-model comparison: 87.98 ± 0.067% test accuracy and 82.38±0.48% macro-F1 exceeded full Virchow2 and UNI2-h among evaluated single-backbone systems.Balanced accuracy reached 81.26 ± 0.49%.
  • Efficiency: 24.96% fewer encoder parameters and 35.20% lower analytical encoder FLOPs produced a 479.40M-parameter deployment with 220.40G FLOPs.The physically compressed encoder contains 473.70M parameters, while the deployed configuration includes the task head.
  • Efficiency: 87.98% accuracy at 479.40M parameters and 220.40G FLOPs compared favorably with fusion systems requiring substantially greater deployment-scale resources.The comparison frames TAP-Path as a favorable operating point within the high-performing large-foundation-model regime.
  • Reliability: 0.1800 ± 0.0005 Brier score and 0.9047 ± 0.0060 failure-detection AUROC indicated informative confidence estimates despite pruning.Brier score and failure-detection AUROC improved over full Virchow2, while UNI2-h remained marginally better in ECE and NLL.
  • Rare-class operating trade-offs: 70.55% validation rare-class balanced accuracy came from Balanced Softmax, whereas logit adjustment achieved 90.25% validation accuracy and raised rare-class balanced accuracy to 69.68%.On the locked encoder, the selected logit-adjusted head increased test rare-class balanced accuracy to 68.64 ± 2.16% while accuracy decreased to 87.13±0.68%.
  • External validation: 91.22 ± 0.83% external accuracy and 91.10 ± 0.81% balanced accuracy were obtained on 433 frozen CPTAC images across three seeds.External failure-detection AUROC was 0.8974 ± 0.0089, but the task represented only two classes and therefore does not establish universal 32-class external generalization.

5 Discussion

TAP-Path’s discussion frames physical structural and token pruning as a favorable accuracy–efficiency operating point, while emphasizing task-dependent depth, reliability, rare-class trade-offs, and the limits of current evaluation.

  • Efficiency: 35.20% per-patch FLOP reduction compounds across whole-slide patch bags, but this estimate is theoretical rather than a latency guarantee.Analytical cost falls from 340.13G to 220.40G FLOPs per encoder invocation; for 10,000 patches, this corresponds to roughly 1.20×10^15 fewer operations.
  • Structural selection: 24-block TaskSparse24 outperformed contiguous Depth24 during validation screening, indicating that task relevance is distributed non-uniformly through the pretrained hierarchy.The result supports preserving selected middle and late blocks rather than relying on naive depth truncation.
  • Token retention: 70% token retention improved screening validation accuracy and reduced compute for TaskSparse24, whereas token removal degraded TaskSparse22.Token sparsity therefore interacts with representational depth and should be selected jointly rather than treated as an independent compression setting.
  • Token retention: Token pruning remains a pathology-specific risk because aggressive compression may remove uncommon but clinically meaningful morphology, despite validation-only selection and explicit rare-class evaluation.The safeguards reduce but do not eliminate token-pruning bias.
  • Reliability: TAP-Path preserves useful probabilistic behavior, achieving the best Brier score and failure-detection AUROC among directly compared large-model configurations.Its failure-detection AUROC of approximately 0.90 supports confidence-guided selective review, while clinical thresholds require prospective workflow-specific evaluation.
  • Rare-class behavior: The rare-aware head increases rare-class balanced accuracy while shifting overall accuracy, exposing an accuracy–minority-sensitivity trade-off.The two operating points are reported separately to preserve metric provenance.
  • External evaluation: Frozen CPTAC evaluation shows high transfer performance, but internal 32-class macro-F1 cannot be directly compared with external two-class macro-F1.The external analysis therefore emphasizes present-class balanced accuracy.
  • Overall discussion: TAP-Path combines competitive predictive performance with reduced model size and computational demand in a joint accuracy–efficiency–reliability evaluation.The evaluation also includes three-seed reproducibility, calibration, failure-awareness analysis, and external performance.

6 Limitations

The paper’s main limitations concern deployment realism, external-evaluation breadth, structural-selection validation, and the statistical rather than clinical nature of reliability analysis.

  • Deployment scope: Analytical FLOPs and 31.85 ms image-level latency do not fully characterize deployment cost, which also depends on hardware, energy, throughput, loading, extraction, and storage I/O.The latency figure describes the evaluated hardware configuration rather than a general systems-level guarantee.
  • External scope: External evaluation currently covers only CCRCC and UCEC, limiting broader pan-cancer robustness analysis across cancer types, scanners, and institutions.Additional cohorts would be needed to extend the transfer assessment.
  • Structural selection: The block-novelty score is a computationally efficient task-selection criterion, but future work should compare it with gradient-, Hessian-, and attribution-based importance measures.The transferability of retained-block patterns across endpoints also remains to be examined.
  • Clinical validation: Reliability evidence is currently statistical and model-centered, leaving prospective clinical evaluation, reader studies, and workflow-specific thresholds for future work.These boundaries limit direct conclusions about clinical operation.

7 Conclusion

TAP-Path converts Virchow2 into a task-specific encoder through validation-constrained structural and token sparsification without knowledge distillation. The resulting model reduces computational demands while retaining predictive and probabilistic performance, supporting compression as complementary to foundation-model scaling.

  • Method: TAP-Path retains 24 of 32 transformer blocks and 70% of patch tokens, recovering multi-depth information through a gated single task head.The method restructures the pretrained model directly rather than training a separate distilled student.
  • Efficiency: 24.96% fewer parameters and 35.20% fewer analytical FLOPs result from the compressed Virchow2 encoder.The reported reductions are 631.24M to 473.70M parameters and 340.13G to 220.40G FLOPs.
  • Results: 87.98±0.067% internal test accuracy and 82.38±0.48% macro-F1 were achieved across three independent single-head runs, matching or slightly exceeding larger baselines on the studied task.The conclusion also reports a Brier score of 0.1800, failure-detection AUROC of 0.9047, and frozen CPTAC accuracy of 91.22±0.83%.
  • Takeaway: The results support removing downstream-task redundancy from large pretrained hierarchies through validation-constrained structural and token sparsification.The conclusion presents compression and foundation-model scaling as compatible rather than mutually exclusive.

Data Availability

The study uses publicly available TCGA and CPTAC histopathology resources, with no external evaluation data used for architecture selection or model optimization.

  • Data sources: TCGA data are accessible through the NCI Genomic Data Commons, while CPTAC resources are available through NCI-supported data repositories.The analyzed histopathology data derive from publicly available resources.
  • Data use: No external evaluation data were used for architecture selection or model optimization.

A.2 Experimental provenance

The manuscript keeps architecture screening, task evaluation, diagnostics, rare-aware analysis, and external validation as distinct experimental protocols. This prevents metrics from being transferred across protocols and distinguishes representation-quality analysis from the deployed TAP-Path task head.

  • The experiments separate architecture screening, primary task evaluation, common-probe diagnostics, rare-aware operating-point analysis, and external validation.
  • Metrics are not transferred between these protocols.
  • Common-probe analysis evaluates representation quality, whereas the original TAP-Path task head defines the proposed deployed system.
Loading 2609.04071v1…