Predicting sub-population specific viral evolution
Authors: Wenxian Shi
Menghua Wu
Regina Barzilay
Venue: Accepted by TMLR ()

Abstract: Forecasting the change in the distribution of viral variants is crucial for therapeutic design and disease surveillance. This task poses significant modeling challenges due to the sharp differences in virus distributions across sub-populations (e.g., countries) and their dynamic interactions. Existing machine learning approaches that model the variant distribution as a whole are incapable of making location-specific predictions and ignore transmissions that shape the viral landscape. In this paper, we propose a sub-population specific protein evolution model, which predicts the time-resolved distributions of viral proteins in different locations. The algorithm explicitly models the transmission rates between sub-populations and learns their interdependence from data. The change in protein distributions across all sub-populations is defined through a linear ordinary differential equation (ODE) parametrized by transmission rates. Solving this ODE yields the likelihood of a given protein occurring in particular sub-populations. Multi-year evaluation on both SARS-CoV-2 and influenza A/H3N2 demonstrates that our model outperforms baselines in accurately predicting distributions of viral proteins across continents and countries. We also find that the transmission rates learned from data are consistent with the transmission pathways discovered by retrospective phylogenetic analysis.

Published in Transactions on Machine Learning Research (03/2025)
Predicting sub-population specific viral evolution
Wenxian Shi
wash edu
Department of Computer Science Massachusetts Institute of Technology
Menghua Wu
rww.@mit.cmmmm edu
Deportment of Computer Science
Massachuselts Institute of Technology
Regina Barzilay
regina Www.al.mm... edu
Deportment of Computer Science
Massachusetls Institute of Technology
Reviewed OpenReview: https://openrevve/ net/ forum? id= Eos
Abstract
Forecasting the change in the distribution of viral variants is crucial for therapeutic de- sign and disease surveillance. This task poses significant modeling challenges due to the sharp differences in virus distributions across sub-populations (e.g., countries) and their dynamic interactions Existing machine learning approaches that model the variant distri- bution as a whole are incapable of making sub-population specific predictions and ignore transmissions that shape the viral landscape. In this paper, we propose sub-population specific protein evolution model, which predicts the time-resolved distributions of viral pro- teins in different sub- nub-populations. The algorithm explicitly models the transmission rates between sub-populations and learns their interdependence from data. The change in pro- tein distributions across all sub-populations is defined through a linear ordinary differential equation (ODE) parametrized by transmission rates. Solving this ODE yields the likelihood of a given protein occurring in particular sub-populations. Multi-year evaluation on both SARS-CoV-2 and influenza A/H3N2 demonstrates that our model outperforms baselines in accurately predicting distributions of viral proteins across continents and countries. We also find that the transmission rates learned from data consistent with the transmis- sion pathways discovered by retrospective phylogenetic analysis. The code is available at httpps///g Seigtth...... com www/wwshh21// http://wssh2133/xxxeeeee///e httpa/////e/reemmammaaaaaanniiii
1 Introduction
Modeling virus evolution is important for vaccine development and disease monitoring. Specifically, the goal is to predict the change in the distribution of viral variants over time. Current approaches aggregate data across all locations of virus collection, resulting in models that represent. monolithic, global distributions of each virus. However, this global landscape is composed of local sub-communities of variants with distinct but interdependent distributions. For example, in January 2021, the most common strain of SARS-CoV. 2 was different in Europe (B.1.1.7) and Asia (B.1.1.214), but two months later, the former had spread rapidly across Asia (Fig. D.. Faithfully modeling these dynamics is essential for understanding how the viral landscape evolves locally and designing vaccines effective for each community, especially those with less access to public health resources and are poorly represented by the data.
In this paper, we aim predict the likelihood of viral variant occurring at particular time in specific sub-populations, such as the geographical location where the viral sample was collected. Viral variants are represented by the amino-acid sequence of the primary antigenic protein responsible for infection. A

---

Published in Transactions on Machine Learning Research (03/2025)
Pr (x 0
Europe, ofhers
ACre Klkas) Ale, Alxas)
Transmission rate matrices
AlG. (knaa)
12 ITA
A(xa Mixes)
guax
01177
amino arids
50 4%
$1.1% 156
8.177
mu of sub-populations
0.1.7
2021-01
2921-03
others
Transformer
5.5%
814173
13 9%
Protein:
M
M
34.0%
34.3%
Sub-population: L.
63:2244
UNaSS good
Date:
Figure 1: SARS-COV-2 clade distributions differ by geographical location and change interdependently over time.
Figure 2: The transmission model: auto-regressive, time-resolved generative model, parametrized in terms of transmission rate matrices.
straightforward approach for modeling location-specific evolution is to separately analyze data for each loca- tion, assuming independent evolution. Inherently, this approach leads to data fragmentation, exacerbating the data sparsity problem. Moreover, this approach fails to capture transitions between sub-communtiies that shape local distributions, especially for highly transmissible discases. Traditional epidemiology liter- ature has sludied transmission dynamics extensively (Sattenspiel & Dielz) 1995
Goel & Sharma 2020:
Balcan et al. 2010 2009 Colizza et al. 2006), but the hand-crafted rules they employ fall short of capturing
the actual complexity of these interactions.
We propose a novel approach to model evolving protein distributions for sub- aob-populations by explicitly cap- turing the transmissions between them. Specifically, the change in the occurrence of a protein sequence over time across all sub-populations is defined by an ordinary differential equation (ODE), parametrized by transmission rate matrix related to that sequence. Solving this ODE yields the time-resolved probabil- ity of the given protein in different sub-populations, which is determined by the transmission rate matrix and boundary conditions. We leverage language models (Radford et al.] 2019) to model the transmission rale matrix and boundary conditions. A practical challenge is that the nunber of transmission rates scales quadratically with the number of sub-populations, rendering eigen-calculations inefficient for fine-grained sub-populations (o.g. countrics). To address this issuc, we also propose a hicrarchical variation of our model, which leverages the inherent structures that relate sub-populations, e.g. countries can be grouped by their geographical proximity. Specifically, we re-parameterize the transmission rate matrix in terms of intra- group and inter-group transmissions. Both versions of our method are advantageous for sub-populations with limited data, which can benefit from knowledge transferred via the transmission rates with other sub- populations. Moreover, the association defined by the transmission rates can be learned from the occurrences of variants in the datasct, climinating the need for hand-crafted knowledge.
We evaluate our approach on two viruses, influenza A/ A/H3N2 and SARS-CoVa across multiple years, in over thirty countries and all six continents. The evaluation shows that our model outperforms thate-of-thh-artt basclines in predicting future distributions of viral proteins. For SARS-CoV-2 the top 500 protein sequences generaled by our model can cover 93.4% of the circulating sequences reported in various countries. Moreover, the transmission rates learned from the data are well aligned with the transmission pathways of viral variants discovered independently by phylogenetic analysis. Finally, our experiments focus primarily on goographical location due to data abundance, but our framework can be easily applied to various types of aub-populations, including individuals of varying ages and vaccination histories, as well as their combinations.
2 Method
Let (,',...,'') E V denote the amino-acid sequence of the viral protein responsible for infections. For example, we model the Hemagglutinin (HA) protein for influenza and the Receptor Binding Domain (RBD)
2

---

Published in Transactions on Machine Learning Research (03/2025)
of Spike protein for SARS-CoV-2. V is the set of amino acids, and I is the length of the protein sequence. Our goal is to model (r; i), the likelihood of protein sequence isolated at time i in sub-population E S. S is the set of m sub-populations, iee, |S| m. Throughout this paper, we use location, defined as the geographical region where the viral sample was collected, interchangeably with sub-population.
2.1 Transmission model
The transmission model parameterizes (r;:;) through the transmission rate x between goographical locations. The model outputs Nx(x), the un-normalized probability (occurrence) of at time m lo- cations: (x) = (a;; i), (e(3:)) We assume that the derivative of occurrence over time is a linear transformation of (z),,, following Kermack & McKendrick (1927); Obermeyer et al. (2022):
dN((2) (x)))) dt
(1)
in which A(x; 0) E RTxx is called the Iransmission rale matrix of protein x. Intuitively, [A(r; O)ls measures the number of people in location infected by one person from location j during dt. In this work, the transmission rate matrix is output by a neural network parameterized by that takes as input. For concision, we drop 8 in the following equations wherever clear. The ordinary differential equation in Eq. [ has closed-form solution,
nx(x;;)= = Σ (xs(=))xp(AA))) c((); (())=
s(+)=-u(s) U(()) (2;
(2)
in which Ax(x)) is the j-th eigenvalue of matrix A(r), and wi((r) is the i-th dimension of the j-th eigenvector of A(z). (x) is determined by the boundary conditions No(x). U-1 is the inverse of eigenvector matrix U. The derivation can be found in Appendix A.1
Intuitively, (x(x)) describes the "fitness" of protein x in a "latent" location j. Larger )y(x)) indicates that strain will reproduce faster in the latent location j. The eigenvectors combine latent locations with different growth rates A, and determine the distributions of x in real locations.
To calculate the probability of the sequence æ in location i at time I, we need to normalize the occurrence ne(2; i), which is intractable for massive space of sequences. Inspired by autoregressive language models, instead of modeling the occurrence of all sequences, model the occurrence of the s-th amino acid x, given the prefix amino acids Nes:
m
=
Σ ug exp(A)(x, os (x.)
(3)
The conditional probability of s-th amino-acid calculated accordingly:
Me(Fs, pr( Exjev (1(2 ic))
(4)
The probability of the whole sequence can be factorized the product of the conditional probabilities of each amino acid. As summarized in Fig. 2 given a sequence sampled from location at time t, for each position 8 in the amino acid sequence, the model takes the previous & - amino acids Ies as input and outputs [V] transmission rate matrices of size m X m. The probabilities of the next amino acid in cach location (a m [V] matrix) are calculated from these transmission rate matrices. For a protein sequence with residues (amino acids), the model first outputs L x [V] matrices of size x m, and then outputs a probability tensor of size L x [V] m.
The parameters 0 are trained by the maximum likelihood objective based on tuples (x, sampled from the training set with empirical distribution Purain (z, i, i,t). The training loss is
((0) -E (x...) log pz (x: i) -E EEpaati (n.,,tt log (x.:
(5)
3

---

Published in Transactions on Machine Learning Research (03/2025)
400
Block
A00
Group-
Group
diagonalize
d
d
dt
dt
(a) Hierarchical transmission model
(b) Block-wise diagonalization
X
(x
Group -oo- location
Mr(+)
X
Intra-group
Inter group
600
Group to group
Intra-group
Inter-group
(c) Inter-group approximation
Figure 3: Hicrarchical transmission model. (a) Instead of modeling transmission between all pairs of coun- tries, we model interactions within and across groups of locations (e.g. continents). (b) This idea can be realized by a block-wise diagonal Å. (c) Two strategies for learning transmissions within each group.
Our model achieves knowledge sharing based on the transmission rate matrix. When the model is trained on a sample from location i, the transmission rate matrix is updated through its eigenvectors and eigenvalues (which relate locations with each other). As a result, the probability that sequence in other locations is also updated. Thus, our model can take advantage of data from different locations learn the specific distribution for each location.
2.2 Hierarchical transmission model
Since cigendccomposition scales as O(m²) with the mumber of locations, applying our model directly to fine-grained locations such as countries becomes highly time-consuming (analysis in Section 4.4). Moreover, the intrinsic structure between sub-populations, such as their geographical proximity, might provide useful inductive biases for modeling their transmissions. Thus, we introduce hierarchical transmission model that leverages the hierarchical structure of locations to approximate the transmission rate matrix and simplify eigen-calculations. As illustrated in Fig. Ba partition locations into groups, e.g. countries to continents. In the hicrarchical transmission model, the occurrence in a specific location is calculated by considering transmissions from locations within the same group and from other groups.
We annotate the group of location as gi E G, where is the set of all groups (e.g. continents). We construct matrix (xx) mmm, which transforms the transition rate matrix A(x) into a block-wise diagonal matrix. Simplifying (wx) as W and A(x) A, multiply by W-1 both sides of Eq. to obtain
dw N.(z) W-! AWW- de
dN.(s)
(.))
(6)
de
in which A(x) (z) is block-wise diagonal matrix with |G| blocks, and Nz(x) (x) is an affine transformation of Nz(x).
Instead of parameterizing the original transmission rate matrix A(x) directly, we parameterize W(z; 0) and the block-wise transmission matrix A(r;6). As proved in Appendix A.2 this re-parameterization retains the same representation capacity as the original matrix, provided that A(x) is diagonalizable. Since A(x)
is block-wisc diagonal, Eq. can be further divided into |G] equations: d.f(()) A² (x). G. As illustrated Fig. 3b. (z E Xm,, is the transformed transmission rate matrix corresponding to group g with mn locations, and No (x) E "h is the transformed occurrence of sequence locations within group g. Calculating the closed-form solution of (z(x)) from A as shown in Eq. 2 is more efficient due to the smaller size of A9.
4

---

Published in Transactions on Machine Learning Research (03/2025)
As illustrated in Fig. 3c
(x(x) can be obtained from the back transformation: N((x) (x)) Thus, the occurrence of x in location is
me(2; = Σ W (1)))
E
Wofne(z;))+ E (x;
(7)
9:83
3:35
(x:4)
Wigne(z;))
(8)
Intra-group transmission term:
Inter-group transmission torm:
That is, we rewrite the occurrence of in the sum of contributions from within i's group (A), and from all other groups (B). Here, assume that Wü: 1 and Wi 0 if 9; = 9i and j * which is reasonable because the transformation of elements within groups is redundant during the block diagonalization. We propose two options to approximate W: group-lo-location (G2L) and group-lo-group (G2G). The intuition behind these approximations is that transmission between distant locations is often dominated by some major routines. For example, when traveling from an Asian country to the United States, the journey may involve a stop at a major port in Asia before reaching the destination (group-lo-locatioon) or it may involve travel between two major ports across continents before reaching the final destination (group-to-group). We provide the formal description of these approximations below.
Group-to-location (G2L). The first strategy approximates the inter- inter-group transformation weight Wy: from location j to location with Wig, from group Oj to location i (Fig. 3cb. The inter-group transmission term is approximated as
Ba2 Ba (x,i,1)=
Σ
Wig(z))
(9)
gEGV(9:>
where W(x) E Rmx[GG| is linear transformation from groups locations, parameterized by the neural network. (x; = Engn=gg āe(x;)) the sum of the transformed occurrences over locations in group g.
Group-to-group (G2G). The second strategy simplifies relations between two locations into a group-to- group term, followed by a term that "distributes" incoming transmissions among locations within a single group. Specifically, we approximate W. with a transformation Čz.9, from group 9j group Ge, and a transformation Dig, from group to constituent location E 9:- More specifically,
Ba i,t) Dr (x)
Σ
())(())
(10)
(
in which Ö( Φ) G the linear transformation from groups to groups, while ń. (. E Rl( is the transformation from group to locations belonging to that group.
In both cases, only the group-level ú. (x: is required, so for efficiency, we use an auxiliary model approximate (z;g)) instead of computing and summing over individual (x(x;)). Similar to Eq. 2
(x g) = Σ ügg (x) exp(Ag (x)t) 2g (x).
(11)
D'EG
In practice, sub-populations be clustered into groups based on their geographical coordinates. While these approxinations perform well in empirical studies, they may fail to fully capture the complexities of transmission dynamics when certain sub-populations exhibit significantly distinct transmission patterns. This issue can be mitigated by refining the group configurations to better separate such sub-populations.
Summary. The occurrence sequences x in location i from group 9: is modeled as either
(r; i, 0) (,( i, 0)
+ ni (z; + BE2L (x. i, t), or <- 7 i) + Beza i,
(12)
The probability Pr(x; 8) is calculated from ne(a; i, 8) using the same autoregressive factorization as Eq. Specifically, given prefix Tees input:
5

---

Published in Transactions on Machine Learning Research (03/2025)
1. The model outputs the transmission rates A (Ee) within group 9r, and we compute occurrences nig (Ix, (esi in location i for each 2s-
2. The model outputs group occurrences ñ (z g) and inter-group transformation weights Wig (Fs, (or equivalent Ö and D) for all other groups # 9t-
The Reg is computed combining the intra- and inter- group contributions, as illustrated in Fq. 12 We obtain pe(salars) by normalizing ni ni( Fes) over all possible values of Is-
The training loss for the hierarchical transmission model is
(()) Perrii (x.iit) log.:: + ACgrou i,t;
(13)
in which À is hyper-paramete and the Egroup is the regression loss for the auxiliary model that predicts the total occurrence within groups:
Cgroup (eroup(x,itt;) (x,1,t;0)=
=
nit (x: ソン
Σ
ne( 9r,
(14)
3 Related work
Differential equations for modeling epidemiology: Classic
(SIR)
models (Kermack & McKendrick] 1927 Beckley et al. 2013 Harko et al. 2014) split. the population into
three sectors, susceptible individuals (S), infectious individuals (I), and recovered individuals (R). Three
ODEs describe how the number of infections in each sector changes over time. To model the spatial spread
of disease, a series of works (Sattenspiel & Dielz 1995 Goel & Sharma 2020 Balcan et al. 2010 2009 Col-
izza et al. 2006) introduce mobility operators between different geographical regions into SIR models using transportation data. However, those models focus on the number of infectious for a specific disease. They do not have the capacity to model how different variants evolve over time, which crucial for vaceine selection and pandemic surveillance. In contrast, model leverages neural network to estimate the infection ratio (i.c. probability) of various protein sequences. This flexibility allows our model to predict the probability of any protein variant.
Predicting fitness for viral variants: One class of work predicts the fitness of single amino acid mu- tations and/or /or predefined clades (Luksza & Lässig, 2014 Neher et al., 2014 Obermeyer et al.! 2022 Maher et al. 2022), but they do not consider fitness based on the whole protein sequence (Gong et al.
2013 and
cannot be applied to unseen single amino acid mutations. Another line of work uses protein language models
to predict the fitness of protein sequences (Hic et al. 2021 Frazer al.) 2021 Thadani et al. 2023 whole genomes ZZvaagn et al. 2023). However, they do not explicitly model epidemiological information, which results in a static landscape incapable of capturing changes in the distribution of variants over time. Shi et al. (2023) predicts the distributional changes in variants based on their protein sequences but does not distinguish aub-populations. To the best our knowledge, our model is the first machine learning model that predicts the temporal distributions of viral proteins in different sub-populations.
Stochastic Differential Equations (SDEs) and ODEs defined by neural networks: A series of studies propose generative models based Stochastic Differential Equations (SDEs) (Song et al. 2021 Lipman et al. 2023 Campbell et al.) 2022). These models introduce an abstract time ranging from 0 to T and define a forward (diffusion) process that transforms the data distribution (at time T) to noisy distribution (at time 0). However, these models are unsuitable for our task, as they inherently model a static distribution at time T, and the diffusion process does not align with the observed time-reversed evolution process. Other work proposes Neural ODE for representation lcarning (Chen et al., 2018), which is not directly applicable our task since the variable evolving in our model is the distribution of discrete sequences, rather than continuous representations.

---

Published in Transactions on Machine Learning Research (03/2025)
4 Experiments
4.1 Settings
Data processing: In this paper, we focus on two viruses, influenza A/H3N2 (abbreviated as FLU) and SARS-CoV-2 (abbreviated as Cov). We predict the change in distributions of Hemagglutinin protein (HA) protein for FLU and the receptor-binding domain (RBD) of Spike protein for Cov. We obtain the amino acid sequences of these proteins from GISAID (Shu & McCaulcy 2017). Samples collected from human hosts with collection date and collection location information are used in the training and lesting. We regard one month (for Cov) or two months (for FLU) as one unit of time.
Evaluation: We evaluated models at two levels of geographic gramularity - continent and country levels. In multi-year retrospective studies, we employed training and testing splits based on realistic vaccine devel- opment processes for both FLU and Cov. For influenza, we emulate the annual recommendation schedule for the northern hemisphere. Since egg-based vaccines require lead times of up to 6 months, we train our models on data collected before February of each year, and evaluate models on the sequences collected from October to March of the next year (winter season), following (Shi et al. (2023). To exclude the influence of the COVID-19 pandemic and ensure adequate training data, we evaluated our continent-level model on four influenza seasons, from the 2015 to 2018 winter seasons, and the country-level model three seasons from 2016 to 2018. In contrast, it takes ~3 months to produce mRNA vaccines, so evaluate on Cov in 3-month intervals. Specifically, we trained four models using data collected before four end-points: 2021-07, 2021-10, 2022-01, and 2022-04. For instance, a model trained on sequences collected before 2021-07-01 will be evaluated on sequences collected between 2021-10-01 and 2022-01-01.
We consider six continents (excluding Antarctica) and countrics with sufficient data. (16 for FLU, 32 for Cov) when training the model. The evaluations are conducted on lesting time and locations with sufficient data (10 countries, 5 continents, excluding South America, for FLU; 32 countries, 6 continents for Cov). More details about data pre-processing and statistics can be found in Appendix A.3
To measure how well our model and baselines predict future distributions of viral proteins, we used two metrics: negative log- liglliklihood (NLL) and reverse negative log-likelihood (reverse NLL) (Strude et al. 2022: Campbell et al., 2024).. The NLL is the average negative log-likelihood of observed sequences occurring at the testing time and locations estimated by our model and baselines. The reverse NLL is the average negative log-likelihood of sequences generated by our model and baselines for given testing times and locations, as estimated by oracle models. The oracle models are trained all available data in cach location independently, including sequences in the lesting periods. The lower NLL favors models with higher diversity, while the lower reverse NLL favors models with higher specificity. By adjusting the temperature of our models and baselincs, we can trade off between diversity and specificity, thereby obtaining the Parcto- frontier in the NLL and reverse NLL space. The Pareto-frontier closer to the lower NLL and reverse NLL corner indicates better model performance. We generated 1500 sequences by our model and baselines with temperatures 0.2, 0.4, 0.6, 0.8, and 1.0. For each model, we obtained the average Pareto frontier by averaging the NLLs and reverse NLLs across testing locations and times with sufficient data. To alleviate the variance caused by the oracle models, we calculated the reverse NLLS by three different oracle models, as illustrated in Appendix A.4
Metrics:
Inspired by Thadani et al. (2023), we also investigate the ability of our model to predict the prevalent proteins in the future. We generate 100, 300 and 500 sequences with the highest probabilities using beam scarch (Sutskever et al., 2014) and calculate their coverage, defined as the ground-truth frequencies of these sequences occurring in the next three months at different locations. In contrast to the focus of Thadani et al.] (2023) on single amino-acid substitutions, we focus on the coverage at the protein sequence level, considering exact matches between the generated sequences and the real-world circulating sequences. We only evaluate performance for Cov due to the sparsity of available data for FLU. We compare the median of coverage across testing times and locations.
Implementation: While the transmission rate matrix is not necessarily symmetric, assuming is a real symmetric matrix is beneficial for training stability and acceleration. Thus, in practice, we parameterize the transmission rate matrix Ag(x) as positive and real symmetric matrix.
7

---

Published in Transactions on Machine Learning Research (03/2025)
For continent-level transmission models, we use a 6-layer GPT-2 (Radford et al. 2019) to parameterize the transmission rate matrix Ag and another 6-layer GPT-2 to model the initial occurrence No(x;; 0). For hierarchical transmission models, we share the transformer layers for group occurrence, transmission rate within groups, and cross-groups. Adam optimizer with learning rates 1c-5 and 50-5 are used for FLU and Cov, and the models are trained for 80,000 steps for FLU and 30,000 for Cov with batch sizes 32 and 256 respectively. The learning rate linearly warmed up from 0 the specified value the first 10% epochs and then decays linearly to zero. More details about training and the hapaeepasammters can be found in Appendix A.5
We use the non-hierarchical transmission model for continent-level tasks and the hicrarchical version for country-level tasks. We use the hierarchical loss function but set the number of groups as one at the continent level, equivalent to the non-hieraree -on-hierall version adding a regularization term. We also present the results for variant of our transmission model, denoted as top-3 eig, where only the three largest eigenvalues are retained for denoising purposes. In hicrarchical transmission models, we perform agglomerative clustering based their geographical coordinates to form groups. We set two groups for FLU and three groups for Cov. The countries in each group can be found in Appendix Tab, 6 and Tab. 7 We present the results using both the G2G and G2L approximations.
Baselines: We compared our model against three classes of baseline models for both predicting future distributions and generating future prevalent sequences. All baselines were implemented using the same transformer architecture as our model.
No location information. Global model: a state-of-the-art evolution model (Shi et al. 2023 trained on all data without using the location information. This model can be regarded as a degenerate version of our model with only one location.
2. Baselines that leverage the location-specific data without directly incorporating the location into the model architecture. Finetune: fine-tuning the global model on local data. We fine-tune 8 epochs with a learning rate of 1e-5. LoRA (Hu et al. 2021 introducing and fine-tuning low-rank adaptations the global model.
Basclines that incorporate location information into the model architecture without considering the interactions between locations. Prepend Label: a modification of the global model, prepending special location token before the amino acid sequences. Both country and continent tokens are prepended for country-level. Add Embed: a modification of the global model, adding location embedding to the hidden state after each transformer layer. Parameter sharing: modification of the global model, sharing the transformer layers but adding different linear output layers for different locations.
We also compared our model with the following non-generative methods, which are applicable only for predicting future prevalent sequences through ranking cxisting sequences. Last: the most frequent sequences occurring within the last 12 months before the end-of-training time in each location. EVEscape (Thadani et al.) 2023): calculating the fitness for single amino acid substitutions based on a combination of protein language model, structural, and biological information. Note that EVEscape is neither temporal model nor sequence generalive model. Thus, utilized it as scoring metric to re-rank the sequences occurring in the last 12 months in each location. To ensure fair comparison, we re-trained their protein language models on our training set.
4.2 Main Results
As shown in Fig. our model achieves the best frontier in the average NLL and reverse NLL space for both FLU and Cov, when predicting protein distributions on both continent and country levels. Compared with the global model, models incorporating location information improve the specificity with lower reverse NLLs, such as Add Embod, Finetune, and Prepend models. However, compared with these baselines which inherently model the evolution in different locations independently, our transmission model achieves even better results in most cases by explicitly considering the transmission between locations. The NLL and
8

---

Published in Transactions on Machine Learning Research (03/2025)
A/H3N2 (continent level)
Cov-RBD (continent level)
1x
model
3x
3
reverse NLL
rewerse NLL
Prepend * Add Embed
2x
2 2x
Param_share
LoRA
Finetune
Transmission
4 10
10-1
2x 10
4x 101006x00 10
10-2 NLL
2 10
NLL
(a) Flu, continent-level
(b) Cov, continent-level
Flu A/H3N2 (country level)
Cov-RBD (country-level)
6 60
4x 10
model
reverse NLL
10-2
reverse NLL
-
Global
Prepend
30-2
Add. Embed
20-2
Param_share share
2
LoRA
Finetune
10-2
Transmlsslon(G2G)
Transmlssion(G2L)
6x
10
2x10
6x 10
2
NLL
NLL
(c) Flu, country-level
(d) Cov, country-level
Figure 4: Average negative log- likelihood (NLL) and reverse negative log- legliikli for FLU and Cov. Lower is better. Error bands represent the 95% confidence interval across different oracle models.
reverse NLL for each continent and country can be found in Appendix A.6 As shown in Appendix Fig. 16 method also achieves the best performance the worst sub sub-populations. The exact values of NLL and reverse NLL for Fig. are provided in Appendix Tab.
In addition, our model generates more prevalent sequences in the future for Cov. As shown in Tab. sequences generated by our model ("Top-3 eig" for continent level and "Hierarchy, G2L for country level) exhibit the highest coverage. The vanilla "Transmission" model also outperforms other baselines, alboit with slightly lower performance when generating the top 500 sequences. Although "Hierarchy G2G" per- forms slightly worse than "Hierarchy G2L,, it still surpasses other baselines, demonstrating the robustness of our hierarchical approach. The improvement in our country-level model is more pronounced than in the continent-level model, showing that considering the transmission is more critical for fine-grained and interdependent sub-populations. Notably, the top-500 sequences generated by our model can cover 90.9% of circulating viral proteins at the continent level and 93.4% at the country level, three months in advance.
4.3 Interpretation of transmission rate matrices evolution pathways
1.4
2018
0.77
1.4
1.2
1.2
1.2
0.91
OC
NA: 48%
11
0.81
1.2
1.3
1.3
0.85
AS
NA
0.75
1.2
1.3
1.2
0.85
AF
-0.8
0.84
EU
AF
A5
8U
NA
SA
SA
(a) Transmission rate (2018)
(b) Maximal spanning tree (2018)
(c) Phylogenetic (2018)
Figure 5: (a) Average transmission rate matrix among sequences collected during the 2018 winter flu season in clade 3C.2alb.2b. (b) The maximal spanning tree obtained from the rates matrices. (c) The phylogenetic tree obtained from the Nextstrain. Africa (AF) is not included due to insufficient data.
9

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Continent top-300
Country top-300
top-100
top-500
top-100
top-500
Last (12M)
0.819
0.874
0.906
0.841
0.882
0.898
EVEscape
0.118
0.408
0.664
0.461
0.842
0.887
Global
0.828
0.855
0.862
0.853
0.882
0.884
Prepend
0.832
0.867
0.891
0.849
0.879
0.891
Add Embed
0.828
0.875
0.889
0.853
0.879
0.894
Param share
0.816
0.870
0.889
0.818
0.870
0.884
LoRA
0.828
0.855
0.863
0.854
0.883
0.885
Finetune
0.839
0.875
0.882
0.855
0.882
0.889
Transmission
0.841
0.888
0.901
Top-3 eig
0.852
0.894
0.909
Hierarchy, G2G
0.859
0.898
0.923
Hierarchy, G2L
0.872
0.925
0.934
Table 1: The coverage (total frequencies of occurrence) of top-100, top-300 and top-500 sequences generated by models for Cov. We take the median over testing times and locations.
Flu AHH32 (continent level)
Cov-RBD (continent level)
AHHN2 (country level)
Cov-RBD (country level)
reverse NLL
Tranonission lein rog
Gabba
model
reverse NLL
4 4
Gibb
Taanmissiion in rug)
model
reverse NLL
36
Mrt Trbbed
model
20
G2uhandom
Fregeed
model
Transnissinn (no regl no reg, shuffe
3
Iranomisiin fro rog! no reg. sn-ffe
Gotc lng
Gatatondimm
30
no reg, tep-3 sip
no reg, 1op-3 di
020eduaeerigg mng
G2stdustertg
14x12-
10:2
4-10- 6x10
10
2x
4x
Exie-
10-1
2-10-
6 20
10-1
2x
10-1
2-10--
NLL
(a) Flu, continent-level
(b) Cov, continent-level
(c) Flu, country-level
(d) Cov, country-level
Figure 6: Ablation study. (a, b) Both the location information and the capacity to factorize the global distribution contribute to the model's performance. (c, d) Similarly, both the inductive biases based on geographic proximity and the ability to model higher-level groups are helpful.
We demonstrate that the transmission rates learned from our method align with the disease spread patterns revealed by phylogenetic analysis (Hadfield et al. 2018). which constructs an evolutionary tree from a set of genetic sequences Fig. 5a show the average transmission rate matrix of sequences collected in the 2018 winter FL.U season and annotated
the clade 3C.2a Ib. The transmission rates within continents are larger than those between continents, which implies the reproduction of viruses within continents is faster than across continents. Fig. 5b] shows the maximal spanning tree obtained from this transmission rate matrix. This suggests that the clade 3C.2alb.2b possibly started in Asia, transmitted to North America, and then to Europe and Oceania. This pathway is consistent with the phylogenetis analysis from Nextflu (Fig. 5ep. In Appendix A.6, Fig. 14 shows another case of FLU 3C.2a2 clade and Fig. 15 shows the case of Cov BA.2 clade. The transmission pathways implied from the transmission rate matrices are also aligned with the phylogenetic trees. It showcases that the transmission rate matrices learned by our method could be used to analyze the transmission pathways and patterns.
4.4 Ablation studies
Our ablation studies dissect the value of incorporating sub-populations as additional signals or through architectural changes (factorizing global distributions into mixtures); the runtime and performance trade-off of hierarchical modeling; and other design choices.
First, we demonstrate that the improvement of the transmission model comes from both the incorporation of location information and the advancement of probabilistic modeling. We randomly shuffled the location labels
The phylogenic obtained from the Nextstrain: attp:://eetteraaa..rrrr/aaaaaaas 2 Annolated by the nextclade.
10

---

Published in Transactions on Machine Learning Research (03/2025)
model G2G/clustering (max m2 10)
0.035
model
ma/1 toleess
0.030
G26/clustering (max 10)
reverse NILL
10²
G2G/random (mg 8)
0.025
G2G/random (m, G2G/random
305
244
255
G2G/random (mg 4)
0.020
(my
hierarchy
0.015
no hierarchy
10a
104
0.010
16
32
0.05
0.06
0.07
number sub-population
NLL
(a) Running time
(b) Performance
Figure 7: Training time (ms) per thousand tokens and performance for hierarchical transmission model and vanilla transmission model for FLU country-level. my is the number of countries in each group. Hierarchical modeling is beneficial for both computing efficiency and performance.
of samples in the training set. In the continent-level results depicted in Fig 6a and Fig 6b the transmission model trained on randomly shuffled data (annotated as shuffle) performs worse. This indicates that our model can effectively utilize location information for more accurate location-specific distribution modeling, However, the model trained on shuffled locations still outperforms the global model, suggesting that the mixture of components with different reproduction rates (represented by different eigenvalues) in Eq. also contributes to the improvement.
We also analyze the impact of grouping locations based meaningful inductive biases (e.g. geographical proximity), compared to hierarchical modeling in general. To eliminate the influence of group size, we kept the number of countries per group consistent with our hicrarchical model but randomly reassigned countries to different groups. As ratest in Fig. and Fig. 6d this random assignment (G2G/random and G2L/random) leads to worse performance, highlighting the advantages of introducing geographical proximity in grouping. However, even with randomly assigned groups, our hierarchical transmission model still outperforms the best baseline model.
Next, demonstrate that the hicrarchical design can improve efficiency significantly while achicving better performance. We compare the training time (milliseconds) per thousand tokens on NVIDIA RTX A6000 for our transmission model, with and without hierarchy. As shown in Fig. for FLU at the country level, our best hierarchical model, G2G/ /ee/ccluternnn.. is 2.6 times faster than the non-hicrarchical model while delivering better performance. We also compare the efficiency of hierarchical models with various group configurations. In the G2G/random models, we split the countries uniformly and randomly into groups, each containing my countries. As expected, the hierarchical model is faster as the number of countries in each group decreases. With my 4, the hicrarchical model is 8 times faster, achieving performance comparable to the non-hierarchical version. In addition, as the nunber of total countries increases, the training time for the non non-hierarchical model grows cubically. In contrast, the hierarchical model's efficiency is dominated by the number of countries in each group, preventing significant time increases with relatively small group sizes. In practice, our hierarchical model can be further accelerated (up to 8 times faster) by leveraging libraries like cuSOLVER NNVITA Corporation) 2023) which optimizes eigen-calculations for smaller matrices.
We studied two additional aspects of our model. We investigated how the number of principal components in the transmission rate matrix affects performance. As shown by the lines annotated top-3 eig in Fig 6a and 6b if we retain only the largest 3 eigenvalues and corresponding eigenvectors of the transmission rate matrix, the performance remains comparable for FLU but better for Cov. This indicates that, in certain cases, performance can be enhanced by de-noising the transmission rate matrix through relaining only the largest eigenvalues. Finally, we found that applying the regularization term Egroup in Eq. slightyl improves the performance for FLU, even without employing hicrarchical structures (a single group).
5 Conclusion
We propose transmission model predict the distribution of viral proteins in different sub- sub- populations. By learning the transmission rate between sub-populations based on protein sequences, our model achieves better
11

---

Published in Transactions on Machine Learning Research (03/2025)
performance in forecasting the distribution of Cov and FLU proteins in different locations. Furthermore, the transmission rate matrices learned from the data are aligned the transmission pathways discovered by phylogenetic analysis of given viral strains.
Our method has the following limitations. First, sub-population annotations are required for the training data, which limits the types of applicable sub-populations. Second, we assume that transmission rates remain constant over time, which may not hold if the environment changes unexpectedly. Third, the assumption of linear relationship between transmission rate and probability in the ODE may oversimplify the dynamics. As part of future work, plan to incorporate non-linear relationships to address this limitation, such leveraging a more general formulation of ODEs (Chen et al. 2018) to offer greater capacity.
Broader Impact Statement
This work provides a tool with significant potential support local public health efforts. The predictions provided by model can be used for discase surveillance and vaccine selection in specifie countries. How- ever, since the model is trained historical data, bias during data collection might cause inaccurate results. Decisions based solely on predictive analytics without considering real-world complexities might have unin- tended consequences, Furthermore, the synthesis and testing of viruses are strictly regulated and controlled by relovant authoritics, thereby minimizing potential misuse.
References
Duygu Balcan, Hao Hu, Bruno Goncalves, Paolo Bajardi, Chiara Poletto, Jose J Ramasco, Daniela Paolotti, Nicola Perra, Michele Tizzoni, Wouter Van den Broeck, et al. Seasonal transmission potential and activity peaks the new influenza a (hInl): a monte carlo likelihood analysis based on human mobility. BMC medicine, 7:1-12, 2009.
Duygu Balcan, Bruno Gonçalves, Hao Hu, José J Ramasco, Vittoria Colizza, and Alessandro Vespignani. Modeling the spatial spread of infectious diseases: The global epidemic and mobility computational model. Journal of compulational science, ((3):132-145) 2010.
Ross Beckley, Cametria Weatherspoon, Michael Alexander, Marissa Chandler, Anthony Johnson, and Ghan S Bhatt. Modeling epidemics with differential equations. Tennessee State University Internal Report, 2013.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucct. continuous time framework for discrete denoising modcls. Advances in Neural Information Processing Systems, 35:28266-28221 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXio preprinl arXiv:2402. 04997,,44999. 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Vittoria Colizza, Alain Barrat, Marc Barthélemy, and Alessandro Vespignani. The modeling of global epidemics: Stochastic dynamics and predictability. Bulletin of mathematical biology, 68:1893-1921, 2006.
Jonathan Frazer, Pascal Notin, Mafalda Dias, Aidan Gomez, Joseph K Min, Kelly P. Brock, Yarin Gal, and Debora S. Marks. Disease variant prediction with deep generative models of evolutionary data. Nature, 2021.
Rahul Goel and Rajesh Sharma. Mobility based sir model for pandemics-wite case study of covid-19. In 2020 TEEE/ACM International Conference on Advances in Social Networks Analysis and Mining (ASONAM), pp. 110-117. IEEE, 2020.
Lizhi Ian Gong, Mare A Suchard, and Jesse D Bloom. Stability-modiated epistasis constrains the evolution of an influenza, protein. Flife, 2:e00631, 2013.
12

---

Published in Transactions on Machine Learning Research (03/2025)
James Hadfield, Colin Megill, Sidney M Bell, John Huddleston, Barney Potter, Charlton Callender, Pavel Sagulenko, Trevor Bedford, and Richard A Neher. Nextstrain: real-time tracking of pathogen evolution. Bioinformatics, 34(23):4121 11211111111,,,, 2018.
Tiberiu Harko, Francisco SN Lobo, and MK3197716 Mak. Exact analytical solutions of the susceptible- rnfeccodd-recoverod (sir) epidemic model and of the sir model with equal death and birth rates. Applied Malhemalics and Compalation, 236:184 1918899944444499 2014.
Brian Hic, Ellen D Zhong, Bonnic Berger, and Bryan Bryson. Learning the language of viral evolution and escape. Science, 371 (6526))222 288, 2021.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, el al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Represen- tations, 2021.
William Ogilvy Kermack and Anderson G McKendrick. A contribution to the mathematical theory of epidemics. Proceedings of the royal sociely of london. Series A, Containing papers of mathematical and physical character, 115(772):700-77 721, 1927.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Fleventh International Conference on Learning Representations, 2023.
Marta Luksza and Michael Lässig. A predictive fitness model for influenza. Nature, 507(7490):57--1, 2014.
M Cyrus Maher, Istvan Bartha, Steven Weaver, Julia Di Iulio, Elena Ferri, Leah Soriaga, Florian A Lempp, Brian I, Hie, Bryan Bryson, Bonnie Berger, et al. Predicting the mutational drivers of future sars-cov-2 variants of concern. Science translational medicine, 14(663):eabk33445 2022.
Richard A Neher, Colin A Russell, and Boris I Shraiman. Predicting evolution from the shape of genealogical trees. Elife, 3:c03568, 2014.
NVIDIA Corporation. cuSOLVER Librory, 2023. URL http:///oo...iiidaacccooodd////d
Ver-
sion 11.7.
Fritz Obermeyer, Martin Jankowiak, Nikolaos Barkas, Stephen F Schaffner, Jesse Pyle, Lconid Yurkovet-
skiy, Matteo Bosso, Daniel Park, Mehrtash Babadi, Bronwyn L, MacInnis, et al. Analysis of 6.4 million sars-cov-2 genomes identifies mutations associated with fitness. Science, 376(6599) 1332,, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Llya Sutskever. Language models are unsupervised multitask learnors. OpenAl Blog, 188):9, 2019.
Lisa Sattenspiel and Klaus Dietz. A structured epidemic model incorporating geographic mobility among regions. Mathematical biosciences, 128(1-2):11-91, 1995.
Wenxian Shi, Jeremy Wohlwend, Menghua Wu, and Regina Barzilay. Improving influenza a vaccine strain selection through deep evolutionary models. bioRwiv, pp. 2023 11, 2023.
Yuelong Shu and John McCaulcy. Gisaid: Global initiative on sharing all influenza fata-from vision to reality. Eurosurveillance, 22(13):30494 2017.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stelano Ermon, and Ben Poole. score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Robin Strudel, Corentin Tallec, Florent Altché, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will Grathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, et al. Self-conditioned embedding diffusion for text generation. arXiv preprint arXio:2211.04. 1.0426, 2022.
Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Nicole N Thadani, Sarah Gurev, Pascal Notin, Noor Youssef, Nathan J Rollins, Chris Sander, Yarin Gal, and Debora S Marks. Learning from pre-pandemic data to forecast viral escape. Nature, 2023.
Maxim Zvyagin, Alexander Brace, Kyle Hippe, Yuntian Deng, Bin Zhang, Cindy Orozco Bohorquez, Austin Clyde, Bharat Kale, Danilo Perez-Rivera, Heng Ma, et al. Genslms: Genome-scale language models revcal sars-cov-2 evolutionary dynamics. The International Journal of High Performance Computing Applicalions, 37(6):683 705, 2023.
14

---

Published in Transactions on Machine Learning Research (03/2025)
A Appendix
A.1 Derivation for transmission model
Assuming that transmission rate matrix Aa(x) has cigenvectors and cigenvalues, and A = UAU Here, A is a diagonal matrix, whose i-th diagonal element is the i-th eigenvalue of A. The i-th column of U is the i-the eigenvector of A. Eq. 1 can be written as
dNc(a) Ao(z)N((z), di dN) UAU- dt dU- N2()) AU- dt dN,(x) AN)
(15)
(16)
(17)
(18)
(x) = exp No))
(19)
in which N₂ U-! exp (At) the matrix exponential of At, which can be easily calculated by taking the exponential in each diagonal element. The Ne(x) can be obtained by
(x(x) (N.(2) = exp (oo))- No(x).
(20)
A.2 Re- pe-aaammttrrizaioon of hierarchical transmission model
We prove that re-parameterization in the hierarchical transmission model does not reduce capacity. This statement is demonstrated under the setting of two groups, and its extension to multiple groups is straight- forward by further partitioning the groups.
Statement 1. If a matriz A E mmxm is diagonalizable, it can be blok-diagonalized as A = W WWW-1, in which Á is a block-diagonal matrix with two blocks.
Proof. If A is diagonalizable, there exists invertible matrix U and diagonal matrix A such that:
= UAU
(21)
Let G e mmx be block-diagonal matrix constructed as follows:
G
(22)
G2
where G1 E Xm and G2 E Rm with m2 + 1222 = m, and both G₁ and G2 are invertible, The inverse of G is also diagoni, given by:
G-1 G
1 c1)
We can write A as:
A UGG-IAGG-IU-1
(23)
(UG)G-)AG(UG)]
(24)
We define W UG and A= G-1AG, so that:
A W -1
15

---

Published in Transactions on Machine Learning Research (03/2025)
in which
(%) G 0
0
(ol
A
0
Gi
0
-
GG -
0
À == G =
G
0
Ag
0
G2
=
0
G2 Ar122G2
(25)
is a block-diagonal matrix.
For any diagonalizable transmission matrix A, there exists invertible matrix W and block-diagonal matrix A such that A = W WAW- This re-parameterization preserves the representation capacity of the original transmission matrix.
A.3 Data pre-processing
A.3.1 Flu
We downloaded the HA sequences from GISAID which was submitted before 2023-03-02. Sequences without the collection time and location information or from non-human hosts are removed. We also discarded the sequences with less than 553 amino acids. Starting from 2003-10, we discretized every two months as one unit of time. We retained the time with more than 10 samples for each continent.
When training the country-level models, we only considered 16 countries with the most available data. For specific testing period (for example, from October 2018 March 2019), we considered the countries and continents with at least 200 viral samples during this period for evaluation, involving 10 countries and 5 continents across all testing periods.
The statistics of data are summarized in Tab. and Tab.
A.3.2 Cov
We downloaded metadata of COVID-19 from GISAID which was submitted before 2023-11-21. We obtain the RBD sequences by applying the annotated mutations the metadata the wild-type Spike protein the mutations lie between position 319 and 541. Sequences without the collection time and location information from non-human hosts removed. Only RBD sequences of length 223 are kept. Starting from 2019-12, discretized every month one unit of time. We retained the time with more than 100 samples for each continent.
When training the country-level models, we considered 36 countrics with the largest number of samples. For a specific testing period (for example, from January 2022 to March 2022), we considered the countries with at least 1000 viral samples during this period for evaluation, involving 36 countries across all testing periods.
The statistics of data are summarized in Tab. 4 and Tab.
A.4 Oracle model
We use the evolution model proposed by Shi et al. (2023) as the oracle model, which defines the probability of sequences at time & as
Pc (<ee) exp(as (K,, xt + ba( xeeo))
(26)
The oracle models are trained for each location (continent or country) independently on data across the training and testing time. To mitigate the bias introduced by the oracle models, we trained three oracle models for each location by changing the pareer-ar.aeters.
For FLU at the continent level, we train three oracle models using a G-layer transformer for as and another G-layer transformer for by with three different random seeds. We trained these models with batch size 32 for 80,000 steps.
For FLU at the country level, we train three oracle models using a 6-layer transformer for as and another 6-layer transformer for by with three different random seeds, for 20,000 steps with batch size 32.
16

---

We use the base GPT-2 model but with the number of layers as 6 and the hidden size as 768. We introduce linear layer to output the transmission rate matrices (size [V] x m) from the last hidden state. To make the transmission rate matrix passive, we apply Softplus function to the outputted transmission rate Published in Transactions on Machine Learning Research (03/2025)
For Cov at the continent level, train oracle model for 100,000 steps with batch size 256 and learning rate le-4, using 6-layer transformer shared by and by. We train the second oracle model with the same architecture but for 30,000 steps. The third oracle model is built on one 6-layer transformer model for ag and another one for bø, and trained with a learning rate of 1c-4 and batch size 160 for 30,000 steps.
For Cov the country level, we train one oracle model with shared 6-layer transformer for 30,000 steps with batch size 256 and learning rate of le-4. The second and third oracle models have the same architecture the first one but are trained by learning rate of 3e-4 for 20,000 steps with two different random sceds.
A.5 hyper-parameter and details of training
matrix. We introduce another linear layer to output the boundary condition (occurrence at lime ( = 0, size [V] x m) from the last hidden state, and apply Softmax function on the vocabulary dimension. We set the À for group regression loss Lyroup to 0.1.
We trained our model on 48G NVIDIA RTX A6000. takes GPU around 18 hours to train a model for FLU the continent level, and GPU 30 to 36 hours train a hierarchical transmission model. For Cov, training the transmission model on the continent level takes GPU around 11 hours. Training the hierarchical transmission model takes 2 GPUs around 15 hours.
A.6 Supplementary results
The NLL and reverse NLL for each continent and country for FLU is shown in Fig. and Fig. and reverse NLL for each continent and country for Cov are shown in Fig. 10 and Fig
The NLL
3
17

---

Published in Transactions on Machine Learning Research (03/2025)
Flu ANH322, africa
Flu 3HH asia
model
model
30000
reverse NLL
Prepend
reverse NLL
Prepend
Add Embed
Add Embed
Param, share
Param_share
LoRA
LoRA
Transmizsion
10-2
Trnsmission
10"2
10
1
6x
10
1x
NLL
NLL
(a)
(b)
Flu A/HEN2, europe
Flu A/HBN2, north america
madel
3
Global
3x10-2
Global
reverse NELL
Prepend
reverse NLL
Pripend
2 Xx.00
Embedd
2 xx10-
Param_share
Param share
LoRA
10-2
Transmission
Tarsmission
4x 10.% 10
10-t
2 10
x 10.1%% 10
10-3
2 10-1
NLL
NLL
(c)
(d)
Flu A/HNN2, oceania
4x10-1
model
3x-10-2
Prepend
reverse
Add Embed
LoRA
Finetune
10-3
Transmission
10-4
NLL
(e)
Figure Average NLL and reverse NLL among years for FLU each testing continent.
18

---

Published in Transactions on Machine Learning Research (03/2025)
Fu AHHENZ, curopelspain
Fl 3/2 asia/china
model
model
4x10
Global
Giobal
raverse NLL
Prepend
reverse NLL
Prepend
Add Embed
Param Shore
Param, shhree
-x10-
LoRA
LoRA
Finctune
Fintune
Tronsmistion(G26)
10
Transmission(G25)
Tra Transmissio((22))
Transmission(
4xx10-xx10-
10
2x10
60000
10
2 xx10
NLL
NLL
(b)
Flu AHHN2 europelgermanyy
MH3N2, europe/russien federation
model
model
Giobal
Global
Prepend
roverse NILL
Prepend
reverse
Add Embeed
Add Embed
LoRA
LoRA
Finetune
Finetune
10
Transmissio((22))
TransmissionG22G)
10a2;
Transmission(G2u)
Transmisslen(G2G)
10
10-3
(c)
(d)
Flu A/HENZ, curope/trance
AHHENZ,
6x 10
10-1
model
Global
Glcbl
reverse NU
-0
reverse NLL
3x 10-
Add Add__mbed Embed
Add Embed
Param, share
Param Share
2x 10-
Anetune
Anetuna
Transmission(220)
10-7-
(on
10-2
4x 6x
10-1 NLL
10
10-2
10"3
NLL
(c)
(I)
Flu A/H3N2, accania/australl
Flu A/HBN2 curopolunite kdngdom
4x 10
model
reverse NLL
30
Prepend
Global
reversee NLL
Prapend
Glotel
Add Embed
Add Embed
2x.00 10
Param share
Param share
LoRA
LoRA
10 7-
Transmission(G2G)
Anetunne
10-2
Finetune
Transmission(G2e)
Transmission((2L)
Transmiss (G22))
6xx35
10"
2xx 10
4x
30
10
10
x 10
NLL
NLL
(g)
(h)
Flu MH3N2, north ameriiaacaaaade
Flu Arl 3N2, north americalunited states
4x
model
model
3
Global
Pregend
Pregend
reverse
Add Embeed
roverse
20=
Add Embedd
Parim_share
LoRA
LoRA
10
Finetune
Finetune
TransmissioniG26)
10-2.
TransmbsloniG26
Transmission(G22)
Transm&ssen(G22U
x
10
(i)
(I)
Figure 9: Average NLL and reverse NLL among years for FLU in cach testing country.
19

---

Published in Transactions on Machine Learning Research (03/2025)
Cov-RBD. africa
Cov- RBD, asia
4x-0-
model
3
Prepend
Prepend
Embad Param, share
Embed
reverse
2x10-
reverse
xx300
Param, share
LoRA
LoRA
Transmission
Transmission
60
10
2
10
10 NUL
x
NLL
(a)
(b)
Cov Cov-RBD, europe
Cov RBD, north_amentt america
4x 10
model
model
--- Global
1x
. Global
reverse NLLL
Prepand
reverse NEL
Prapend
2x
Param_shape share LoRA
<8 Param share
LoRA
Transmission
Transmission
4xi0-%xxo-1
10-1
x
4 10.0%% x ii
10
2 10
NLL
NLL
(d)
Cov-RBID oceania
Cov-RBD, south aamerica
4x
model
3x10
model
revers NLL
Pregend Add Embeed share
roverse NLL
Prepend
2x-2
2x
2x
Add Embed
LoRA
LoRA
Finetune
Transmission
Transmission
4x10-6x100
xxi0
NLL
NLL
(e)
(f)
Figure Average NLL and reverse NLL among years for Cov each lesting continent.
20

---

Published in Transactions on Machine Learning Research (03/2025)
asia/china
Coy-RRD africaNouth
4x10-1
model
Gobal
Global
Propend
Sx 10
Propand
roverte
2010
Add Embed
reverse
Add Embed
Fieetne
2x10
Rnetune
Transmission(220)
Transmission(020)
10-3
Transmission(C23)
Transmisdon((22)
(a)
(b)
Cov-RBD RDD surppenorraa
CoyRBD europeitureey
10
model
4
model
6kx10-)
Global
33x10-
Giobal
4x.0
Add
maverrse
Parsm_shas
Param share
LoRA
LoRe
2x100
Anctune
Aiettne
Transmissicn(C2C)
Trans Trammmisnio((G2))
Transmission(C2L)
Transmission(22L)
10m3
4x
10
2 103%
NLL
NLL
(c)
(d)
sumpe//paland
Cov-RRD
4x
3x10-
Prepend
Glbbal
30
Prepend
2x10
Param share
Add Embed
reverse
10
Add Embed
Faram share
LoRA
Transmission(G20)
Transmission(226
Transmission(221
Tranamislon(C2U)
10.4
10
Cov-RB euro pewwiitziderrdnn
Cov-RBD purpe//uembburg
6x
model
4x
model
4
Glbba
Prepend
NLL
30
Giobal
Prepend
3x10-1
Ambod
Param_share
Param_share
LoRA
LoRA
2x100
Anetune
Aictune
Tracamision(G22c)
Tranamissin(G22))
Transmission(G22)
Transmission(G2))
10
1641
(g)
(b)
Cov-RBD europeiczech republic
Cov-RBD euro
model
6x
model
Glabal
Giobal
4x10-
roverse NLL
4x
3x10-
Param share
Add Embed
3x.10
Add Embed
Param share
LoRA
LoRA
2x10
2x10-2
Rieeune
Yransmissian((26)
Transmission(G2a)
TransmisstonGG2))
Transmisslon(G2L)
6x
10
2x100:w00
10'
NLL
NUL
(i)
(i)
Figure 11: Average NLL and reverse NLL among years for Cov in each testing country (Part T).
21

---

Published in Transactions on Machine Learning Research (03/2025)
OoyRRBD aslafndonessi
OswRBD north amedica/mesic
4x
model
3x
Gobal
Global
Propend
3x10-2
Propand
roversa
Add Embed
Add Embed
2x10--
2x10
Fnetune
Rnetune
Transmission(020)
Transmission((23)
Transmisdin((22)
10
(a)
(b)
Cov-RBD south americaiperu
Cov-RBD RBD eurgpeintterrenns
model
model
4x10
Global
Giobal
4x
3x10
Add
2x10-1
Parsm_shas
LoRA
2x
Param share
Lote
Anctune
Aieetne
TTa0Gmission(C2C)
Tranamisiio(G22
Transmission(2L)
10"4
Transmission(22L)
6x100
10th
2x 10
20 NLL
NLL
(c)
(d)
Cov-BB sumpe/italy
Cov-RBD americadbrazil
28x
4x
2.4x10-2
204
3x10
2 222x10
Presend
reverse NLL
Propend
2x10--
Add imbed
Add Embed
18x10
Faram, share
2x
Param 2har
166x10
LoRA
LoRA
14x10-
TransmissioniG26)
Transmission((22)
Tranamissinn(C2))
Tranemissien G2L)
4x10
10
2x10
6
(f)
Cov-RBI europe/belgium
Cav-HDD eurpeespaan
model
model
Glbba
3x
Giobal
Prepend
Prepend
Ambee
Ambod
Param_share
Param_share
LoRA
LoRA
Anetune
Anctune
Tracamision(G22c)
Tranamissin(G22))
Transmission(G22)
Transmission(G2))
10m1
10
10-3
(g)
(h)
Cav-RBD europe/ussia
Cou-RBD surope/senden
4x10
Global
model
4x
Global
model
3x10
roverse NLL
Prepend
Add Embed
1x
Add Embed
Param share
Param share
LoRA
LoRA
2x
10
Yransmissian(G26)
Transmission(G2a)
Transmisston(G21)
Transmisslon(G2L)
4x
1000
10:
2x
6x
10
2x103xx0
NLL
NLL
(i)
(i)
Figure 12: Average NLL and reverse NLL among years for Cov in each testing country (Part 2).
22

---

Published in Transactions on Machine Learning Research (03/2025)
Cov-RBD aslarindia
RBD aslarsouth
model
naverse NLL
Prepend
3x
Prepend
2x-10
Ad2 Embed
Param, share
reverse
20"30
Add Embed
Param
LoRA
LoRA
10-2
Transmlostenl226)
10-1
2
10%10
Transmission
10-1
Transmission((23)
NLL
NIL
(a)
(b)
noeanis/australa
model
Global
Global
Prepend
NLL
33x10
Prepend
1x10
Add Embeed
Add Embed
Lott
Anetune
2x
Fineure
Transmission(226)
Transmission(220)
Transmission(G2L)
TransmissionG22))
(c)
(d)
Cov europelaustra
Cov.RBD europe/denmark
6x10
model
model
Global
Giobal
Add
Param, share
3x
Param, share
2xx10-
LoRA
Latte
Rowture
Firelune
Transmtssion(G2G)
Transmisslon(22a
Transmission(22))
Transmission(221)
6'x 10-1
2x 10 ic
10-1
NIL
NLL
(c)
(f)
Cov-RBD americalcanada
Cos-BBD
6x
3x
Global
4x-0-1
Prepend
reverse NLL
Prepend
Add Emreed
Add Eimbed
Parer shere
Panam share
2xx10-
TransmbssionGG26)
Transmisslon(22a
Transmission(G2L)
Transmission(G2))
10.00exx 30 36
2x 165%
10
(h)
CBv-RBD surppegemanyy
Cov Cov-RBD RBD curopelunited kingdom
model
model
3x10-
Global
Globel
Prepend
Preperd
Embee
reverse
Ambed
2x10-
2
Param_shane
L2BA
LottA
Finsurne
Fintune
Transmistio((22G)
Transmission(G2G
Transmission(G2))
Transmission(G22)
10-1
2
2x 107xx
(1)
(i)
Cov-BD asiafopan
CovARED north america/usa
10-1
model
3x
model
reverse NUL
noverse NUL
Prepend
Add trebed
Embtd
Param share
Param share
LoRA
LoRA
Aoddune
Firelune
Transmission(G2G)
Transmission622))
Tranemission(G21)
10'3
6x
10
2x 10
NLL
NLL
(k)
(I)
Figure 13: Average NLL and reverse NLL among years for Cov in each testing country (Part 3).
23

---

Published in Transactions on Machine Learning Research (03/2025)
11
82%
0,90
0.75
0.74
0.7
0.68
0:81
Flu 2017
1.0
NA
EU
OC: 55%
0.74
0:83
1.2
8:89
0.87
0.73
0.9
-0.44
0.9
0.89
1.1
0.99
0.85
OC
SA: 94%
8
00.99
-088
SA
AS
$
0.82
0:75
0.74
0.86
0.87
-0.7
AS
NA
AF
(a) Transmission rate (2017)
(b) Maximal spanning tree (2017)
(c) Phylogenctic tree (2017)
Figure 14: (a) Average transmission rate matrix among sequences collected during 2017 winter FLC scason in clade 3C.2a2. (b) The maximal spanning tree obtained from the rates matrices. (c) The phylogenetic tree obtained from the Nextstrain.
atrica
0.41
0.35
0.36
0.34
0.33
0.46
Polacr
0.48
europe-B 0.35
0.42
0.4
0.4
0:37
0.35
0.42
Europe
north, america -0.
0.43
0.4
0.46
0.41
0.38
0.38
South America
North America
nouth_ameria 0.33
0.39
0.35
0.38
0.33
0.39
-0.44
aftca-
sii
europe
north meerca
ccanno
Asia
52%
Africa
Oceania
(a) Transmission rate
(b) Maximal spanning tree
(c) Maximal spanning tree
Figure 15: (a) Average transmission rate matrix among sequences collected from 8022-04 to 2022-06 in lineage BA.2 Cov. (b) The maximal spanning tree obtained from the rates matrices. (c) The phylogenetic obtained from the Nextstrain. Both phylogenetic analysis and our model show that the clade originated from Asia.
24

---

Published in Transactions on Machine Learning Research (03/2025)
Number unique sequences
Number of samples
2015
2016
2017
2018
2015
2016
2017
2018
Europe
1250
1492
2681
3350
3116
3761
7553
9576
North America
2708
3142
4078
5443
8315
10261
13716
19740
Asia
2301
2890
3687
4584
5076
6361
8236
10533
South America
314
439
503
797
919
1221
1339
2088
Oceania
591
954
1255
1718
1459
2364
3482
5275
Africa
354
499
621
792
705
1017
1339
1719
Total
6842
8499
11550
14932
19590
24985
35665
48931
Asia
China
436
604
721
927
904
1251
1478
1890
Asia
India
48
48
53
98
92
93
99
168
Asia
Japan
570
683
887
1099
1073
1296
1705
2076
Asia
Singapore
234
270
367
435
561
627
845
1013
Europe
France
189
219
385
150
276
327
772
922
Europe
Germany
91
116
254
343
146
178
401
562
Europe Italy
65
145
153
119
255
Europe Netherlands
68
79
135
187
92
110
200
274
Europe Russian Federation
184
222
317
420
816
Europe Spain
109
129
185
230
208
249
372
462
Europe Switzerland
22
167
36
41
568
733
Europe United Kingdom
199
225
404
605
533
592
1088
1671
North America Canada
327
369
535
763
649
744
1183
1801
North America United States
2263
2617
3394
4523
7182
8874
11763
16808
Oceania / Australia
467
751
996
1432
1097
1840
2733
4366
South America Brazil
119
192
235
380
232
383
439
752
Table 2: Data statistics for FLU training set.
Number unique sequences
Number samples
2015
2016
2017
2018
2015
2016
2017
2018
Europe
169
1340
603
1043
360
1292
1782
2893
North America
350
1168
1122
812
1123
4271
4691
4026
Asia
380
666
459
631
758
1399
1028
1558
South America
22
100
73
56
36
187
184
162
Oceania
171
280
527
Africa
72
71
41
158
132
162
84
417
Asia China
132
106
74
125
274
185
105
215
Asia Japan
78
200
130
151
138
398
228
369
Europe France
17
180
63
155
25
488
108
427
Europe / Germany
25
178
56
84
33
323
84
213
Europe / Russian Federation
62
124
307
156
352
Europe / Spain
60
58
84
129
125
200
Europe / United Kingdom
209
74
547
673
168
North America / Canada
45
217
249
109
91
572
663
324
North America / United States
965
901
699
3573
3872
3592
Oceania Australia
54
180
98
147
84
479
258
452
Table 3: Data statistics for FLU in testing set.
25

---

Published in Transactions on Machine Learning Research (03/2025)
Number of unique sequences
Number of samples
2021-07
2021-10
2022-01
2022-04
2021-07
2021-10
2022-01
2022-04
Europe
2851
3980
7897
16673
1.4e++06
2.5e++06
3.9e-1066
5.8e++06
North America
2837
3956
7214
12354
997557
1.9e++06
2.9e++06
3.7e++06
Asia
1503
2234
3991
7496
227000
416156
522226
769393
South America
784
1228
1962
3448
88737
159608
217509
285418
Oceania
205
880
25016
45896
73489
113998
Africa
763
1009
1933
2704
48254
75797
101557
120091
Total
6041
8221
15329
29042
2.8e++06
5.1e+066
7.7e+006
1.1e+07
AF AF/South Africa
262
391
885
1275
16586
26623
34622
40382
AS/China
26
40
50
83
1638
2233
2810
4248
AS/India
749
1075
2333
4206
67449
96526
138287
190182
AS/Indonesia
146
214
322
821
6502
11063
13131
24446
AS AS/Israel
190
294
1447
13982
23934
38847
70558
AS/Japan
284
571
604
876
83524
185183
190916
278109
AS/South Korea
145
655
11299
22719
34588
18793
EU/Austrias
220
337
437
884
32811
51220
65025
93616
EU/Belgium
1284
33861
58923
86258
127964
BU/Czech Republic
133
185
317
684
8397
14115
27053
41362
EU/Denmark
340
536
1659
2824
119817
174989
295263
470155
EU/France
492
883
2048
4056
57448
132715
237321
361986
EU/G /Grrmaay
718
1007
2057
5395
150395
230313
366323
599381
EU/ /Teeland
132
230
312
487
21683
38484
52945
69117
EU/Ttaly
433
633
1019
1962
15197
69711
95017
120276
STU/Luxebbuurg
99
136
280
602
11121
15183
21616
31026
EU BU/Netherlande
335
495
666
902
45127
67948
91917
115284
EU/Norway
121
204
288
465
20707
31330
45992
63759
EU/ EU/Poland
204
279
478
1130
18213
23484
44564
82162
EU/Russia
209
304
1139
11873
17524
30290
40634
RU/Slovenis
164
227
409
662
21097
32613
51553
67634
RU/Spain
387
619
938
1822
53870
87005
114204
145622
BU BU/Swede /Sweden
337
469
784
1524
89590
117489
150429
193078
EU /Siiiizzerl
368
533
882
1358
46577
73502
109349
137736
BU/ CUrTurkey
215
347
625
1297
8659
66958
80697
90021
EU BU/United Kingdom
1106
1649
2874
5073
557291
1e-+06
1.7e+-06
2.6e-406
NA/Canada
405
570
1000
1449
112240
167621
237356
306369
NA/ HA/Mexts /Mexico
294
486
1021
1859
20390
36595
48750
61074
NA/U Usa
2578
3638
6518
11207
855032
1.7e++06
2.6e++06
3.3e-+06
OC/ Australia
171
245
432
782
21317
39399
61434
95005
SA/Brat /razil
541
1229
2462
57629
98196
127431
173465
SA 3A/Per
101
153
269
487
5706
10350
16421
22137
Table 4: Data statistics for Cov in training set.
26

---

Published in Transactions on Machine Learning Research (03/2025)
Number of unique sequences
Number of samples
2021-07
2021-10
2022-01
2022-04
2021-07
2021-10
2022-01
2022-04
Europe
4873
11185
6029
5440
1.4e+066
1.9e+-06
677986
489526
North America
4213
6997
4920
4844
976124
875411
585858
513890
Asia
2262
4302
3386
5308
106070
247167
169727
243692
South America
941
1906
1088
1065
57901
67909
41000
39846
Oceania
240
485
405
636
27593
40510
41305
37553
Africa
1066
1106
708
977
25760
18534
17193
10721
AF/South Africa
540
548
348
240
7999
5760
6566
1564
AS/ AS/India
1542
2347
1795
2964
41761
51895
17106
22667
AS AS/Indonesia
136
554
311
504
2068
11315
5056
10019
AS/Tsr //Tsrael
409
962
832
956
14913
31711
32205
33706
AS AS/Japan
77
295
257
483
5733
87193
63937
119228
AS AS/South Korea
228
402
11869
14205
22064
24779
EU/ Austria
166
514
551
606
13805
28591
38237
42594
EU/Belgium
357
688
447
394
27335
41706
20057
15598
BU/Czech Republic
171
433
354
399
12938
14309
7277
7795
BU/Denmark
1232
1681
538
344
120274
174892
70330
47944
BU/France
1382
2566
1576
1039
104606
124665
91470
49631
EU/Germnn /U/Germany
1311
4035
2024
1643
136010
233058
147251
78441
BU EU/Ireland
136
198
167
145
14461
16172
15874
8820
EU/Ttaly
197
1176
859
713
25306
25259
18314
14109
EU RU/Luxembourg
159
395
266
214
6433
9410
8339
7433
EU/Netherlands
269
319
263
280
23969
23367
14001
13988
EU/Norway
122
208
166
164
14662
17767
6602
3477
EU/Poland
251
760
108
149
21080
37598
1571
4001
EU/Russia
239
749
149
428
12766
10344
1912
20857
EU/ BU/Slovenia
220
310
174
173
18940
16081
3822
5653
1081
807
759
27199
31478
27913
19124
RU/Sweden
396
920
355
525
32940
12649
10915
20999
EU/Switzerlan
450
627
246
253
35847
28387
6903
6659
EU/Turke Turrkr
349
745
240
223
13739
9324
4462
3805
EU/United Kingdom
1723
2803
1106
1027
705745
893246
149120
82428
NA/Canade
537
621
427
511
69735
69013
53255
49505
NA/Mexice
649
1093
497
358
12155
12324
8610
11072
NA/Use
3777
6321
4529
4501
885278
785517
512906
445673
OC/ Australia
222
445
381
579
22035
33571
34465
28458
SA/ /razill
557
1478
654
516
29235
46034
21768
14964
SA/ /Per
143
268
232
316
6071
5716
5118
11160
Table 5: Data statistics for Cov in testing set.
Group Countries
Australia, China, India, Japan, Russian Federation, Singapore Brazil, Canada, France, Germany, Italy, Netherlands, Spain, Switzerland, United Kingdom, United States
Table Setting of groups for FLU.
27

---

Published in Transactions on Machine Learning Research (03/2025)
Group Countries
Austria, Belgium, Czech Republic, Denmark, France, Germany, Ireland, Isracl, Italy, Luxembourg, Netherlands, Norway, Poland, Slovenia, South Africa, Spain, Sweden, Switzerland, Turkey, United Kingdom Australia, China, India, Indonesia, Japan, Russia, South Korea Brazil, Canada, Mexico, Peru, USA
Table 7: Setting of groups for Cov.
Flu A/H3N2 (continent level)
Cov-RBD (continent level)
10
4 10-
model
10-1
3x100
Global
Pregend
reverse
10-2
reverse
2 10-2
Param_share
LoRA
Finetune
1
Transmission
6 10
10
2x100 10-1
10-2
10
2x
NLL
NLL
(a) Flu, continent-level
(b) Cov, continent-level
A/H3N2 (country level)
Cov-RBD (country-level)
40-2
model
3.
Global
reverse NLL
Prepend
Add Embed
2 x.10
Param, Param_share share
LoRA
Finetune
10
Transmission(G2G)
Transmission(G2L)
6x 10
10
1xx101
x 1o 111103xx1000 3x1xiio
NLL
NLL
(c) Flu, leeetry-level
(d) Cov, country-level
Figure 16: Negative log-likelihool (NLL) and reverse negalive log-likelihood for FLU and Cov on the wors! nub-population. Lower is better. Error bands represent the 95% confidence interval across different oracle models. The worst sub-population for each model corresponds to the one with the highest average NLL at temperature of 1.
28

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Temperature
NLL
reverse NLL (95% CI)
Global
0.2
0.1921
0.0099 (0.0096, 0.0102)
Prepend
0.2
0.2003
0.0097 (0.0095, 0.0099)
add__mmbed
0.2
0.1951
0.0090
(0.0089, 0.0091)
Param share
0.2
0.2210
0.0239
(0.0234, 0.0244)
LoRA
0.2
0.1914
0.0099
(0.0097, 0.0101)
Finetune
0.2
0.1964
0.0089
(0.0088, 0.0091)
Transmission
0.2
0.1951
0.0095
(0.0091, 0.0099)
Global
0.4
0.0967
0.0105
(0.0103, 0.0107)
Prepend
0.4
0.1010
0.0100
(0.0097, 0.0102)
Add Embed
0.4
0.0984
0.0098
(0.0096, 0.0099)
Param share
0.4
0.1113
0.0250
(0.0247, 0.0253)
LoRA
0.4
0.0964
0.0105
(0.0104, 0.0107)
Finetune
0.4
0.0990
0.0097
(0.0095, 0.0098)
Transmission
0.4
0.0981
0.0097 (0.0094, 0.0100)
Global
0.6
0.0653
0.0124 (0.0122, 0.0126)
Prepend
0.6
0.0684
0.0113 (0.0110, 0.0115)
Add Embed
0.6
0.0666
0.0114 (0.0112, 0.0115)
Param share
0.6
0.0752
0.0257 (0. (0.0254, 0.0259)
LoRA
0.6
0.0651
0.0124 (0.0122, 0.0125)
Finetune
0.6
0.0670
0.0113
(0.0110, 0.0115)
Transmission
0.6
0.0661
0.0105 (0.0102, 0.0108)
Global
0.8
0.0503
0.0186
(0.0183, 0.0189)
Prepend
0.8
0.0527
0.0152
(0.0150, 0.0154)
Add. Embed
0.8
0.0515
0.0158
(0.0156, 0.0160)
Param share Share
0.8
0.0576
0.0274 (0.0271, 0.0277)
LoRA
0.8
0.0502
0.0186 (0.0183, 0.0189)
Finetune
0.8
0.0517
0.0158
(0.0155, 0.0161)
Transmission
0.8
0.0505
0.0132
(0.0130, 0.0134)
Global
0.0424
0.0382
(0.0373, 0.0391)
Prepend
1.0
0.0445
0.0264
(0.0260, 0.0268)
Add Embed
1.0
0.0437
0.0282
(0.0275, 0.0289)
Param share
1.0
0.0479
0.0324
(0.0321, 0.0328)
LoRA
1.0
0.0424
0.0385
(0.0376, 0.0395)
Finetune
1.0
0.0436
0.0300
(0.0291, 0.0308)
Transmission
1.0
0.0421
0.0244 (0.0236, 0.0252)
Table 8: Average negative log-likelihood (NLL) and reverse negative log- likeliidddood for FLU at the continent level. Lower is better. The 95% confidence intervals (CI) across different oracle models are illustrated.
29

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Temperature
NLL
reverse NLL (95% CI)
Global
0.2
0,1880
0.0160 (0.0154, 0.0165)
Prepend
0.2
0.2128
0.0133 (0.0126, 0.0140)
add__mmbed
0.2
0.1895
0.0138 (0.0131, 0.0145)
Param share
0.2
0.1964
0.0156 (0.0131, 0.0182)
LoRA
0.2
0.1880
0.0160 (0.0154, 0.0166)
Finetune
0.2
0.1911
0.0142 (0.0136, 0.0148)
Transmission
0.2
0.1663
0.0140 (0.0129, 0.0151)
Global
0.4
0.0968
0.0167 (0.0161, 0.0172)
Prepend
0.4
0.1093
0.0140 (0.0134, 0.0147)
Add Embed
0.4
0.0976
0.0147 (0.0140, 0.0154)
Param share
0.4
0.1002
0.0158 (0.0138, 0.0178)
LoRA
0.4
0.0967
0.0167 (0.0161, 0.0172)
Finetune
0.4
0.0984
0.0152 (0.0147, 0.0158)
Transmission
0.4
0.0857
0.0154 (0.0144, 0.0165)
Global
0.6
0.0670
0.0191 (0.0185, 0.0196)
Prepend
0.6
0.0755
0.0156 (0.0149, 0.0162)
Add Embed
0.6
0.0679
0.0167 (0.0160, 0.0175)
Param share
0.6
0.0691
0.0167 (0.0151, 0.0183)
LoRA
0.6
0.0670
0.0189 (0.0184, 0.0195)
Finetune
0.6
0.0682
0.0172 (0.0167, 0.0178)
Transmission
0.6
0.0599
0.0177 (0.0167, 0.0188)
Global
0.8
0.0531
0.0245
(0.0239, 0.0251)
Prepend
0.8
0.0595
0.0188
(0.0181, 0.0195)
Add. Embed
0.8
0.0540
0.0207 (0.0199, 0.0215)
Param share Share
0.8
0.0545
0.0186
(0.0173, 0.0199)
LoRA
0.8
0.0531
0.0243 (0.0237, 0.0249)
Finetune
0.8
0.0540
0.0214
(0.0209, 0.0220)
Transmission
0.8
0.0480
0.0217
(0.0207, 0.0227)
Global
0.0459
0.0347
(0.0339, 0.0354)
Prepend
1.0
0.0509
0.0258
(0.0249, 0.0267)
Add Embed
1.0
0.0469
0.0288
(0.0278, 0.0299)
Param share
1.0
0.0467
0.0227
(0.0218, 0.0236)
LoRA
1.0
0.0459
0.0344
(0.0336, 0.0351)
Finetune
1.0
0.0466
0.0302
(0.0294, 0.0310)
Transmission
1.0
0.0420
0.0297 (0.0285, 0.0308)
Table 9: Average negative lieeliodd (NLL) and reverse negative log- ligeliiid for Cov at the continent level. Lower is better. The 95% confidence intervals (CI) across different oracle models are illustrated.
30

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Temperature
NL.
reverse NTL (95% CT)
Global
0.2
0.2061
0.0146 (0.0142, 0.0150)
Prepend
0.2
0.2294
0.0103
(0.0094, 0.0111)
Add Embed
0.2
0.2243
0.0097
(0.0092, 0.0102)
Param share
0.2
0.2494
0.0179
(0.0172, 0.0185)
LoRA
0.2
0.2058
0.0142
(0.0138, 0.0146)
Finelune
0.2162
0.0111
(0.0107, 0.0115)
Transmission G2G)
0.2
0.2113
0.0098
(0.0089, 0.0107)
Transmission (2L)
0.2
0.2213
0.0101
(0.0094, 0.0108)
Global
0.4
0.1038
0.0151
(0.0148, 0.0155)
Prepend
0.4
0.1155
0.0106
(0.0099, 0.0114)
Add Embed
0.4
0.1131
0.0106
(0.0102, 0.0110)
Param share
0.4
0.1256
0.0185
(0.0180, 0.0191)
LoRA
0.4
0.1036
0.0146
(0.0142, 0.0149)
Finetune
0.4
0.1089
0.0117
(0.0113, 0.0122)
Transmission G2G)
0.4
0.1063
0.0100 (0.0092, 0.0108)
Transmission G2LL
0.4
0.1112
0.0104 (0.0097, 0.0110)
Global
0.6
0.0702
0.0177
(0.0174, 0.0181)
Prepend
0.6
0.0780
0.0123
(0.0116, 0.0131)
Add Embed
0.6
0.0765
0.0124
(0.0121, 0.0127)
Param share
0.6
0.0847
0.0199
(0.0195, 0.0204)
LoRA
0.6
0.0700
0.0171
(0.0168, 0.0174)
Finetune
0.6
0.0736
0.0137
(0.0133, 0.0141)
Transmission G22))
0.6
0.0716
0.0111
(0.0105, 0.0117)
Transmission G22L
0.6
0.0749
0.0113 (0.0108, 0.0118)
Global
0.8
0.0540
0.0275
(0.0272, 0.0278)
Prepend
0.8
0.0598
0.0170
(0.0162, 0.0177)
Add Embed
0.8
0.0588
0.0164
(0.0163, 0.0165)
Param share
0.8
0.0648
0.0230
(0.0227, 0.0233)
LoRA
0.8
0.0539
0.0272
(0.0269, 0.0275)
Finetune
0.8
0.0566
0.0207
(0.0204, 0.0210)
Transmission G2G)
0.8
0.0549
0.0148
(0.0143, 0.0152)
Transmission G2L)
0.8
0.0573
0.0148
(0.0144, 0.0152)
Global
1.0
0.0455
0.0626
(0.0620, 0.0632)
Prepend
1.0
0.0499
0.0294
(0.0282, 0.0306)
Add Embed
1.0
0.0491
0.0291
(0.0290, 0.0293)
Param share
1.0
0.0537
0.0302
(0.0299, 0.0305)
LoRA
1.0
0.0453
0.0594
(0.0589, 0.0600)
Finetune
1.0
0.0474
0.0465
(0.0463, 0.0467)
Transmission G2G)
1.0
0.0458
0.0337
(0.0335, 0.0339)
Transmission G22)
1.0
0.0476
0.0314
(0.0312, 0.0316)
Average negative ligeliiel (NLL.) and reverse negative log-likelihood for FLU at the country level. Lower is better. The 95% confidence intervals (CI) across different oracle models are illustrated.
Table 10:
31

---

Published in Transactions on Machine Learning Research (03/2025)
Model
Temperature
NL.
reverse NLL (95% CT)
Global
0.2
0.2103
0.0150 (0.0148, 0.0152)
Prepend
0.2
0.2293
0.0165 (0.0164, 0.0166)
Add Embed
0.2
0.3318
0.0157
(0.0155, 0.0158)
Param share
0.2
0.2478
0.0403
(0.0396, 0.0411)
LoRA
0.2
0.2104
0.0151
(0.0149, 0.0153)
Finelune
0.2
0.2129
0.0150
(0.0148, 0.0152)
Transmission G2G)
0.2
0.1984
0.0144
(0.0142, 0.0146)
Transmission G2L)
0.2
0.1923
0.0146
(0.0145, 0.0147)
Global
0.4
0.1079
0.0167
(0.0165, 0.0169)
Prepend
0.4
0.1173
0.0169
(0.0168, 0.0169)
Add Embed
0.4
0.1682
0.0160
(0.0159, 0.0161)
Param share
0.4
0.1261
0.0388
(0.0382, 0.0395)
LoRA
0.4
0.1079
0.0167
(0.0166, 0.0169)
Finetune
0.4
0.1092
0.0165
(0.0163, 0.0167)
Transmission G2G)
0.4
0.1013
0.0152
(0.0150, 0.0153)
Transmission G2L)
0.4
0.0983
0.0151
(0.0151, 0.0152)
Global
0.6
0.0749
0.0192
(0.0190, 0.0194)
Prepend
0.6
0.0809
0.0179
(0.0178, 0.0179)
Add Embed
0.6
0.1146
0.0169
(0.0168, 0.0170)
Param share
0.6
0.0865
0.0379
(0.0373, 0.0384)
LoRA
0.6
0.0750
0.0191
(0.0190, 0.0193)
Finetune
0.0759
0.0187
(0.0185, 0.0188)
Transmission G2G)
0.6
0.0698
0.0166
(0.0165, 0.0167)
Transmission G2L)
0.6
0.0678
0.0164
(0.0163, 0.0165)
Global
0.8
0.0596
0.0233
(0.0232, 0.0235)
Prepend
0.8
0.0637
0.0200
(0.0199, 0.0201)
Add Embed
0.8
0.0886
0.0188
(0.0188, 0.0189)
Param share
0.8
0.0675
0.0380
(0.0375, 0.0385)
LoRA
0.8
0.0596
0.0232
(0.0231, 0.0234)
Finetune
0.8
0.0603
0.0225
(0.0224, 0.0227)
Transmission G2G)
0.8
0.0549
0.0194
(0.0193, 0.0194)
Transmission G2L)
0.8
0.0534
0.0189
(0.0188, 0.0189)
Global
1.0
0.0517
0.0315
(0.0314, 0.0315)
Prepend
1.0
0.0545
0.0252
(0.0249, 0.0254)
Add Embed
1.0
0.0739
0.0234
(0.0233, 0.0235)
Param share
1.0
0.0569
0.0402
(0.0397, 0.0406)
LoRA
1.0
0.0517
0.0310
(0.0310, 0.0310)
Finetune
1.0
0.0522
0.0298
(0.0297, 0.0299)
Transmission G2G)
0.0468
0.0273
(0.0272, 0.0274)
Transmission G2L)
1.0
0.0458
0.0279
(0.0278, 0.0281)
Average negative log-likelihood (NLL.) and reverse negative log-likelihood for Cov at the country level. Lower is better. The 95% confidence intervals (CI) across different oracle models are illustrated.
Table 11:
32


Reviews:
[{"review_id": "j2Bp8hfktr", "reviewer": "Reviewer_2RdL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n* The paper is well polished. It reads clearly and figures are a great support for detailing the moving parts.\n* The mathematical formulation of the methodology is clear and rigorous and proofs are provided for the necessary parts (thanks to recent revisions).\n* Results are strong, clearly outperforms baselines.\n* Ablation results are well constructed and creates trust in the methodology.\n* Have strong notion of their own limitation (as described in the conclusion).\n\n**Weaknesses**\n* Related works section seems a bit narrow.\n* Few of the baselines are published, notably for baselines that incorporate location information.\n* I see a bit of a parallel/opportunity here with the Neural Ordinary Differential Equations field of research which might provide useful modeling tool for this approach in terms of precision which has not been mentioned here.\n\n[1] Chen, Ricky TQ, et al. \"Neural ordinary differential equations.\" Advances in neural information processing systems 31 (2018).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The manuscript presents an approach to model the propagation of virus (amino-acid sequence) strains across both time and sub-populations (continents, countries). While common approaches model the evolution of such viruses either as a whole or independently for each subpopulation, they present methodology that allows to model the propagation in aggregate across subpopulations while maintaining the ability to predict fine-grained prediction for each country. They achieve this through hierarchical modeling of earth because of the intractability of modeling interactions across each pair of sub-populations. They show state-of-the-art results on the Pareto-frontier in the negative log-likelihood and reverse negative log-likelihood space compared to evaluated baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Could you add a description of what is a protein sequence residue in section 2.1?  \n* In section 2.2, I think a comment about the specific limitations of group-to-location and group-to-group setting would be useful here to understand the compute/expressivity tradeoff. e.g., is the a scenario where these simplification would be incapable of modeling a location-to-location interaction and to which degree? \n* In the ablation study (section 4.4), could you comment on why is the no hierarchy model worse than the G2G clustering model? Shouldn’t it be at least as expressive?\n* A common measure for sub-population distribution shift problem is performance on worst subpopulation. Are there settings here in the results where agglomerated approaches (global, or param sharing) that outperform the transmission approach on their worst subpopulation.\n* Could you elaborate a bit further on how were groups chosen in table 6 and 7? Specifically why not make more than 2 groups for the Flu model and have something closer to the Cov groups?\n* I see the code has been provided in the supplementary material. Would be good to have a sentence in the main text pointing to a repository (Placeholder for now)\n\n**nit**\n* Some of the plots are readable, but a bit hard to read (some points are obscured by others). Perhaps something could be do for that, but not a hard requirement", "broader_impact_concerns": "I see no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0OqYmxvFgC", "reviewer": "Reviewer_W9g8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nI think the formulation of the problem is good. It is clear the authors are familiar with the biology  of the domain, as well as the approaches used in the field.\n\nTaking advantage of the geographic structure is very interesting and reasonable for some aspects of this work. With the amount of intercontinental travel, do variants stay geographically isolated for long periods of time?\n\n“Since egg-based vaccines require lead times of up to 6 months, we train our models on data collected before February of each year, and evaluate models on the sequences collected from October to March of the next year (winter season),” - This is a good evaluation!\n\n“We generate 50 and 500 sequences with the highest probabilities using beam search (Sutskever et al., 2014) and calculate their coverage, defined as the ground-truth frequencies of these sequences occurring in the testing time and locations.” This is also a very interesting way of evaluating performance.\n.\nThe different model ablations and modifications (Global, finetune, prepend, add embed, and parameter sharing) are also very interesting.\n\n“To ensure a fair comparison, we re-trained their protein language models on our training set.” This is great too, and probably a lot of work.\n\nWeaknesses\n\nThe notation is messy and all terms should be defined more clearly in the initial description of terms.\n\n“Since eigendecomposition scales as O(m3 ) with the number of locations, applying our model directly to fine-grained locations such as countries becomes highly time-consuming…” Are you actually doing any eigendecomposition with this model?I thought it was parameterized by a neural network? If so, why even mention this? Or does the model output some data, and then you do eigendecomposition?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors create a biologically inspired machine learning model to predict how viral proteins will evolve over space and time.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "“Throughout this paper, we use location interchangeably with sub-population.” I feel like this is a bit confusing. I had originally thought that ‘location’ meant ‘site’, as in the site of mutation within the protein. Sub-population is seemingly the cleanest and most technically correct.\n\n“ The model outputs Nt(x), the un-normalized probability (occurrence) of x at time t in all locations: Nt(x) = [nt(x; 1), ..., nt(x; m)].” What is m?\n\n“ Intuitively, [Aθ(x)]ij measures the number of people in location i infected by one person from location j during dt.” What is j?\n\n“...outputs |V | transmission rate matrices…” Why the ‘|’ character around the V?\n\nBecause this model is autoregressive, it will generate L different transmission matrices per protein? This isn’t explicitly stated between equations 4 and 5.\n\n“We obtain pt(xs|x<s) by normalizing nt(xs, x<s) over all possible values of xs.” Do you actually normalize over all possible values, or do you just use the autoregressive likelihood.", "broader_impact_concerns": "On a darker note, model could also be used to generate sequences that are highly novel in a subpopulation that are likely to be transmissible, and therefore more deadly.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lOu4znqMr8", "reviewer": "Reviewer_QJzL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The integration of deep learning with differential equations to model viral evolution presents a novel approach.\n- The model's ability to capture the dynamic interactions between sub-populations is important especially for diseases with complex transmission patterns.\n\nWeaknesses:\n- To what extent does the model account for the potential non-linear relationships between transmission rates and viral protein distributions?\n- Could the model's reliance on hierarchical structures lead to overfitting, particularly when the inherent structures relating sub-populations are not accurately represented?\n- What steps are taken to ensure that the hierarchical variation of the model, which re-parameterizes the transmission rate matrix, does not lose critical information when reducing the complexity of the transmission dynamics?\n- How does the model account for the potential variability in viral mutation rates, which could significantly impact the accuracy of protein distribution predictions?\n- Can the authors elaborate on the model's sensitivity to changes in the underlying assumptions about viral transmission pathways, and how this might impact the model's predictive accuracy over time?\n- What steps are taken to ensure that the hierarchical variation of the model, which re-parameterizes the transmission rate matrix, does not lose critical information when reducing the complexity of the transmission dynamics?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript introduces a deep time-series graph neural ordinary differential equation model for predicting the distribution of viral proteins in different sub-populations, aiming to improve upon traditional finite element method simulations in multi-physics field predictions during the curing process of laminated composites.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See Weaknesses.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a well-presented and throughout solid paper that shows an interesting model with thorough empirical evaluation. All reviews are positive. The paper presentation is sufficient for accepting as is. This paper fits well with TMLR.
Decision: Accept (poster)