basilboy/microbiome-model

Microbiome Transformer (Set-Based OTU Stability Model)

This repository provides Transformer checkpoints for microbiome set modeling using SSU rRNA OTU embeddings (ProkBERT-derived vectors) and optional text metadata embeddings.

Please see https://github.com/the-puzzler/microbiome-model for more information and relevant code.

Model summary

• Architecture: MicrobiomeTransformer (see model.py)
• Input type 1 (DNA/OTU): 384-d embeddings
• Input type 2 (text metadata): 1536-d embeddings
• Core behavior: permutation-invariant set encoding via Transformer encoder (no positional encodings)
• Output: per-token scalar logits (used as stability scores)

Available checkpoints

Filename	Size variant	Metadata variant	d_model	num_layers	dim_feedforward	nhead
small-notext.pt	small	DNA-only	20	3	80	5
small-text.pt	small	DNA + text	20	3	80	5
large-notext.pt	large	DNA-only	100	5	400	5
large-text.pt	large	DNA + text	100	5	400	5

Shared dimensions:

• OTU_EMB = 384
• TXT_EMB = 1536
• DROPOUT = 0.1

Input expectations

1. Build a set of OTU embeddings (ProkBERT vectors) per sample.
2. Optionally build a set of text embeddings (metadata) per sample for text-enabled variants.
3. Feed both sets as:
   • embeddings_type1: shape (B, N_otu, 384)
   • embeddings_type2: shape (B, N_txt, 1536)
   • mask: shape (B, N_otu + N_txt) with valid positions as True
   • type_indicators: shape (B, N_otu + N_txt) (0 for OTU tokens, 1 for text tokens)

Minimal loading example

import torch
from model import MicrobiomeTransformer

ckpt_path = "large-notext.pt"  # or small-notext.pt / small-text.pt / large-text.pt
checkpoint = torch.load(ckpt_path, map_location="cpu")
state_dict = checkpoint.get("model_state_dict", checkpoint)

is_small = "small" in ckpt_path
model = MicrobiomeTransformer(
    input_dim_type1=384,
    input_dim_type2=1536,
    d_model=20 if is_small else 100,
    nhead=5,
    num_layers=3 if is_small else 5,
    dim_feedforward=80 if is_small else 400,
    dropout=0.1,
)
model.load_state_dict(state_dict, strict=False)
model.eval()

Intended use

• Microbiome representation learning from OTU sets
• Stability-style scoring of community members
• Downstream analyses such as dropout/colonization prediction and rollout trajectory experiments

Limitations

• This is a research model and not a clinical diagnostic tool.
• Outputs depend strongly on upstream OTU mapping, embedding resolution, and cohort preprocessing.
• Text-enabled checkpoints expect compatible metadata embedding pipelines.