Screening of Biosecurity Features in Metagenomic Data with Evo 2 Probes

Technical Breakdown§

This paper evaluates the linear decodability of biosecurity-relevant signals from the frozen representations of the Evo 2 genomic foundation model. The core idea is to train minimal probes (linear classifier or single-head attention) on the layer-26 token embeddings without any fine-tuning of the base model.

Methodology§

Embedding Extraction: For a given DNA sequence of length $L$, we extract the hidden states at layer 26 of Evo 2, producing a sequence of vectors $\mathbf{h}_1, \ldots, \mathbf{h}_L \in \mathbb{R}^d$ ($d=512$ for the 1B-param model). We consider two aggregation strategies:

• Mean Pooling: $\bar{\mathbf{h}} = \frac{1}{L} \sum_{i=1}^L \mathbf{h}_i$
• Attention Pooling: A learned single-head attention that computes $\mathbf{q} = \text{Linear}(\bar{\mathbf{h}})$, then $\alpha_i = \text{softmax}(\mathbf{q}^\top \mathbf{h}_i)$ and $\mathbf{z} = \sum_i \alpha_i \mathbf{h}_i$.

Probe Training:

• Linear Probe: A logistic regression (or linear layer followed by sigmoid) on the pooled embedding: $\hat{y} = \sigma(\mathbf{w}^\top \bar{\mathbf{h}} + b)$.
• Attention Probe: The pooled representation $\mathbf{z}$ is passed through a linear classifier with cross-entropy loss.

Training minimizes binary cross-entropy for each region-level label (e.g., presence/absence of AMR gene family). The probes are trained on a dataset of metagenomic contigs with known functional annotations.

Evaluation: Held-out test sets from the same metagenomic distribution. Performance measured by region-level ROC-AUC. The best linear probe (mean-pool) achieves 0.888, and the attention probe reaches 0.977 for AMR detection. For virulence, region-level AUC is 0.833.

Read-Level Transfer: The same probe, when applied directly to features extracted from simulated short reads (without retraining), yields a read-level ROC-AUC of 0.898 (mean-pool), comparable to full-region results.

Sparse Autoencoder Analysis: To interpret the model, a sparse autoencoder is trained on the layer-26 activations. However, the recovered features are less consistent than the supervised probes.

Code Snippet (PyTorch-style Linear Probe)§

import torch.nn as nn

class LinearProbe(nn.Module):
    def __init__(self, d_model=512, num_classes=1):
        super().__init__()
        self.fc = nn.Linear(d_model, num_classes)
        
    def forward(self, embeddings):
        # embeddings: (B, L, d_model) -> (B, d_model)
        pooled = embeddings.mean(dim=1)  # mean pooling
        logits = self.fc(pooled).squeeze(-1)
        return logits

Key Equations§

• Mean-pooled representation: $\bar{\mathbf{h}} = \frac{1}{L} \sum_{i=1}^L \mathbf{h}_i$
• Attention-pooled: $\alpha_i = \text{softmax}\left( \mathbf{q}^\top \mathbf{h}_i \right)$, $\mathbf{z} = \sum_i \alpha_i \mathbf{h}_i$, where $\mathbf{q} = \mathbf{W}_q \bar{\mathbf{h}}$
• Binary cross-entropy loss: $\mathcal{L} = -\frac{1}{N} \sum_{i=1}^N [y_i \log(\hat{y}_i) + (1-y_i) \log(1-\hat{y}_i)]$

Conclusions§

Using frozen Evo 2 embeddings with lightweight probes provides a fast, inexpensive first-pass detection layer for metagenomic biosurveillance, achieving strong AMR detection without fine-tuning or assembly. The approach is limited by weaker virulence signal and the need for full-length sequences for optimal performance.