AntigenLM: Structure-Aware DNA Language Modeling for Influenza

Technical Summary: AntigenLM: Structure-Aware DNA Language Modeling for Influenza

Published at: ICLR 2026

1. Problem Statement

Predicting the evolution of influenza viruses is a critical public health challenge, as rapid antigenic drift necessitates frequent vaccine updates. Current forecasting methods generally fall into three categories:

• Classical Phylogenetic Models: Rely on tree dynamics and lineage ranking (e.g., LBI), which often fail to account for genome-wide coordinated evolution.
• Site-wise Deep Learning Models: Predict mutation fitness at individual sites (e.g., beth-1) but treat mutations as independent events, ignoring the coordinated interactions between different genomic segments.
• General-Purpose Foundation Models: DNA/Protein models (e.g., DNABERT, ESM) capture local patterns but are trained on heterogeneous multi-species data, often failing to capture the specific, higher-order structural constraints and co-evolutionary dependencies (e.g., RNA-RNA interactions or polymerase-antigen compatibility) unique to influenza.

Furthermore, protein-only models are "blind" to nucleotide-level evolutionary drivers such as synonymous mutations, codon usage, and non-coding regulatory elements, which are vital for viral fitness.

2. Methodology

The authors propose AntigenLM, a generative DNA language model designed to capture the full genomic context of Influenza A.

• Architecture: A decoder-only Transformer (GPT-2 style) with 6 layers, 384 hidden dimensions, and 6 attention heads. Crucially, it supports a long context window of up to 13,000 nucleotides, allowing for the modeling of the entire influenza genome (~13kb) without truncation or the need for k-mer tokenization.
• Structure-Aware Pretraining: Unlike general models, AntigenLM is pretrained on the full eight segments of the influenza genome (PB2, PB1, PA, HA, NP, NA, MP, NS) concatenated in a fixed order. This preserves the biological integrity of functional units and enables the model to learn long-range co-evolutionary dependencies across segments.
• Multi-Task Learning: The model utilizes a dual-head design:
  1. Language Modeling (LM) Head: An autoregressive head for predicting the next nucleotide, used for sequence forecasting.
  2. Classification Head: A head that projects hidden states into subtype logits for sequence-level discrimination.
• Fine-Tuning Strategy:
  • Forecasting: The model is fine-tuned on temporal windows of HA and NA sequences. It takes three historical blocks (past months) as a prompt to autoregressively generate the "future" block (next month or next season).
  • Classification: It uses sentinel tokens (e.g., <HA>, <NA>) to delimit regions, guiding the model to identify subtypes.

3. Key Contributions

1. Functional-Unit-Aware Pretraining: Demonstrates that preserving segment order and genomic structure is superior to segment-wise or fragmented pretraining.
2. Nucleotide-Level Generative Framework: Provides a model that captures both protein-coding changes and non-coding evolutionary signals.
3. Integrated Forecasting & Classification: A unified framework that achieves state-of-the-art performance in both predicting future sequences and identifying viral subtypes.

4. Results

• Short-term Forecasting: AntigenLM achieved extremely low amino acid (AA) mismatches in HA (<1%) and NA (<0.5%) when predicting next-month sequences, significantly outperforming general-purpose models (HyenaDNA, ProtGPT2) and ablation variants.
• Long-term (Next-Season) Forecasting: In testing against post-2022 Japanese data, AntigenLM outperformed the current WHO selection methods and the state-of-the-art beth-1 model. It reduced AA mismatches by over 70% compared to the current system and 50% compared to beth-1.
• Generalization:
  • Cross-Subtype: The model successfully generalized to the minor subtype H7N9, despite it representing less than 0.3% of the fine-tuning data.
  • Geographic: When trained only on Europe and Asia, the model maintained high accuracy on unseen U.S. data, correctly predicting clade transitions in 90% of cases.
• Ablation Insights: The study proved that "Incomplete-genome" (fragmented) and "Segment-wise" (isolated) pretraining significantly degrade performance, confirming that whole-genome context is essential for capturing evolutionary constraints.

5. Significance

AntigenLM represents a shift from "sequence-only" modeling to "structure-aware" biological modeling. By incorporating the physical and functional organization of the genome into the Transformer architecture, the authors provide a blueprint for building more effective foundation models for rapidly evolving pathogens. This has direct implications for vaccine strain selection, real-time genomic surveillance, and predictive genomics in public health.