Microbiome World Modelling

Your gut can increase your IQ, give you better cardio, let you eat ice cream all day and not get fat and make you feel euphoria. All you need is a good microbiome.

The food you eat plays a massive role, but so do factors outside your control. You inherit your microbiome from your mother at birth, then from your environment as you grow up. You are covered in nooks and crannies in which microbes form resilient communities. Most of the time, they are neutral or beneficial, but sometimes they can fall into negative stable states. Due to the enclosed nature of these environments, these states can embed themselves for the long haul and then be passed on to your children. In fact, it is thought that over the past 200 years of modernisation, we have seen an extinction in microbiome diversity, potentially irreparable.

Clearly there is something to study here, and even better, model. As of today, the space is wide open. It is still not clear how these microbes have such a profound effect on us. In this post, I will outline my research on this topic and attempts to create a microbiome world model that understands the complexities of the abstract interactions between microbes across as many contexts as possible. This should allow the model to build a general understanding of microbial interactions.

What is a Microbiome?

A microbiome is a community of microorganisms—including bacteria, viruses, fungi, and other microbes—that inhabit a specific environment. In humans, the microbiome primarily resides in the gut but can also be found on the skin, in the mouth, and in other areas of the body. Basically, anywhere there is a surface or crevice, there is a microbiome. The more protected it is from the environment, the more embedded and stable that microbiome will be.

They interact with each other in complex ways and form communities with emergent properties. For example, in your mouth, bacteria work together to form a biofilm that protects them from the environment and provides a safe space for consuming nutrients (and causing cavities). Biofilm example

What can we aim to achieve?

If we can build a model that understands the interactions between microbes, we can build targeted pro-biotics Herculean enough to reverse degenerative stable states to:

Increase our IQ
Get rid of diseases like IBS, eczema, auto-immune etc,
Becomes marathon runners
Eat pizza without consequence
Become drunk 24/7 (although technically already possible)

Current methods struggle with this because simply sending random “good” microbes into the gut is not enough. We need a targeted approach that takes into account the existing microbial community, displacing only the organisms we don't want and ensuring the ones we do want can plausibly survive there.

Data

Before we can start modelling we need some data. For this project I have been using data from MicrobeAtlas (paper).

MicrobeAtlas is the largest microbiome database to date. It contains millions of samples from various environments, including human, animal, and environmental sources. For each sample, we also have extensive textual metadata we can use.

A sample consists of a list of present microbes and their abundances (how many there were), but we will treat them only as present/absent here.

Representation

An important question to ask is: how do we represent microbes in a way that a model can understand? This is a tricky question because microbes are not like words in a sentence; they do not have a fixed order or structure. Even defining a microbe can be somewhat tricky (ultimately, they're just bags of DNA).

Here I choose to use DNA language models to embed 16S rRNA sequences (a widely accepted, uniquely identifying DNA marker) into a common space. This way, we can use a parameter-efficient transformer architecture without requiring learned embeddings.

From MicrobeAtlas, we now have a vocabulary size of ~100k unique DNA pieces/microbes. Plotting them in DNA space now looks like this: 16S input space The different clusters represent Archea (e.g. Methanobrevibacter), Bacteria (e.g. lactobacillus), and Eukaryota (e.g. human, amoeba, etc.).

Model Objective

The model was trained using a cross-entropy (CE) loss objective to predict the presence or absence of microbes in a sample given the other microbes present. In practice, the model sees a putative sample comprising the DNA embeddings of the present microbes and textual metadata (embedded with an LLM). For each embedding, it then outputs a confidence score indicating how likely that microbe truly belongs to the sample.

Thus, to make this difficult, I used a bernouli noising strategy. Adding embeddings to confuse the model and as such make it learn the true relationships.

Why no Auto-Regression?

I opted not to use auto-regression and instead use this quasi-contrastive approach because microbes in a sample are not ordinal. It doesn't make sense to impose an order on them when they more naturally exist as unordered sets. This means there is no positional embedding whatsoever; the transformer operation is entirely permutation-invariant.

This method also avoids projecting back to vocabulary space, which would be very expensive given the size of the vocabulary. It also means the model is more generalised beause it has an open vocabulary (any DNA embedding can go in, even if here we stick to a predefined set from MicrobeAtlas).

Training

The model was trained on ~2M samples from MicrobeAtlas. I used a batch size of 32 and trained for 1 epoch. The model was trained using AdamW with a learning rate of 1e-4. The model was trained on a single V100 for ~7 hours. I trained two versions of the model. Large:800k params, and small:50k params.

Results/Eval

Now for the interesting part. What did the model learn? First lets look at how the models representation changes as we pass a noised sample through its layers. Below is a UMAP for first the small then the large model. The first subplot on both represents the DNA space.

Small Model UMAP comparison

Large Model UMAP comparison

In these figures the triangles represent the noise microbes and the colours represents the final logits they were assigned (remember lower logit == less likely to be here). We can clearly see in both cases, the triangles are relativly evenly distrubuted in the DNA space, but as the model processes them, they are slowly removed.

Wow, so we can clearly see that the larger model has a much better go at figuring out which microbes were original and which were added. You'll also notice if you look closely that some triangles are not given lower logits whilst some circles are. This is likely because even the ground truth is somewhat noisey and stochastic.

Gingivitis Eval

The next question I have is, can we use this model, that was never trained on any dynamics or concepts of time, to predict which microbes will be present at another time step and which will not? I call this test: 'drop-out or colonise?'

To do this we will be using a held out gingivits dataset. This dataset consists of samples from patients with induced gingivitis over the course of a month. There is no finetuning or classification head. We simply take the model and run 0-shot.

Drop-Out Test

Wow again! Not only is the larger model (latter figures) better at this task, but both models are doing far better than random chance! This is very exciting because it means the model has learned some underlying representation of how microbes interact with each other. This is despite never being trained on any temporal data. In otherwords: given a set of microbes, the model can predict which ones are likely to drop out in the future/past. Hence the model has learned: given this set of microbes: these ones are stable and these ones are not. Note the model is doing this over diverse conditions in this dataset (first patients have 'healthy' oral biomes, then ginigivitis, then 'healthy' again).

Colonise Test

Another good result! Not quite as good but this test is much harder. Here the model must predict which microbes will newly appear in the next time step. I tested this by taking microbes that were present in other samples but not at a given time point t. Then I would place the microbes into t and see if the model predicting higher logits on these putative microbes corresponded with higher liklihood of that microbe appearing in the future.

What does this mean?

Both of these results will be super useful for designing targeted pro-biotics as they mean we can cheaply run millions of experiments in silico and optimise the logit space for the microbes we want to drop out/colonise. It should also be noted that the dataset itself is inherently noisey and the nature of microbiomes is probably stochastic, so perfect predictions are not to be expected.

Infants

Lets change tack here and see whether the final embeddings produced by the model might also be good as embeddings/representations of samples. To do this, I will use a dataset of samples taken from infant guts at different ages, some born by C-section and some vaginally. The goal here is to see whether the models learned embeddings can predict age and birth mode.

Given the implications of the microbiome for health outcomes, if we can tell them apart, that means the model easily finds patterns that differentiate the states. This would justify the recently adopted, somewhat medieval, practice of vaginal seeding for C-section born infants.

First lets take a look at a PCA of the embeddings produced by our models for the samples, this will show us whether the model has learned any interesting structure without any fine-tuning or classification heads.

Small Model PCA of embeddings

Large Model 3D PCA of embeddings

The first PCA comes from the small model and the second from the large model. I used 2d for the large model because it was easier to understand. In both cases we can see that the model has learned some very obvious structure. It clearly inherently differentiates between the different age groups and shows an almost monotonic time progression between them. Again this data was totally unseen and this is 0-shot (no finetuning or classification head).

Now lets see if we can use these embeddings to classify whether an infant was born by C-section or vaginally and what age they are. I did a 5x cross validation with a logistic regression head on top of the frozen embeddings.

Again we see that the larger model is better, but both models are doing far better than random chance. For comparison I show the results of a comparable method from MGM (paper). Their model has ~2M parameters and they fine-tuned on the infants data (unlike my frozen model approach).

What does this mean?

This means that the models embeddings are rich, and in this embedding space, there is a clear structure that differentiates the different age groups and birth modes. In fact we can see that even at 5Ys of age the model embeddings differentiate between C-section and vaginally born infants. Clearly, the birth mode is having long lasting effects on the microbiome.

Linking back to our discussion on health, IQ, cardio, etc, who knows what the long term consequences of this difference are? Given our results in colonization, maybe we could make a per-indivudal pro-biotic to help C-section born infants? (Although the vaginal seeding practice is probably enough).

Conclusion

In this post I have outlined my work on building a microbiome world model. The model was trained on millions of samples from MicrobeAtlas and was able to learn interesting representations of microbial interactions. These representations were then shown to be useful for predicting which microbes would drop out or colonise in a held out gingivitis dataset. Furthermore, the final embeddings produced by the model were rich enough to be useful for classifying whether an infant was born by C-section or vaginally and their age.

One day, this model could be used to design targeted and individual pro-biotics to help us all become super-humans.

There is still much work to be done in this field, but I believe that this approach has great potential. I think the use of the DNA embeddings was crucial to the success of this model.

Code