Logo image
Domain-Informed Attention Pooling: Abundance-Aware Set Transformers for Robust Representation Learning
Abstract

Domain-Informed Attention Pooling: Abundance-Aware Set Transformers for Robust Representation Learning

Hyunwoo Yoo, Mohammad Saleh Refahi and Gail L. Rosen
Data Science: Foundations and Applications, pp 534-545
2027
Featured in Collection :   Drexel's Newest Publications

Abstract

Abundance-Aware Learning Domain-Informed Attention Microbiome Representation Robust Representation Learning Set Transformer
Biological and environmental datasets often consist of variable length sets of elements whose quantitative importance is domain-specific, such as microbial sequence abundance. Conventional aggregation strategies like mean pooling or unweighted attention fail to reflect these priors, resulting in unstable or biologically diluted representations. We propose Domain-Informed Attention Pooling, an abundance-aware Set Transformer that integrates quantitative priors into attention weights without modifying the core Transformer architecture. Our method preserves permutation invariance while emphasizing elements proportional to their biological relevance. We evaluate our method on three real-world microbiome classification tasks, namely tumor tissue prediction, pathogenic co-occurrence, and cross-study soil detection. Across these tasks, our method demonstrates consistent performance gains over mean pooling and unweighted Set Transformers, including perfect accuracy in subtle co-occurrence prediction and improved robustness under distributional shift. This study highlights how domain-informed attention mechanisms can generalize beyond microbiomics, offering a scalable framework that exposes attention scores amenable to interpretation for robust set representation learning across heterogeneous data domains.

Metrics

1 Record Views

Details

Logo image