Thank you for your feedback!

Concern 1: Acknowledgements

To our knowledge, we did not include an Acknowledgements section in the posted version of the paper.

Concern 2: Additional Baselines

• We add in SPIN as a baseline, a non-parametric architecture using linear attention. SPIN was also evaluated on a variant of genotype imputation and was compared against other attention mechanisms and showed competitive or the best performance on their tasks. We incorporate SPIN into our workflow and evaluate SPIN on over 9k untyped variants in Table 2. Note that their tasks only considered a relatively small number of SNPs at a time (O(100)) where we consider a much larger problem more representative of real world use cases. This may account for performance discrepancies between their reported values and this work.
• We additionally add in another HMM baseline, Minimac and update results reporting. Full details are in the General comment, but briefly:
  • We bootstrap our results to get confidence bounds
  • We report the error $(1-r²)$ to make the performance gains even more apparent (e.g a $\\sim 12\\%$ decrease in error rate for NPSSM over the best existing HMM).
  • We add in more datasets (Table 10 on a different Chromosome and Table 11 on a different dataset and different chromosome).

Table 2

Concern 3: Unclear Loss Formulations

Let's consider as an example the imputation problem ignoring batch size to clean up the notation. Our input is $X \\in \\mathbb{R}^{N \\times (A+L)}$ where the first row ($X[0,\\ldots]$ = $[x_1, x_2, \\text{<MASK>}, \\ldots, x_{a+l}]$) is our partially masked query sequence and the rest is our partially masked context set. Our output is $\\hat{Y} \\in \\mathbb{R}^{N \\times (A+L) \\times V}$ for some vocabulary of length $V$ and our labels are $Y \\in \\mathbb{R}^{N \\times (A+L)}$. Our MLM loss is just a standard cross entropy loss calculated only on the reference set $\\mathcal{L}^{\\text{MLM}} = \\text{CEL}(\\hat{Y}[1:,\\ldots], Y[1:, \\ldots])$, and the Aux loss is a weighted version of the cross entropy loss calculated only on the query sequence (dropping the weight term for simplicity), $\\mathcal{L}^{\\text{Aux}}=\\text{CEL}_{\\text{weighted}}(\\hat{Y}[0,\\dots], Y[0, \\ldots])$ .

Therefore the full loss would look like:
$

\mathcal{L}^{\text{total}} &= (1-\lambda) \mathcal{L}^{\text{MLM}} + \lambda\mathcal{L}^{\text{Aux}} \\ &= (1-\lambda) \text{CEL}(\hat{Y}[1:,\ldots], Y[1:, \ldots]) + \lambda \text{CEL}_{\text{weighted}}(\hat{Y}[0,\ldots], Y[0, \ldots])\\
$

for some weight parameter $\\lambda$. Conceptually, we index out the context set and calculcate a MLM loss over that, and we index out some sample/attribute of interest and calculcate a different loss over that (in this case a weighted cross entropy loss). In practice $\\mathcal{L}^{\\text{Aux}}$ can be any arbitrary loss, BCE, MSE, $L^1$, etc, as long as you index the appropriate portions of $\\hat{Y}, Y$ to calculate the loss over and reshape the outputs and targets as necessary (e.g. flattening $Y$ and $\\hat{Y}$ ). We hope this example cleared things up, but if not please let us know where the confusion is and we can attempt to address it.

Questions

Can this approach be applied to other tasks, such as language tasks?

Approaches like this could be applied to language tasks, but they require having a structured/aligned input. In biological sequences, alignments like MSAs have been commonly used for a long time and often occur during analysis, making them a natural use case for models of this sort. Natural language tasks that have a structure that lends itself to an alignment would be a viable task to apply this sort of model to, but if the task doesn't naturally have some sort of structure/alignment, this type of model may not be suitable for that task, leading to this architecture being of particular use in biological tasks.

We hope that if this response adequately addressed the reviewers concerns, they will consider updating their score.

Referenced Tables

Below we attach some of the tables references in the rebuttal for convenience.

Table 2

Table 10 - Chr14 1k Genomes

Imputation performance ($r^2$) evaluated on $19,369$ untyped variants (dev set) from $516$ haplotypes on chromosome 14. The Non-Parametric Models were trained on chromosome 20, and given the nature of the problem there should be no information flow between chromosomes.

Table 11 - Chr14 HapMap

Imputation performance ($r^2$) evaluated on $962$ untyped variants from $400$ haplotypes from Hapmap on chromosome 14. The Non-Parametric Models were trained on chromosome 20 on 1000 Genomes. Only 250 variants were taken per bootstrap.

Table 9 - Ablations

Imputation Performance ($r^2$) evaluated on 9159 variants from 516 haplotypes on chromosome 20. Models were trained for 100,000 steps, each with the same model configuration. Flattened is taking the input H_ref and flattening it down into a single sequence (i.e. laying all the sequences in $\text{H}_{\text{ref}}$ back to back).

Questions

1. For Table 3, was the HMM trained on chromosome 20?

For Table 3- All HMM have parameters specific to the chromosome their being evaluated on. However, the different imputation methods are a mix of closed form parameters given some input and parameters determined by EM, details follow: The HMMs parameters for Impute and Beagle have a closed form based on the Li and Stephens model, and these parameters are determined for each chromosome. Impute and Beagle take in a “genetic map” that describes the structure of a given chromosome to calculate the specific parameters for their transition matrices at various locations along a given chromosome. Given this map, Impute and Beagle define all their transition matrices (not multiple transition matrices, these are example of non-homologous HMMs) .Minimac4 runs EM to fit the transition parameters based on the observed data, while Impute and Beagle don’t do any further fitting of the transition matrices. The emissions matrices are parametrized at the start for all models, but Beagle and minimac also run EM to further fit these parameters based on the observed chromosome data, although in general the parameter usually ends up in the 1e-3 to 1e-5 range. The Non-parametric models were trained on a subset of chr20, they do not have an equivalent of chromosome specific parameters, after training their weights remain fixed from chr20 for all results.

Typo: “Each these methods “ -> “Each of these methods “

These corrections have been made, thank you!

HMMs seem to have similar performance to the non-parameteric methods in Table 2 (differences seem small and not statistically significant unless you show otherwise).. Can you comment on why HMM is an inferior solution in this case?

As noted above, we now report confidence intervals for our results and show that some of the non-parametric models are significantly better than existing HMM methods. The exact algorithm differences from existing HMMs are indeed an interesting question, but are hard to assay. We hypothesize that the following are some of the possible reasons for the increased performance observed.

1. Mutation Rate modeling (e.g. fixed mutation rate) - Existing HMM have an emissions matrix they use to model mismatches between a given observed variant in the query sequence and the reference sequences. The emissions matrix is parameterized by a single value (e.g. 1e-4 for Impute, ~1e-5 for Beagle), and while these methods are robust to reasonable err values [todo: cite impute 4, figure S4], the genome is known to have varying mutation rates along a given chromosome that existing HMMs likely do not model well. In Figure 1 we observe that non-parametric models outperform HMM based methods more at variants with higher MAFs, that is variants that are more likely to have a mutation are better modeled by non-parametric methods. Note that NPSSM still matches or exceeds HMM based methods at rare-variants too.
2. Linear interpolation of untyped variants - Existing HMM methods only model copy states at the typed (observed) variants. The copy states for the unobserved variants are then interpolated between the bounding observed variants. Modeling these untyped variants as a convex combination of the bounding variants state probabilities ignores information that could be used for each untyped variant, such as the mutation rates for each of these untyped variants.
3. In addition, one avenue that remains unexplored is that these non-parametric architectures admit richer inputs. For example, adding in recombination information, demographic information, or phenotype data as auxiliary inputs to this model could be as simple as adding in a few more columns as in PNPT. Incorporating some of this information into existing HMM methods is non-trivial, and this may lead to even further performance gains or the ability to apply these models to a wide range of imputation problems. We look forward to studying these types of additions in follow-up work and further showing how this class of models supplants the existing HMMs for this use case.

Another natural baseline, which I don't see but maybe I missed, is to just train on the target dataset D_C, with an SSM model (but still perform some pre-training with masking).

This baseline should correspond to the SSM_attr only ablation in Table 9, which is a SSM model with D_C used only as a training set and not added as an input during inference. Not having D_C at inference time seems to make the problem harder.

We hope that if this response adequately addressed the reviewers concerns, they will consider updating their score.

We thank the reviewer for their feedback

Concern 1: Novelty

Our novelty lies in the following contributions:

• We show for the first time that state-space models can be applied to a new task: non-parametric modeling of datapoints and sequence alignments. This expands the scope of tasks that can be addressed with SSMs.
• We show that SSMs over MSAs are more scalable and performant than transformers over MSAs. The MSATransformer has been immensely important in biological applications (e.g., protein LMs, AlphaFold); improving it has potential for high impact. In that regards, our novelty relative to Mamba is analogous to the novelty of MSATransformer relative to the transformer.
• We additionally derive a novel probabilistic formulation of these models and a meta-learning objective that goes beyond classical masked training of MSA models. We report that this training is crucial to solve the imputation task (see Table 5, a model trained only on a MLM objective version achieves only $85\\%$ of the performance our full setup achieves ).
• While our architecture builds on top of BiMamba, producing a performant model requires solving a number of technical challenges. We report in the following sections the effect of these changes.
• Designing a loss function (Figure 4)
• Annealing schedules (Table 5)
• Masking strategies (Table 4, Figure 3)

Genotype imputation has largely worked off the same model (Li and Stephens HMM) for over 2 decades now, and we show that these models have the potential to supplant these existing models. This however did not come from applying SSMs to the problem, a significant amount of engineering work was required.

Concern 2: Limited Eval

We do the following to address our insufficient evaluation:

2a. Add Baslines:

1. We Add Semi-Parametric Inducing Point Networks (SPIN) as a baseline. This model is another non-parametric model, but it achieves linear scaling in the context size and was also evaluated in a Genotype Imputation setting against multiple other linear scaling architectures.
2. We Add in Minimac4 as a baseline, another HMM based method. This includes all commonly used methods for Imputation.

2b. Add Datasets We add further evaluation (Tables 10 and 11).

1. (Addition: New Dataset) Table 11 (results on chr14 on HapMap genomes ) - We add results taking the models trained on a portion of chr20 on 1k genomes and evaluating them on chr14 of HapMap, an entirely different dataset. We split the HapMap dataset into a HapMap specific reference panel and a HapMap specific evaluation set. Given the different chromsome from training, and the context set being used, we argue there is no possiblity for any sort of data leakage and this adequately asses generalization. Note that this dataset is a worst case scenario for long context models (see Questions 3 for details), Yet we still observe HMM methods < NPSSM $\\approx$MSA-Transformer on the set of data.
2. (Addition: Different Chromosome) Table 10 (results on chr14 on 1k genomes) - We add results taking the models trained on a portion of chr20 and evaluate them on chr14. We argue that chr14 shows entirely different patterns of relationships between variants and samples (e.g. even if 2 samples/haplotypes/datapoints were full copies of each on chr20, this need not hold for a different chromosome), and so chr14 should be considered as it's own dataset for the purposes of assesing generalizability. We observe HMMs < MSA Transformer < NPSSM on this set of data.

2c. Add Statistical Significance. We re-compute by bootstraping our results (10 bootstrap samples of $\sim 20\\%$ of the data by default). We then report the mean $r^2$ and $\\sigma$ of these trials and generally show HMM Methods < MSA-Transformer < NPSSM. For example, we now see that NPSSM is significantly better than HMMs in all experiments. In addition, we now also show the error rate$(1-r²)$ to better illustrate the performance gains. In Table 2 we see NPSSM is $\sim 12\\%$ better than the best HMM in terms of error rate (0.057 vs 0.050)

We thank the reviewer for their feedback

Concern 1: Novelty

Our novelty lies in the following contributions:

• We show for the first time that state-space models can be applied to a new task: non-parametric modeling of datapoints and sequence alignments. This expands the scope of tasks that can be addressed with SSMs.
• We show that SSMs over MSAs are more scalable and performant than transformers over MSAs. The MSATransformer has been immensely important in biological applications (e.g., protein LMs, AlphaFold); improving it has potential for high impact. In that regards, our novelty relative to Mamba is analogous to the novelty of MSATransformer relative to the transformer.
• We additionally derive a novel probabilistic formulation of these models and a meta-learning objective that goes beyond classical masked training of MSA models. We report that this training is crucial to solve the imputation task (see Table 5, a model trained only on a MLM objective version achieves only $85\\%$ of the performance our full setup achieves ).
• While our architecture builds on top of BiMamba, producing a performant model requires solving a number of technical challenges. We report in the following sections the effect of these changes.
• Designing a loss function (Figure 4)
• Annealing schedules (Table 5)
• Masking strategies (Table 4, Figure 3)

Genotype imputation has largely worked off the same model (Li and Stephens HMM) for over 2 decades now, and we show that these models have the potential to supplant these existing models. This however did not come from applying SSMs to the problem, a significant amount of engineering work was required.

Concern 2: Limited Eval

We do the following to address our insufficient evaluation:

2a. Add Baslines:

1. We Add Semi-Parametric Inducing Point Networks (SPIN) as a baseline. This model is another non-parametric model, but it achieves linear scaling in the context size and was also evaluated in a Genotype Imputation setting against multiple other linear scaling architectures.
2. We Add in Minimac4 as a baseline, another HMM based method. This includes all commonly used methods for Imputation.

2b. Add Datasets We add further evaluation (Tables 10 and 11).

1. (Addition: New Dataset) Table 11 (results on chr14 on HapMap genomes ) - We add results taking the models trained on a portion of chr20 on 1k genomes and evaluating them on chr14 of HapMap, an entirely different dataset. We split the HapMap dataset into a HapMap specific reference panel and a HapMap specific evaluation set. Given the different chromsome from training, and the context set being used, we argue there is no possiblity for any sort of data leakage and this adequately asses generalization. Note that this dataset is a worst case scenario for long context models (see Questions 3 for details), Yet we still observe HMM methods < NPSSM $\\approx$MSA-Transformer on the set of data.
2. (Addition: Different Chromosome) Table 10 (results on chr14 on 1k genomes) - We add results taking the models trained on a portion of chr20 and evaluate them on chr14. We argue that chr14 shows entirely different patterns of relationships between variants and samples (e.g. even if 2 samples/haplotypes/datapoints were full copies of each on chr20, this need not hold for a different chromosome), and so chr14 should be considered as it's own dataset for the purposes of assesing generalizability. We observe HMMs < MSA Transformer < NPSSM on this set of data.

2c. Add Statistical Significance. We re-compute by bootstraping our results (10 bootstrap samples of $\sim 20\\%$ of the data by default). We then report the mean $r^2$ and $\\sigma$ of these trials and generally show HMM Methods < MSA-Transformer < NPSSM. For example, we now see that NPSSM is significantly better than HMMs in all experiments. In addition, we now also show the error rate$(1-r²)$ to better illustrate the performance gains. In Table 2 we see NPSSM is $\sim 12\\%$ better than the best HMM in terms of error rate (0.057 vs 0.050)

We attach a copy of the relevant tables below for your convenience.

Concern 3: Ablations We incorporate these suggestions as Table 9 (attached below for convenience). Briefly, we trained 5 separate models with the exact configurations (i.e. same number of parameters, same losses, etc) for 100k steps. The only differences for each model were:

• The order of the layers
  • This has a noticeable effect on the performance of the model, although it is possible that this difference may disappear for longer training budgets.
• Flattening the input sequence. Instead of having an input of [num_datapoints, num_attributes, embed_dim], we flatten this into [num_datapoints * num_attributes, embed_dim]. The information contained is the same as the base model, but without the MSA structure.
  • Removing the structure of the sequence severely impacts performance. It may be possible that training for long might result in a decreased performance gap, but at 100k steps there is a significant decrease in performance when removing the structure of the problem.
• SSM_attr only - We model each relationship between each attribute (variants) in a sample sample.
  • This is equivalent to ablating away the context set, the model is given a batch of partially masked inputs and is trying to unmask each sequence given only the unmasked variants of each individual sequence.
• SSM_data only - We model the relationship between all the data points for a given variant.
  • This model doesn’t learn much, the performance plateaus after ~10k steps. Given that the input sequence to this model is the sequence of a given attribute, it is likely that this version of the model is essentially learning a weighted average of the provided sequences. Given that the performance of this model is very similar to that of the KNN baseline, is is probable that this is the case. Given the results of these ablations, we believe it to be the case that the structure input and both the attribute and datapoint level layers are important to the model.

We hope that if this response adequately addressed the reviewers concerns, they will consider updating their score.

Table 9 - Ablations

Imputation Performance ($r^2$) evaluated on 9159 variants from 516 haplotypes on chromosome 20. Models were trained for 100,000 steps, each with the same model configuration. Flattened is taking the input H_ref and flattening it down into a single sequence (i.e. laying all the sequences in H_ref back to back).

Table 2

Table 10 - Chr14 1k Genomes

Imputation performance ($r^2$) evaluated on $19,369$ untyped variants (dev set) from $516$ haplotypes on chromosome 14. The Non-Parametric Models were trained on chromosome 20, and given the nature of the problem there should be no information flow between chromosomes.

Table 11 - Chr14 HapMap

Imputation performance ($r^2$) evaluated on $962$ untyped variants from $400$ haplotypes from Hapmap on chromosome 14. The Non-Parametric Models were trained on chromosome 20 on 1000 Genomes. Only 250 variants were taken per bootstrap.

Concern 3: PNPT Experiments

3a. Performance seems to be invariant to context set size K

• This is true for some tasks. Certain tasks did see an increase in performance with model size (e.g. TAT_HV1BR_Fernandes_2016 @ k=1000 = 0.577, k=1500 = 0.595, k=2000 = 0.605), but on average most tasks we evaluated did not show an appreciable difference based on context size. However we do show performance matching existing methods on these tasks while consuming less memory, and this was important for the experiments in Imputation.

3b. Why not bigger models (at least for NPSSM)

• For the initial experiments, we wanted to keep the number of parameters as close as possible. If we train a larger NPSSM model, we gain modest performance gains in a few tasks, but on average scaling up to 3M parameters does not result in a performance gain.

3c. Memory Scaling not as expected

• The 2x memory scaling does not hold for larger $k$. For example k=2000 uses ~29 Gb of Memory.

Questions

1. From Table 3 and Figure 2 it is not clear how the MSA Transformer can reach a r2=0.956 when it runs out of memory before reaching r2=0.95 in Figure 2?

   Figure 2 and Table 3 are on different sets of data, Figure 2 and Table 2 are both on a subset ($1.5\\%$ of chr20, 9k untyped variants) and Table 3 is on the entirety of chr20 (670k untyped variants).

2. Table 2 and Figure 2 are on a subset of Table 3. Why not considering the experiment in Table 6, but starting from the best model in Table 3? Note also that the value of k used in the main experiments is not noted.

   The results from Table 3 are the best models from Table 6. Specifically, the models in Table 2, Table 3, and Figure 1 are the best performance models from Table 6. There is an Errata for Table 6, the best MSA Transformer model (with k=650) should be 0.9461 (consistent with the 0.946 in Table 2 and Figure 1).

3. Why does there seem to be a marginal return on context set size?

Context size is well known to be an important factor in imputation performance, but depending on the dataset used its benefits can be attenuated. For example, if my reference panel contains the mother and father of the query, I should be able to nearly perfectly impute the query with a context size of 4, and more context would yield marginal returns. However, if instead we have 1000s of samples where eachonly matches the query at $\sim 1\\%$, then the context size is very important since the larger the context size the more likely we can find the relevant information we need to impute the query. For example, the HapMap dataset (Table 11) is known to contain many smaller groups of related individuals, so the benefit of larger context sizes is limited. But even then we still see better performance from the Non-parametric models compared to the HMMs. In practice however, imputation is usually performed with large set of moderately related individuals, where context size really helps.

We hope that if this response adequately addressed the reviewers concerns, they will consider updating their score.

Table 2

New Figures

Table 10 - Chr14 1k Genomes

Imputation performance ($r^2$) evaluated on $\sim 19,369$ untyped variants (dev set) from $516$ haplotypes on chromosome 14. The Non-Parametric Models were trained on chromosome 20, and given the nature of the problem there should be no information flow between chromosomes.

Table 11 - Chr14 HapMap

Imputation performance ($r^2$) evaluated on $962$ untyped variants from $400$ haplotypes from Hapmap on chromosome 14. The Non-Parametric Models were trained on chromosome 20 on 1000 Genomes. Only 250 variants were taken per bootstrap.

We thank the reviewer for their feedback

Concern 1: Novelty

Our novelty lies in the following contributions:

• We show for the first time that state-space models can be applied to a new task: non-parametric modeling of datapoints and sequence alignments. This expands the scope of tasks that can be addressed with SSMs.
• We show that SSMs over MSAs are more scalable and performant than transformers over MSAs. The MSATransformer has been immensely important in biological applications (e.g., protein LMs, AlphaFold); improving it has potential for high impact. In that regards, our novelty relative to Mamba is analogous to the novelty of MSATransformer relative to the transformer.
• We additionally derive a novel probabilistic formulation of these models and a meta-learning objective that goes beyond classical masked training of MSA models. We report that this training is crucial to solve the imputation task (see Table 5, a model trained only on a MLM objective version achieves only $85\\%$ of the performance our full setup achieves ).
• While our architecture builds on top of BiMamba, producing a performant model requires solving a number of technical challenges. We report in the following sections the effect of these changes.
• Designing a loss function (Figure 4)
• Annealing schedules (Table 5)
• Masking strategies (Table 4, Figure 3)

Genotype imputation has largely worked off the same model (Li and Stephens HMM) for over 2 decades now, and we show that these models have the potential to supplant these existing models. This however did not come from applying SSMs to the problem, a significant amount of engineering work was required.

We argue that the linear scaling is analyzed, particularly in Figure 2 where we show performance as a function of context set and Figure 5 where we show how the linear scaling of SSMs allows them to outperform Transformer based architectures through the Virtue of having a larger context set in imputation.

Concern 2: Imputation Experiments

We do the following to address our insufficient evaluation:

2a. Add Baslines:

1. We Add Semi-Parametric Inducing Point Networks (SPIN) as a baseline. This model is another non-parametric model, but it achieves linear scaling in the context size and was also evaluated in a Genotype Imputation setting against multiple other linear scaling architectures.
2. We Add in Minimac4 as a baseline, another HMM based method.

2b. Concern: Insufficient Datasets - Add Datasets We add further evaluation (Tables 10 and 11).

1. (Addition: New Dataset) Table 11 (results on chr14 on HapMap genomes ) - We add results taking the models trained on a portion of chr20 on 1k genomes and evaluating them on chr14 of HapMap, an entirely different dataset. We split the HapMap dataset into a HapMap specific reference panel and a HapMap specific evaluation set. Given the different chromsome from training, and the context set being used, we argue there is no possiblity for any sort of data leakage and this adequately asses generalization. Note that this dataset is a worst case scenario for long context models (see Questions 3 for details), Yet we still observe HMM methods < NPSSM $\\approx$MSA-Transformer on the set of data.
2. (Addition: Different Chromosome) Table 10 (results on chr14 on 1k genomes) - We add results taking the models trained on a portion of chr20 and evaluate them on chr14. We argue that chr14 shows entirely different patterns of relationships between variants and samples (e.g. even if 2 samples/haplotypes/datapoints were full copies of each on chr20, this need not hold for a different chromosome), and so chr14 should be considered as it's own dataset for the purposes of assesing generalizability. We observe HMMs < MSA Transformer < NPSSM on this set of data.

2c. Concern: Unclear Gains - Add Statistical Significance. We re-compute by bootstraping our results (10 bootstrap samples of $\sim 20\\%$ of the data by default). We then report the mean $r^2$ and $\\sigma$ of these trials and generally show HMM Methods < MSA-Transformer < NPSSM. For example, we now see that NPSSM is significantly better than HMMs in all experiments. In addition, we now also show the error rate$(1-r²)$ to better illustrate the performance gains. In Table 2 we see NPSSM is $\sim 12\\%$ better than the best HMM in terms of error rate (0.057 vs 0.050)

2d. Concern: Unclear Hyper parameter Choices - Add Details - MSA-Transformer and NPSSM use many of the same hyperparameters (e.g. they both use the same annealed loss, the same weighted aux loss, loss schedules, masking strategies, etc). The hyper param choices are always the ones that showed the best performance in section 4. The largest differences between MSA-Transfromer and NPSSM in terms of hyperparms is the context set size, which has been added to the relevant tables.