Self-supervised learning enables robust microbiome predictions in data-limited and cross-cohort settings

preprint OA: gold CC-BY-NC-4.0
📄 Open PDF Full text JSON View at publisher

Abstract

The gut microbiome plays a crucial role in human health, but machine learning applications in this field face significant challenges, including limited labeled data availability, high dimensionality, and batch effects across different cohorts. To address these limitations, we developed representation learning models for gut microbiome metagenomic data, drawing inspiration from foundation models approaches based on self-supervised and transfer learning principles. By leveraging a large collection of 85,364 metagenomic samples, we implemented multiple self-supervised learning methods, including masked autoencoders with varying masking rates and adapted single-cell RNA sequencing models (scVI and scGPT), to generate embeddings from bacterial abundance profiles. These learned representations demonstrated significant advantages over raw bacterial abundances in two key scenarios: first, when training predictive models with very limited labeled data, improving prediction performance for age (r = 0.14 vs. 0.06), BMI (r = 0.16 vs. 0.11), visceral fat mass (r = 0.25 vs. 0.18), and drug usage classification (PR-AUC = 0.81 vs. 0.73); and second, when generalizing predictions across different cohorts, consistently outperforming models based on raw abundances in cross-dataset evaluation. Our approach provides a valuable framework for leveraging self-supervised representation learning to overcome the data limitations inherent in microbiome research, potentially enabling more robust and generalizable machine learning applications in this field.
Full text 57,918 characters · extracted from oa-pdf · 6 sections · click to expand

Abstract

20 The gut microbiome plays a crucial role in human health, but machine learning applications in this field 21 face significant challenges, including limited labeled data availability, high dimensionality, and batch 22 effects across different cohorts. To address these limitations, we developed representation learning 23 models for gut microbiome metagenomic data, drawing inspiration from foundation models approaches 24 based on self -supervised and transfer learning principles. By leveraging a large collection of 85,364 25 metagenomic samples, we implemented multiple self -supervised learning methods, including masked 26 autoencoders with varying masking rates and adapted single -cell RNA sequencing models (scVI and 27 scGPT), to generate embeddings from bacterial abundance profiles. These learned representations 28 demonstrated significant advantages over raw bacterial abundances in two key scenarios: first, when 29 training predictive models with very limited labeled data, improving prediction performance for age (r = 30 0.14 vs. 0.06), BMI (r = 0.16 vs. 0.11), visceral fat mass (r = 0.25 vs. 0.18), and drug usage classification 31 (PR-AUC = 0.81 vs. 0.73); and second, when generalizing predictions across different cohorts, 32 consistently outperforming models based on raw abundances in cross-dataset evaluation. Our approach 33 provides a valuable framework for leveraging self -supervised representation learning to overcome the 34 data limitations inherent in microbiome research, potentially enabling more robust and generalizable 35 machine learning applications in this field. 36 37 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint

Introduction

38 The gut microbiome exerts widespread influence on human health through its interconnected effects on 39 metabolism, immunity, and neurological functions. This dynamic community of microorganisms operates 40 through complex ecological networks and host interactions, which remain only partially understood. 41 Analysis of metagenomic data from gut microbiome samples offers valuable insights into these complex 42 microbial communities. 43 44 A significant challenge in microbiome research is the limited availability of data. Shotgun metagenomic 45 sequencing remains costly, and studies often include relatively small cohorts. This creates a fundamental 46 obstacle for machine learning applications, as the data is inherently high -dimensional, typically 47 represented as relative abundances across thousands of microbial species and millions of genes. The 48 combination of high dimensionality with small sample sizes leads to models that struggle with 49 generalizability and robustness. 50 51 A natural solution to limited sample sizes would be to combine multiple cohorts, but this approach 52 introduces its own challenges. Batch effects in metagenomic data make it difficult to combine multiple 53 cohorts or compare different populations. These technical variations can mask true biological signals, 54 limiting our ability to leverage diverse datasets to increase effective sample size. 55 56 Self-supervised and transfer learning approaches have shown remarkable success in other biological 57 domains, including protein structure prediction, genomic analysis, and single -cell transcriptomics 1 2 3 4. 58 These deep-learning methods leverage large amounts of unlabeled samples to learn the structure of the 59 data, training with labels that are generated from the data itself. By incorporating knowledge learned from 60 the large datasets, they can be used to create meaningful sample representations that transfer well to 61 downstream tasks with limited labeled data. Importantly, these approaches can mitigate overfitting to 62 dataset-specific signals, including batch effects, by learning more generalizable patterns across diverse 63 data sources. 64 65 These principles are particularly relevant for biological applications where labeled data is scarce but 66 unlabeled data is abundant. They enable the processing of samples from new cohorts with limited sample 67 sizes, leveraging patterns learned from larger datasets to extract meaningful features. In the microbiome 68 research field, large datasets are starting to accumulate — but they typically include very minimal host 69 information, which limits supervised learning applications. While recent work has begun exploring 70 microbiome representation learning 5 6 7 8, mainly using 16S sequencing data (which offers lower resolution 71 than metagenomics), this emerging field still offers significant opportunities for methodological 72 advancement. Critically, many representation learning works in biology fail to demonstrate clear 73 advantages over standard approaches used in the field, leaving their practical utility unclear 9,10. 74 75 To address these challenges, we developed representation learning models for gut microbiome shotgun 76 metagenomics that leverage large -scale unlabeled datasets totalling 85,364 samples. We explored 77 multiple architectural approaches, including masked autoencoders and cross -domain adaptation of 78 established single-cell RNA sequencing models. Our study makes three key methodological contributions: 79 (1) benchmarking that demonstrates clear advantages over the standard practice in microbiome machine 80 learning --- prediction from raw bacterial relative abundances; (2) novel cross -domain transfer of 81 successful single-cell RNA sequencing models to microbiome analysis; and (3) a pretrained model that 82 can be used to extract representations from new metagenomic samples. 83 84 Our learned representations demonstrate clear advantages over standard microbiome analysis 85 approaches in two critical scenarios: enhanced performance in limited-data settings that reflect real-world 86 study constraints, and improved generalizability across different populations and technical frameworks. 87 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint These advances directly address critical barriers to machine learning applications in microbiome research 88 and establish a framework for adapting successful biological representation learning methods across 89 related domains. 90 Representation learning models 91 To create meaningful representations of gut microbiome samples, we employed self-supervised learning 92 (SSL), in which deep learning models learn from unlabeled data by creating learning objectives directly 93 from the data structure itself. We gathered metagenomic data from three studies: a cohort of 66,151 94 participants from both the US (6,206; “1000_US”) and Israel (59,945; “1000_IL”) 11 , a cohort of 11,084 95 participants from Israel (“10_IL”) 12,13, and a cohort of 8,129 participants from the Netherlands (“32_NE”) 96 14. These datasets span different geographies, ancestries, collection periods and laboratories, providing a 97 robust foundation for learning generalizable microbiome representations. 98 99 All datasets included metagenomic sequencing data alongside basic host information (age, sex, Body 100 Mass Index (BMI)), with one cohort (10_IL, 13) also providing diverse clinical and phenotypic 101 measurements, including body composition profile from medical imaging, drug usage patterns, and 102 disease diagnoses. Our approach involved a two-stage process: first, we used only metagenomic data to 103 train the models and learn representations. Subsequently, we evaluated the effectiveness of these learned 104 representations by testing their impact on machine learning models predicting host phenotypes. 105 106 For each sample, we mapped sequencing reads to a reference database of bacterial genomes and 107 retained the 902 most prevalent species after filtering (requiring species to be present in at least 500 108 samples of the deeply phenotyped cohort; Methods). We calculated relative species abundances, creating 109 a species abundance vector (902 dimensions) for each sample — a standard representation of 110 microbiome samples. This vector serves as both our baseline representation (“raw'', Table 1) and the input 111 to our representation learning models. 112 113 To learn sample representations, we implemented multiple SSL approaches. First, we trained a masked 114 autoencoder. Masking is a common SSL approach, in which a proportion of the sample inputs are 115 obscured and the model is trained to reconstruct them. This approach is designed to learn underlying 116 microbial community patterns by forcing the model to predict missing species based on observed ones, 117 with the expectation of capturing ecological relationships and functional redundancies within microbial 118 communities. We masked the abundances of 10, 30, or 70% of the species in each sample, testing low, 119 medium, and high masking fractions. Starting with a basic deep -learning architecture, we trained multi -120 layer perceptron (MLP) models to reconstruct the hidden values, and used the encoded representation 121 from the middle layer as the sample representation. We set aside 10% of the data as a validation set, and 122 used it to choose a stopping point for training. In the self -supervised task of reconstructing masked 123 abundances, the 10%, 30%, and 70% masking models achieved an explained variance (EV) of 26.4%, 124 24.5%, and 19.1%, respectively, in the validation set averaged across all species (Supp. figs. 1-3). 125 126 We also evaluated a transformer -based masked autoencoder. Transformers, originally developed for 127 natural language processing, use self -attention mechanisms that can capture complex dependencies 128 between all input features simultaneously. This architecture offered the potential to model ecological 129 interactions between bacterial species better than the simpler feedforward networks. In our 130 implementation, we stratified the relative abundances of each species into quantiles, assigned a token per 131 species quantile, and applied the same masking-reconstruction task as with the MLP models. To extract 132 sample representations for downstream tasks, we mean- and max-pooled the output of the trained model 133 (Methods). 134 135 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint Additionally, recognizing the shared characteristics between metagenomic and single -cell RNA 136 sequencing (scRNA-seq) data – both involve high-dimensional, sparse (zero-inflated), tabular data, with 137 features (genes or species) exhibiting highly variable abundance distribution, and pronounced batch 138 effects – we explored the potential of adapting well -established representation learning models from the 139 scRNA-seq domain. While representation learning for microbiome data remains in its nascent stages, the 140 scRNA-seq field has advanced significantly in recent years, with numerous foundation models designed 141 specifically to handle high -dimensional, sparse biological data. We hypothesized that the underlying 142 principles of these models, designed to capture complex biological variations and mitigate technical noise 143 in scRNA-seq data, could be effectively transferred to the analysis of gut microbiome compositions. To 144 this end, we adapted and evaluated two prominent scRNA-seq representation learning models: scVI 3 and 145 scGPT 15. 146 147 scVI (single-cell Variational Inference) is a variational autoencoder (VAE)-based model that learns a low-148 dimensional latent space by modeling observed data as arising from a probabilistic generative process. In 149 our adaptation, we treated bacterial species abundances analogously to gene expression levels, allowing 150 the model to capture the complex distribution of microbiome compositions. This approach enables explicit 151 incorporation of batch information, helping to disentangle technical variation from biological signals. The 152 probabilistic framework of scVI is particularly well -suited for addressing zero -inflation, inherent 153 measurement noise, and the varying abundance distributions across different bacterial species that 154 characterize both scRNA-seq and metagenomic data, potentially yielding more robust representations of 155 microbial communities. 156 157 Conversely, scGPT represents a foundation model approach for single -cell omics data, utilizing 158 transformer architectures to capture complex patterns in sparse, high-dimensional biological data. Its self-159 attention mechanisms, originally designed to model gene-gene interactions in single cells, could potentially 160 capture ecological relationships between bacterial species in microbiome communities. We adapted 161 scGPT for microbiome analysis, treating bacterial species abundances analogously to gene expression 162 values, with the hypothesis that its specialized architecture for modeling biological interactions would 163 effectively capture patterns of bacterial co-occurrence and mutual exclusion in gut microbiome samples. 164 165 In addition to species abundance information, we provided the scVI and scGPT models with batch and 166 cohort information, demonstrating a multi -modal approach that jointly models biological features with 167 technical covariates to handle technical variation and biological heterogeneity. 168 Sample embeddings improve phenotype predictions from microbiome 169 We designed comprehensive evaluations to rigorously test whether our learned representations provide 170 advantages over the current standard practice in microbiome machine learning: using raw bacterial 171 relative abundances as features for predictive models. We leveraged a dataset with rich host information 172 to train tree-based models 16 to predict host phenotypes, comparing the performance of models trained on 173 sample embeddings to those trained on raw bacterial relative abundances — the standard approach in 174 microbiome machine learning. We performed 4 repetitions of 5-fold cross-validation to robustly assess the 175 predictive performance of different representations. We trained models to predict host age, sex, BMI, 176 Visceral Adipose Tissue (VAT) mass, Proton Pump Inhibitors (PPI) intake, hyperlipidemia diagnosis, and 177 fatty liver disease diagnosis. 178 179 For the first benchmarking analysis, we used the entire deeply phenotyped cohort (n = 11,084; Table 1). 180 Age prediction based on the 30 and 70% mask-AE models embedding was better than the prediction from 181 the raw relative abundances, with mean Pearson correlations of 0.34 for the raw features, 0.37 for the 182 30% mask-AE (q-value, FDR adjusted = 0.0003), and 0.36 for the 70% mask -AE (q-value = 0.002). The 183 superiority of these models’ embeddings in predicting Visceral Adipose Tissue (VAT) mass was marginally 184 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint significant, with a Pearson correlation of 0.38 for both these models, and 0.37 for the raw abundances (q-185 value = 0.05, 0.09). In other cases, we did not observe a significant advantage for models trained on the 186 learned representations rather than the bacterial abundances. For the transformer model, the different 187 pooling strategies we tested yielded similar results (Supp. table 1). 188 189 190 191 Table 1 - Prediction performance across different microbiome representations. Performance 192 evaluation was based on four repetitions of a 5-fold cross-validation, resulting in a total of 20 train-test split 193 experiments. The table shows the performance (continuous: Pearson correlation, binary: PR-AUC) mean 194 and standard deviation across all test sets. Q -values are based on the Mann -Whitney U test comparing 195 each performance distribution to the raw model’s performances, and corrected for multiple testing using 196 FDR. 197 198 In the second benchmarking analysis, we tested the benefit of the learned representations in limited-data 199 scenarios. We hypothesized that in scenarios in which there are much fewer labels to learn from, learning 200 features in an SSL scheme will be more advantageous, as small cohorts lack the diversity and sample 201 size needed to learn complex microbiome relationships from scratch. To simulate this, we downsampled 202 our cohort to 100 training samples — a typical scale for many microbiome studies, clinical trials in particular 203 — and used them to train the phenotype predictors from either the representations or the raw bacterial 204 relative abundances (Table 2, Supp. table 2). 205 206 Indeed, the advantage of representation learning became much more pronounced in this data -limited 207 setting. Predicting age from the learned representations improved the performance from a Pearson 208 correlation of 0.06 to 0.14 (scVI, q-value<10-5). The Pearson correlation for BMI prediction improved from 209 0.11 to 0.16 (mask-AE 30%, q-value=0.003), and VAT mass prediction improved from 0.18 to 0.25 (mask-210 AE 30%, p-value=0.0004). In the classification of individuals using proton pump inhibitor (PPI) drugs, PR-211 AUC improved from 0.73 to 0.81 (mask-AE 70%, p-value<10-11). In the classification of sex and fatty liver 212 disease, the mean PR-AUC achieved with the learned representations was also better than that gained 213 with the raw abundances, but the difference was not statistically significant. This demonstrates the 214 advantage of learned representations from SSL models in training supervised models on small datasets. 215 216 217 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 218 Table 2 - Prediction performance across microbiome representations with limited (n=100) training 219 data. Train cohorts consisted of 100 randomly selected samples. Performance distribution is based on 220 100 repetitions. The table shows the performance (continuous: Pearson correlation, binary: PR -AUC) 221 mean and standard deviation across all test sets. Q -values are based on the Mann -Whitney U test 222 comparing each performance distribution to the raw model’s performances, and corrected for multiple 223 testing using FDR. 224 Microbiome predictors trained on learned representations are more 225 generalizable to new cohorts 226 Microbiome datasets differ due to both biological and technical factors. Sample collection and storage, 227 DNA extraction, PCR amplification and library preparation, sequencing platforms and runs, and data 228 processing protocols all affect the observed species composition. In addition, environmental variables and 229 health differences between populations are reflected in different microbiome profiles. These differences 230 impair the ability to apply machine learning models trained on one dataset to label samples from another 231 dataset — a critical requirement for clinical deployment. We hypothesized that using features learned by 232 the SSL models, rather than the raw bacterial abundances, would enable models to generalize better 233 between datasets. 234 235 To test if the learned representations improve the generalizability of supervised models, we separated our 236 dataset to four subsets (by study and country), trained tree -based models predicting age, sex, and BMI 237 on each subset, and evaluated it on the other three (Supp. fig. 4; Methods). Notably, we did not finetune 238 the supervised models on the target cohort — only predicted the labels of the target cohort using the 239 models that were fitted on the source cohort, testing how well the models generalize across different 240 populations. Overall, across five representation models and three phenotypes, the learned representations 241 improved cross -cohort prediction in 148 of 180 comparisons (82%) relative to raw features, with an 242 average improvement of 10% ± 15% (Fig. 1, Supp. Fig. 4). When transferring age and BMI models to 243 different cohorts, scVI and the three mask -AE representations each outperformed the raw -based model 244 in 10 to 12 of the 12 train -test cohort pairs (q -values=0.03 to 0.001). The highest improvement was 245 demonstrated when applying the age prediction model that was trained on the 32_NE cohort to predict 246 age in the 10_IL cohort — the model based on mask -AE 30% representations outperformed the raw -247 based model by 81%. For sex classification, improvements were more modest, achieving an increase in 248 PR-AUC of up to 14% (q -values=0.09 to 0.007). The only model whose representations did not 249 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint significantly outperform the baseline was scGPT. These results demonstrate the advantage of learned 250 microbiome representations as features for models that generalize better across cohorts. 251 252 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 253 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint Figure 1 - cross-cohort generalizability of models compared to the baseline model. Each point 254 represents a train-test cohort pair, showing the ratio of representation -based model performance to raw 255 abundance model performance. Values >1.0 indicate superior performance of learned representations. 256 Q-values are FDR-adjusted. Original performance metrics of each model are detailed in Supp. Fig. 4. 257 Model application to new studies 258 To facilitate reproducibility and enable broader application of our findings, we make our best -performing 259 masked autoencoder model available to the microbiome research community. The pretrained 30% 260 masked autoencoder model, which demonstrated superior performance in our evaluations, is available 261 along with example code showing how to generate embeddings from bacterial abundance profiles. 262 263 The pretrained model accepts the relative abundances of bacterial species in samples and outputs the 264 corresponding embedding vectors that can be used for downstream machine learning tasks. These 265 embeddings may be particularly valuable for improving performance in scenarios with limited training data 266 or when attempting to develop models that generalize across different cohorts. 267 Discussions 268 In this study, we developed representation learning models for gut microbiome metagenomic data and 269 demonstrated that self -supervised learning can significantly improve microbiome -based predictive 270 modeling in two critical scenarios: limited -data settings and cross -cohort generalization. These 271 improvements address fundamental challenges in microbiome machine learning with clear implications 272 for clinical translation. 273 274 Self-supervised learning and transfer learning approaches have recently emerged in microbiome research 275 5 6 7 8. Contributing to this emerging area, our work leverages a large metagenomic dataset of 85,364 276 samples to train multiple representation learning models and systematically evaluates their performance 277 against the current standard practice in microbiome machine learning — using raw bacterial abundances 278 with standard machine learning methods — across diverse prediction scenarios, addressing a 279 fundamental need to demonstrate practical utility over established methods 9,10. 280 281 Particularly noteworthy was the consistent superiority of representation -based models in cross -cohort 282 prediction tasks, improving prediction accuracy by 10% on average and up to 81% compared to models 283 trained on raw features. The ability of our models to generalize across different populations and data 284 collection protocols addresses one of the most significant challenges in microbiome machine learning: 285 batch effects and population -specific signatures that typically impair model transferability. By learning 286 more robust, generalizable features that capture underlying biological signals rather than technical 287 variations, our approach offers a promising avenue for developing more universal microbiome -based 288 predictive models. This enhanced generalizability addresses one of the most significant barriers to clinical 289 translation of microbiome research: the need for models that perform reliably across diverse populations, 290 healthcare systems, and technical protocols. 291 292 The significant improvements in data -limited settings demonstrate how self -supervised pretraining 293 enables effective knowledge transfer, allowing models to leverage patterns learned from large unlabeled 294 datasets when supervised training data is scarce. While the performance gains were modest when using 295 the full dataset, they became substantially more pronounced when restricting training to just 100 samples. 296 This pattern aligns with findings in other domains where representation learning shows the greatest 297 advantage in low-data regimes, effectively leveraging patterns learned from larger unlabeled datasets to 298 improve performance on specific tasks with limited labeled examples, and is particularly relevant for 299 microbiome research. While our deeply-phenotyped cohort contains over 10,000 samples, the majority of 300 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint microbiome studies operate with cohorts 1 -3 orders of magnitude smaller, typically consisting of tens to 301 hundreds of samples. These smaller studies often focus on specific conditions or populations that cannot 302 feasibly be expanded to larger scales. This capability could enable meaningful discoveries in studies that 303 were previously underpowered due to sample size constraints. 304 305 Among the representation models we evaluated, the 30% masked autoencoder consistently demonstrated 306 the best overall performance across prediction tasks, followed by the 70% masking model. This suggests 307 that intermediate levels of feature masking may provide an optimal balance between preserving sufficient 308 information for meaningful reconstruction while forcing the model to learn robust underlying patterns. 309 Despite the theoretical advantages of transformer architectures in capturing complex inter -species 310 relationships through self -attention mechanisms, our transformer -based model did not outperform the 311 simpler MLP-based approaches. This pattern was mirrored in the adapted single -cell models: while scVI 312 showed competitive performance and even outperformed other methods in specific scenarios (age 313 prediction in small cohorts and BMI prediction in cross -cohort analyses), the more sophisticated 314 transformer-based scGPT model generally underperformed relative to both scVI and our custom models. 315 316 Our work builds upon recent successes of representation learning in other biological domains, particularly 317 in single -cell RNA sequencing. The scRNA -seq field has developed sophisticated computational 318 approaches to overcome challenges that are similar to those in microbiome research: high dimensionality, 319 technical variation, zero -inflated sparse data, and limited sample sizes. Models like scVI have become 320 standard tools for integrating and analyzing diverse scRNA-seq datasets, suggesting similar approaches 321 could benefit microbiome research. Our successful adaptation of single -cell RNA sequencing model 322 architectures to microbiome analysis demonstrates the value of cross-domain methodological transfer in 323 biological data analysis, leveraging mature computational methods from one biological domain to advance 324 another. 325 326 These advances have practical implications for microbiome research applications. By enabling smaller 327 studies to achieve meaningful predictive performance and improving model transferability across 328 populations, our approach could accelerate the identification of diagnostic microbiome signatures and 329 support the development of models that perform consistently across diverse cohorts. This enhanced 330 analytical capability may be particularly valuable for studies involving rare conditions, controlled 331 interventions, or unique environmental exposures, where large sample collection is challenging. To 332 support these applications, we provide our best-performing pretrained model for community use. 333 334 While our training dataset is large by microbiome standards, it remains modest compared to foundation 335 model datasets in other domains. Future work could benefit from even larger datasets, representing a 336 bigger variety of geographies and health states. Other models built upon this framework and an even 337 larger dataset may include a larger selection of species — including the rare ones — as well as greater 338 resolution of features, like strains and genes. Our representations also sacrifice some interpretability 339 compared to raw abundance approaches, though the performance gains may justify this trade-off in many 340 applications. 341 342 In conclusion, we demonstrate that self -supervised learning on large metagenomic datasets can 343 substantially enhance model generalizability across diverse populations and improve prediction accuracy 344 in data-limited scenarios. This work establishes representation learning as a viable solution for overcoming 345 fundamental challenges in microbiome-based machine learning, opening new possibilities for leveraging 346 larger unlabeled datasets to improve predictions in specialized studies where large labeled cohorts remain 347 challenging to obtain. 348 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint

Methods

349 Data 350 We based our work on data from three previously published studies: a cohort of 66,151 individuals from 351 the US and from Israel (“1000_US” and “1000_IL”, accordingly) 11, a cohort of 8,129 individuals from the 352 Netherlands (“32_NE”) 14, and a cohort of 11,084 individuals from Israel (“10_IL”) 12,13. All datasets 353 included gut microbiome metagenomic data alongside basic host information: age, sex, and BMI. From 354 the 10_IL cohort, we also obtained data on body composition, drug usage, and medical diagnoses. We 355 only included the earliest sample of each participant. 356 Preprocessing the microbiome data 357 We processed the metagenomic samples and aligned them to bacterial reference genomes as described 358 in Leviatan et al. 17 to infer species composition. We filtered the species by prevalence, including species 359 seen in at least 500 samples in the deeply -phenotyped (“10_IL”) cohort (out of 11,084 samples) — 360 resulting in 902 species. For these 902 species, we calculated their relative abundances in each sample 361 (summing to 1). Our pipeline’s detection threshold is set to 0.0001, and we truncated any smaller 362 abundance value to this minimum. 363 Masked autoencoder 364 To create sample representations using masked autoencoders, we implemented multi -layer perceptron 365 (MLP) architectures trained on a self-supervised reconstruction task. We preprocessed bacterial relative 366 abundance data by applying a log10 transformation and shifting values by 2 to convert the range from [ -367 4, 0] to [-2, 2], facilitating more stable neural network training. 368 369 We employed a masking scheme where we randomly obscured 10%, 30%, or 70% of input features for 370 each sample during training, replacing masked positions with a fixed value of -2. The masking pattern was 371 randomly generated for each sample in every training batch. We computed the mean squared error loss 372 only on the obscured positions. 373 374 We optimized hyperparameters through a systematic hyperparameter sweep, testing combinations of 375 masking fractions, encoding dimensions, number of layers, learning rates, batch sizes, dropout rates, 376 activation functions, masking values, whether to shift values by 2 or not, and loss computation strategies 377 (loss on all features vs. masked features only). For each masking fraction, we selected the best-performing 378 hyperparameter combination, resulting in models with different optimal configurations. The 10% and 70% 379 masking models employed Mish activation with 10% dropout, while the 30% masking model used GELU 380 activation with 25% dropout. All final models used a learning rate of 0.0005 with Adam optimization, 381 training for 300-350 epochs. Our autoencoder architecture consisted of an encoder-decoder structure with 382 single-layer components and batch normalization. The encoder transformed the 902 -dimensional input 383 (representing bacterial species abundances) to a 1024 -dimensional latent representation, while the 384 decoder reconstructed the original input dimensions. 385 386 We split the dataset into 90% training and 10% validation sets, monitoring reconstruction performance 387 using mean squared error loss and explained variance. The learned representations were extracted from 388 the encoder's output layer after training was completed. 389 Masked transformer 390 We also implemented a transformer -based masked autoencoder using self -attention mechanisms to 391 capture potential interactions between bacterial species. For input representation, we tokenized bacterial 392 abundance values using species-specific quantile binning. For each bacterial species independently, we 393 divided the abundance range into 10 bins: one bin for the minimal value (-4), and 9 bins based on quantiles 394 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint of non-minimal abundances. Each species-abundance combination was assigned a unique token, creating 395 a vocabulary of 9,021 tokens (902 species × 10 bins + 1 mask token). This quantile -based approach 396 preserves relative abundance information while creating discrete tokens suitable for transformer 397 processing. 398 399 The transformer architecture consisted of 2 encoder layers with 256-dimensional embeddings, 16 attention 400 heads, and feedforward layers with 512 hidden units. We employed a 10% dropout and trained the model 401 as a regression task, predicting continuous abundance values rather than discrete tokens. Each bacterial 402 species had its own linear decoder head that converted the transformer's output embeddings to 403 abundance predictions. 404 405 We applied the same masking strategy as the MLP models, randomly masking 30% of input features and 406 replacing them with a special mask token. The mean squared error loss was computed only on the masked 407 and reconstructed inputs. We used Adam optimization with a learning rate of 0.001, a batch size of 256, 408 and trained for 260 epochs with a 90% -10% train -validation split. The learned representations were 409 extracted by max -pooling, mean -pooling, or by concatenating mean -pooling and max -pooling of the 410 transformer's output across the species dimension, resulting in 512 -dimensional sample representations 411 (256 dimensions each from mean and max pooling). 412 scVI 413 To create sample representations with scVI 3, we used scVI-tools version 1.2.2.post2 with Python 3.11.5. 414 To align with the expected input format of the scRNA -seq models, which are typically applied to count 415 data, we converted bacterial relative abundances to pseudo-counts by multiplying each sample's relative 416 abundances by its total mapped read count, then converted the resulting values to integers. We replaced 417 abundance values of 0.0001 with zero to reflect the detection limit of our data. 418 419 We constructed batch identifiers combining country, study type, DNA extraction kit, library preparation 420 method, and collection year to account for technical variation across our diverse datasets. We configured 421 scVI with a 2 -layer encoder-decoder architecture using 512 hidden units and 256 latent dimensions. To 422 model the count distribution, we used the zero -inflated negative binomial (ZINB) distribution with gene -423 batch-specific dispersion parameters — to model the sparsity and the variability of distributions between 424 species and batches that is common in microbiome data. We enabled covariate encoding and used 425 embedding-based batch representation with 5-dimensional batch embeddings. 426 427 We trained the model from scratch for up to 50 epochs using a batch size of 128 and a validation set 428 comprising 5% of samples. We implemented early stopping with a patience of 3 epochs and used a 429 learning rate of 0.001 with KL divergence warmup over 6 epochs. Learning rate reduction on plateau was 430 enabled with a patience of 2 epochs. After training, we extracted the latent representations from the trained 431 model to serve as sample embeddings for downstream analysis. 432 scGPT 433 To create sample representations with scGPT 15, we used scGPT 0.2.1 with Python 3.11.3 and Torch 434 2.3.1. We trained on each sample’s vector of species abundance, tokenizing each species abundance so 435 that a sequence of abundance values became a sequence of tokens. We used a masking ratio of 0.4 and 436 25 bins to tokenize the relative abundance values. As we did not have a cell type feature, we disabled the 437 elastic cell similarity objective by setting its relative weight to zero. We used the sequencing run batch 438 label (“RunName”) as the batch identifier for batch integration. 439 440 We chose the transformer hyperparameters based on the lowest masked MSE evaluation loss (L2 norm 441 between true and predicted values for masked tokens), and tested embedding dimensions of 128 to 768, 442 and depths between two to six. We trained the final model from scratch using a learning rate of 1e -3, a 443 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint batch size of 64, and set the number of transformer encoder layers to three, using 32 heads and a hidden 444 token dimension of 256. Automatic mixed precision was used. We fixed the random seed for the training 445 experiment to 42. We used the masking value of -1. We set abundance values of 0.0001 to zero in keeping 446 with the detection limit, and ignored ‘cell type’ as it was not relevant for our use case. We created a 447 vocabulary based on species names as opposed to gene names. We used the data in log1p form. We 448 trained for 10 epochs, using all other default settings. We extracted the CLS token (cell embeddings in the 449 original paper) after training as the sample representation for downstream analysis. 450 Full cohort phenotype predictions 451 We conducted an analysis to evaluate the advantage of the various models for generating features for 452 supervised models. For this analysis, we used the data from the 10_IL cohort, which includes information 453 about host health. We trained tree-based models to predict host age, sex, Body Mass Index (BMI), Visceral 454 Adipose Tissue (VAT) mass, Proton Pump Inhibitors (PPI) intake, hyperlipidemia diagnosis, and fatty liver 455 disease diagnosis, from either the raw microbiome abundances, or the sample embeddings extracted from 456 the pretrained models. It is important to note that the embeddings were extracted from the self-supervised 457 models, which were trained without the sample labels. We used the LightGBM package 18 with 2000 458 estimators, a learning rate of 0.001, max_depth of 3, min_child_samples of 30, 70% samples sampling, 459 60% features sampling, and an initial random seed of 42. We used a five -fold, stratified (for binary 460 phenotypes), cross-validation scheme, and used 15% of the train samples as a validation set to enforce 461 early stopping (after 80 rounds without improvement). We evaluated model performances on the test set 462 using Pearson’s correlation for continuous phenotypes and the area under the precision recall curve (PR-463 AUC) for binary phenotypes. We conducted four repetitions of these experiments, for a total of 20 sets of 464 train, validation, and test samples (over five folds and four repetitions) — from which we derived the mean 465 and the standard deviation. To compare the test prediction performances of each model, we compared 466 this 20-long vector with that of the predictor trained on species relative abundances using the Mann -467 Whitney U test 19 - (alternative=’greater’). We adjusted for multiple testing across all model -phenotype 468 comparisons using the Benjamini-Hochberg False Discovery Rate (FDR) procedure 20. 469 Downsampling analysis 470 To evaluate prediction performances of models trained on 100 labeled samples, we conducted an 471 experiment similar to the one described in the previous paragraph, with a few changes. For each model, 472 we randomly selected 100 samples for the training set and 100 samples for the test set. For the binary 473 phenotypes, we sampled 50 samples from each of the label groups. In this experiment, we did not apply 474 early stopping, but reduced the number of estimators to 1000. We also did not use cross validation. For 475 each pretrained model we compared and each target phenotype, we repeated this experiment 100 times, 476 and used the Mann-Whitney U test to compare the performances to those gained by the model trained on 477 species relative abundances. We adjusted for multiple testing across all model -phenotype comparisons 478 using Benjamini-Hochberg False Discovery Rate (FDR) procedure 20. 479 Cross-cohort predictions analysis 480 To test how well models trained on different representations predict phenotypes in cohorts different from 481 those used to train them, we each time trained a prediction model on one cohort, and tested its prediction 482 performance in other cohorts. To train the models, we used 85% of the cohort samples, and used the 483 remaining 15% both for the within-cohort evaluation, and for performing early stopping (after not improving 484 for 50 rounds). Then, to evaluate the trained model on a different target cohort, we tested it on 100% of 485 the target cohort samples. For this reason, in this analysis we only performed a single evaluation per 486 cohort pair, and did not have a distribution of performance values over folds and repetitions like in the 487 previous analyses. For the supervised models we used LightGBM, with 2000 estimators and the same 488 parameters used in the previous analyses. We adjusted for multiple testing across all model -phenotype 489 comparisons using Benjamini-Hochberg False Discovery Rate (FDR) procedure. 490 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint Model availability 491 The best performing trained model, MAE-30%, is available at https://github.com/LironZa/MBEmbed along 492 with code to create representations of new metagenomic samples for machine learning applications with 493 new microbiome datasets. 494 Competing Interests 495 The authors declare no competing interests. 496 Author contribution 497 L.Z. conceived the project, designed and conducted all analyses, interpreted the results, and wrote the 498 manuscript. Z.L. performed data analysis and reviewed the manuscript. A.G. performed data analysis. 499 V.D., R.D., K.S.P. and A.W. reviewed the manuscript. E.S. conceived, directed and designed the project 500 and analyses, interpreted the results and wrote the manuscript. 501

Acknowledgements

502 We thank the members of the Segal group for fruitful discussions. E.S. is supported by the Crown Human 503 Genome Center; the Larson Charitable Foundation New Scientist Fund; the Else Kröner Fresenius 504 Foundation; the White Rose International Foundation; the Ben B. and Joyce E. Eisenberg Foundation; the 505 Nissenbaum Family; Marcos Pinheiro de Andrade and Vanessa Buchheim; Lady Michelle Michels; Aliza 506 Moussaieff; and grants funded by the Minerva Foundation, with funding from the Federal German Ministry 507 for Education and Research and by the European Research Council and the Israel Science Foundation. 508

References

509 1. Jumper, J. et al. Highly accurate protein structure prediction with AlphaFold. Nature 596, 583–589 510 (2021). 511 2. Ji, Y., Zhou, Z., Liu, H. & Davuluri, R. V. DNABERT: pre-trained Bidirectional Encoder 512 Representations from Transformers model for DNA-language in genome. Bioinformatics 37, 2112–513 2120 (2021). 514 3. Lopez, R., Regier, J., Cole, M. B., Jordan, M. I. & Yosef, N. Deep generative modeling for single-cell 515 transcriptomics. Nat. Methods 15, 1053–1058 (2018). 516 4. Theodoris, C. V. et al. Transfer learning enables predictions in network biology. Nature 618, 616–517 624 (2023). 518 5. Pope, Q., Varma, R., Tataru, C., David, M. & Fern, X. Learning a deep language model for 519 microbiomes: the power of large scale unlabeled microbiome data. BioRxiv (2023) 520 doi:10.1101/2023.07.17.549267. 521 6. Zhang, H. et al. MGM as a large-scale pretrained foundation model for microbiome analyses in 522 diverse contexts. BioRxiv (2025) doi:10.1101/2024.12.30.630825. 523 7. Oh, M. & Zhang, L. DeepGeni: deep generalized interpretable autoencoder elucidates gut 524 microbiota for better cancer immunotherapy. Sci. Rep. 13, 4599 (2023). 525 8. Oh, M. & Zhang, L. DeepMicro: deep representation learning for disease prediction based on 526 microbiome data. Sci. Rep. 10, 6026 (2020). 527 9. Boiarsky, R. et al. Deeper evaluation of a single-cell foundation model. Nat. Mach. Intell. 6, 1443–528 1446 (2024). 529 10. Wong, D. R., Hill, A. S. & Moccia, R. Simple controls exceed best deep learning algorithms and 530 reveal foundation model effectiveness for predicting genetic perturbations. Bioinformatics 41, 531 (2025). 532 11. Rothschild, D. et al. An atlas of robust microbiome associations with phenotypic traits based on 533 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint large-scale cohorts from two continents. PLoS ONE 17, e0265756 (2022). 534 12. Zahavi, L. et al. Bacterial SNPs in the human gut microbiome associate with host BMI. Nat. Med. 535 29, 2785–2792 (2023). 536 13. Reicher, L. et al. Deep phenotyping of health-disease continuum in the Human Phenotype Project. 537 Nat. Med. (2025) doi:10.1038/s41591-025-03790-9. 538 14. Gacesa, R. et al. Environmental factors shaping the gut microbiome in a Dutch population. Nature 539 604, 732–739 (2022). 540 15. Cui, H. et al. scGPT: toward building a foundation model for single-cell multi-omics using generative 541 AI. Nat. Methods 21, 1470–1480 (2024). 542 16. Grinsztajn, L., Oyallon, E. & Varoquaux, G. Why do tree-based models still outperform deep 543 learning on typical tabular data? Advances in Neural Information Processing Systems (2022). 544 17. Leviatan, S., Shoer, S., Rothschild, D., Gorodetski, M. & Segal, E. An expanded reference map of 545 the human gut microbiome reveals hundreds of previously unknown species. Nat. Commun. 13, 546 3863 (2022). 547 18. Ke, G. et al. LightGBM: A Highly Efficient Gradient Boosting Decision Tree. Advances in Neural 548 Information Processing Systems (2017). 549 19. Virtanen, P. et al. SciPy 1.0: fundamental algorithms for scientific computing in Python. Nat. 550

Methods

17, 261–272 (2020). 551 20. Benjamini, Y. & Yekutieli, D. The Control of the False Discovery Rate in Multiple Testing under 552 Dependency. Annals of statistics (2001). 553 554 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint Supplementary Data 555 556 Supplementary figure 1 — performances of the MAE -10% model reconstructing masked 557 abundances. Top: train loss, validation loss, and explained variance on the validation set per 558 epoch. Explained variance was calculated per species, and then averaged. Bottom: distribution of 559 the average explained (left) or correlation between original and reconstructed abundances (right) 560 of each species, calculated on the validation set in the final epoch. 561 562 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 563 Supplementary figure 2 — performances of the MAE -30% model reconstructing masked 564 abundances. Top: train loss, validation loss, and explained variance on the validation set per 565 epoch. Explained variance was calculated per species, and then averaged. Bottom: distribution of 566 the average explained variance (left) or correlation between original and reconstructed 567 abundances (right) of each species, calculated on the validation set in the final epoch. 568 569 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 570 Supplementary figure 3 — performances of the MAE -70% model reconstructing masked 571 abundances. Top: train loss, validation loss, and explained variance on the validation set per 572 epoch. Explained variance was calculated per species, and then averaged. Bottom: distribution of 573 the average explained variance (left) or correlation between original and reconstructed 574 abundances (right) of each species, calculated on the validation set in the final epoch. 575 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 576 Supplementary table 1 — Prediction performance across different microbiome representations, 577 including multiple pooling strategies. Performance evaluation was based on four repetitions of a 5-fold 578 cross-validation, resulting in a total of 20 train -test split experiments. The table shows the performance 579 (continuous: Pearson correlation, binary: PR-AUC) mean and standard deviation across all test sets. Q -580 values are based on the Mann-Whitney U test comparing each performance distribution to the raw model’s 581 performances, and corrected for multiple testing using FDR. 582 583 584 Supplementary table 2 — Prediction performance across microbiome representations with limited 585 (n=100) training data, including multiple pooling strategies. Train cohorts consisted of 100 randomly 586 selected samples. Performance distribution is based on 100 repetitions. The table shows the performance 587 (continuous: Pearson correlation, binary: PR-AUC) mean and standard deviation across all test sets. Q -588 values are based on the Mann-Whitney U test comparing each performance distribution to the raw model’s 589 performances, and corrected for multiple testing using FDR. 590 591 592 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint 593 Supplementary figure 4 — test performance of supervised predictors trained on one cohort and 594 tested on another. 595 .CC-BY-NC 4.0 International licenseperpetuity. It is made available under a preprint (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in The copyright holder for thisthis version posted October 20, 2025. ; https://doi.org/10.1101/2025.10.19.683269doi: bioRxiv preprint

Text is read by the "Ask this paper" AI Q&A widget below. Extraction quality varies by source — PMC NXML preserves structure cleanly, OA-HTML may include some navigation residue, and OA-PDF can have broken hyphenation. The publisher copy (via DOI) is the canonical version.

My notes (saved in your browser only)

Ask this paper AI returns verbatim quotes from the full text · source: oa-pdf

Answers must be backed by verbatim quotes from this paper's full text. Hallucinated quotes are dropped automatically; if no verbatim passage answers the question, we say so. How this works

Citation neighborhood (no data yet)

We don't have any in-corpus citations linked to this paper yet. This is a recent paper (2025) — citers typically take a year or two to land, and the OpenAlex reference graph may still be filling in.

Source provenance

europepmc
last seen: 2026-05-20T01:45:00.602351+00:00
unpaywall
last seen: 2026-05-21T05:10:58.409756+00:00
License: CC-BY-NC-4.0