Chronological age estimation from human microbiomes with transformer-based Robust Principal Component Analysis

Abstract

Deep learning for microbiome analysis has shown potential for understanding microbial communities and human phenotypes. Here, we propose an approach, Transformer-based Robust Principal Component Analysis(TRPCA), which leverages the strengths of transformer architectures and interpretability of Robust Principal Component Analysis. To investigate benefits of TRPCA over conventional machine learning models, we benchmarked performance on age prediction from three body sites(skin, oral, gut), with 16S rRNA gene amplicon(16S) and whole-genome sequencing(WGS) data. We demonstrated prediction of age from longitudinal samples and combined classification and regression tasks via multi-task learning(MTL). TRPCA improves age prediction accuracy from human microbiome samples, achieving the largest reduction in Mean Absolute Error for WGS skin (MAE: 8.03, 28% reduction) and 16S skin (MAE: 5.09, 14% reduction) samples, compared to conventional approaches. Additionally, TRPCA's MTL approach achieves an accuracy of 89% for birth country prediction across 5 countries, while improving age prediction from WGS stool samples. Notably, TRPCA uncovers a link between subject and error prediction through residual analysis for paired samples across sequencing method (16S/WGS) and body site(oral/gut). These findings highlight TRPCA's utility in improving age prediction while maintaining feature-level interpretability, and elucidating connections between individuals and microbiomes.

Fig. 1. Preprocessing and model overview to transformer-based RPCA.

Samples represented as count tables are visualized and converted to RPCA vectors. RPCA vectors are input as sequences into a transformer encoder model with multi-head attention. The transformer model outputs are provided to a classification (CLS) or regression (REG) head for classification, regression, or Multi-task learning.

Fig. 2. Regression for age prediction by sequencing method and body site.

a–c are regressions for 16S data and d–f for WGS data. a, d, b, e, c, f indicate regressions for skin, oral, and gut microbiome samples respectively.

Fig. 3. MTL performance on country of birth classification and age prediction for WGS Stool samples from curatedMetagenomicData.

a, b are the confusion matrix and age prediction regression for the TRPCA MTL model. c, d are RF models trained individually for the same classification and regression tasks as the TRPCA MTL model.

Fig. 4. Top 50 normalized and clustered feature importances for 16S Skin microbiome features.

The feature importances are derived from the dot product of the PCA and SHAP matrices with rows as samples and columns as assigned taxonomy from influential ASVs. Redundant taxonomy values are indicative of AVSs which map to the same taxonomy classification. Clustering of columns indicate influential taxa groupings for age prediction, whereas clustering of rows highlight individuals with similar significant features.

Fig. 5. Interpretable feature importances.

a Sample level feature importances for WGS Skin microbiome samples. Features colored in blue indicate a feature that influences the prediction of that sample to be younger, whereas red is indicative of a feature that drives the prediction to be older. b Feature level comparison of feature importances for WGS skin between TRPCA and RF model. c A pairwise comparison of the Pearson correlation between feature importances for each model architecture. SVR, GBR, KNN, RF, and TRPCA feature importances are highly correlated for WGS Skin microbiome feature importances.

Fig. 6. Paired sample age residual errors.

a TRPCA predictions for 16S stool samples from the THDMI and FINRISK cohorts. b TRPCA predictions for WGS stool samples from the THDMI and FINRISK cohorts. c The correlation between prediction errors in paired 16S and WGS stool microbiome samples (R2 = 0.632). d TRPCA predictions for WGS oral microbiome samples from curatedMetagenomicData. e TRPCA predictions for WGS stool microbiome samples from curatedMetagenomicData. f The correlation between prediction errors in paired WGS oral and WGS stool microbiome samples (R2 = 0.339). The error in model prediction for paired samples (WGS Stool/16S Stool and WGS Stool/WGS oral) implies host associated attributes may be associated with residual error.