Non-Parametric State-Space Models Over Datapoints and Sequence Alignments

Authors: Edgar Mariano Marroquin, Yair Schiff, Volodymyr Kuleshov
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Non-Parametric Models, State-Space Models, Genotype Imputation, Comp Bio

Abstract:
Non-parametric models are flexible and can leverage a context set to express rich mappings from inputs to outputs. However, these methods often scale super-linearly in context size, e.g., attention-based
methods scale quadratically in the number of data points, which in turn limits model expressivity.  In this work, we leverage advances in state-space modeling and introduce Non-Parametric State
 Space Models (NPSSM). We find that NPSSMs attain similar performance to existing non-parametric attention-based models while scaling linearly in the number of datapoints. We apply NPSSMs to the task of genotype imputation, where the linear scaling enables larger context sets resulting in competitive performance relative to other methods and widely used industry-standard tools. We also demonstrate the effectiveness of
NPSSMs in the context of meta-learning where the ability to efficiently scale to larger training sets provides more favorable compute-to-accuracy tradeoffs.

Decision: reject
Reviewer scores: 3, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730058733972

thread_type:
official_review

<thread_object>
creator: reviewer_jfin

<content>
summary:
In this work, the authors apply SSMs, an efficient sequence modeling technique, to non-parametric transformers. The experiments validate that the Non-Parametric State Space Models (NPSSM) attain similar performance to exising attention-based methods.

strengths:
- The method is well-motivated. Given that current non-parametric models typically rely on the self-attention mechanism, replacing attention with SSMs is a reasonable approach.
- This submission validates the effectiveness of SSMs in modeling, contributing to a better understanding of SSMs in the field.

weaknesses:
- (**major**) The paper includes an acknowledgment section, at the end of the submission, which potentially reveals author information and violates the double-blind review policy.
- Since SSM can be viewed as an efficient sequential modeling technique, the authors should compare their method with other efficient attention algorithms, including linear attention works.
- It should help if authors provide a formal formulation of \mathcal{L}_{MLM} and \mathcal{L}_{Aux}. I'm still a little confused about the method pipeline.

questions:
Can this approach be applied to other tasks, such as language tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730345285854

thread_type:
official_review

<thread_object>
creator: reviewer_namg

<content>
summary:
This work proposes an SSM approach to non-parameteric learning. i.e., the setting where the dataset itself is used as input to the model, rather than learning a model based on the training data and then discarding the data. This setting is particularly advantageous for meta-learning where one seeks a model for new datasets. 
In previous work, transformers have been proposed for this task (e.g., MSA transformer and NP transformers, mentioned in the paper). The current paper proposes to instead use state space models for this, motivated by their ability to handle longer contexts with less memory requirements.
On the one hand, the idea is reasonable, and the empirical results confirm that performance is competitive and does allow use of longer contexts. On the other hand, it seems like the novel component here is mostly replacing transformers with SSM, but the other conceptual aspects of using it for meta-learning (e.g., various masking losses) have been introduced in the previous NP transformer works.

strengths:
As mentioned above, it is good to see that the advantages of SSM are also manifested in this setting.

weaknesses:
It seems like conceptually the approach largely follows that of NP transformers, and thus the main contribution is to show that state space models are a viable alternative in this case. If the main claim is empirical, I would expect more extensive experiments and clearer gains.

questions:
1. For Table 3, was the HMM trained on chromosome 20?
2. Typo: “Each these methods “ -> “Each of these methods “
3. 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?
4. 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).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721675051

thread_type:
official_review

<thread_object>
creator: reviewer_xwvj

<content>
summary:
1. The paper introduces NPSSM, which adapts state-space models (SSMs) into a non-parametric model for  genotype imputation and protein mutation prediction tasks.
2. NPSSM replaces attention with bidirectional Mamba-based SSM layers which achieves a linear complexity in reference set size, an advantage over Transformers which scale quadratically.
3. NPSSM demonstrates competitive performance with transformer-based models

strengths:
1. The paper is well written.
2. NPSSM achieves competitive performance against SOTA transformer model.
3. Memory usage is low due to the use of SSMs, a known benefit that allows for efficient handling of large datasets and contexts.

weaknesses:
**Limited Novelty**: The primary contribution appears to be the replacement of transformers with SSMs, gaining known benefits like linear scalability, handling of long-range dependencies through selectivity. However, these advantages largely derive from established SSM properties rather than novel methodological improvements.

**Ablating the use of Attribute-Specific Components**: The introduced attribute and data-specific SSM layers may benefit from ablation studies. For instance, it is unclear how much of the heavy-lifting is done by the two components and can we remove one of them. Should the order of SSM application be Attribute and then Data or vice versa. Can we instead flatten the sequence.

**Possibly Insufficient Evaluation**: The model is tested on a limited set of baselines. I am not an expert in these tasks but I feel that expanding the number of tasks and baseline models would strengthen the claims of generalizability.

questions:
N/A please see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730750379266

thread_type:
official_review

<thread_object>
creator: reviewer_zj3c

<content>
summary:
The authors introduce the non-parametric state space model (NPSSM), which scales linearly with the number of data point unlike other attention-based models. The proposed approaches is applied in the context of genotype imputation and meta learning.

strengths:
The proposed approach is simple and inherits the linear scaling of state space models such as BiMamba.

weaknesses:
The methodological component of the paper reads excessively like a pure application (without methodological novelty). First the model basically applies BiMamba twice without further modification. Second, a big  emphasis of the proposed approach is its linear scaling, however, this is not discussed or analyzed in the methods, basically, because it translates directly from BiMamba.

The first experiment (protein analysis) is underwhelming. First they show that increasing k does not improve the performance of neither of the approaches considering, thus defeating the need of a more expressive, but more expensive, model. Second, the model considered had 1M parameters vs. 3.5M in Notin et al. (2023b). Why not consider the larger model at least for NPSSM? Also, from Table 1, it doe not look like PNPT scales (in memory) much worse than NPSSM? One would expect the factor (~2) not to remain that similar for K=1000 and K=1500.

The second experiment is not very convincing because the advantage of NPSSM over MSA Transformer is not clear. This because i) results (Tables 2 and 3) are only presented on a single dataset, ii) the variability and dependency of the main performance characteristics on hyperparameter choices is not clear (though explored via ablation), and iii) computational cost is briefly illustrated in Figure 5. However, is is not clear which context size is used for the main experiments in Tables 2 and 3 and why given the relationship between context size and performance (in Figure 2), the proposed model is not more much better than MSA Transformer in terms of r2.

questions:
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?

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732291525595

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Common Concerns

comment:
## Common Feed Back
We thank the reviewers for their helpful feedback. Many reviewers shared a few common points of criticism so we'll address them here in addition to individually responding to every reviwer.
### Concern 1:  Novelty
Multiple reviewers identified the limited novely of this work as a weakness. The claim is that a significant portion of our advantages come for free from applying state space models to this problem (e.g. linear scaling). 

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 below the effect of these changes.
Designing a loss function (Figure 4)
Annealing schedules (Table 5)
Masking strategies (Table 4, Figure 3)
etc.

 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: Insufficient Evaluation
Multiple reviewers identified that our evaluation was insufficient. We agree and do the following:
1. 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.
2. Add Datasets. A common criticism is that we do not have enough varied datasets to accurately assay generalization. We Add Tables 10 and 11 to address this. 
    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 reviewer ZJ3C concern 3 context size). 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.
    3. (Clarification) Table 3 - Table 3 (results on full chr20) is included to show the generalizability of Non-parametric methods. NPSSM and MSA-Transformer in Table 3 are trained on the first $\sim 1.5\\%$ of chr20, and then evaled on all of it. While it's the same chr, the $\sim 98.5\\%$ of chr20 that the models were not trained on exhibits different relationships between attributes/variants than was observed during training. 
3. **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)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732291631109

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Changelong and Updated Tables

comment:
# Updates
We list the changes we make and the updated Tables for easy reference.
## Changelog
1. Added bootstrapping. 
2. Added SPIN baseline (constant attnetion DL model), Minimac4 (another HMM method) and $(1-r²)$ column to Table 2. 
3. Added context set sizes (k) .
4. Added Ablations on layer ordering and ablating specific layers of NPSSM (Table 9). All models ablated have the same configuration (loss, number of parameters, etc), the changes are mainly in the structure/order of the input.
5. Added imputation results on chr14 (Table 10). Models were trained on chr20, no further training was done. 
6. Added imputation results on HapMAP dataset on chr14 (Table 11) using HapMAP as the reference panel. Models were trained on chr20 on 1k genomes.  
7. Errata in Table 6: MSA Transformer train=650, eval=650: 0.9481 -> 0.9461. 
8. Additionally in the text we make clearer that Tables 2,3 and Figure 1 all use the models with the best performance from Table 6.
9. Typos and misc formatting fixes.

## Updated Figures

### Table 2 (New Baselines)
| Class                       | Method             | $k$                   | $r^2$                | $(1-r²)$
|------------------------------|--------------------------|-----------------------|----------------------|-----|
| Traditional ML        | KNN                              | 4388                  | $0.758 \\pm 0.005$    |  $0.242 \pm 0.005$ |
|                                | LR                               | NA                    | $0.882 \\pm 0.006$    | $0.118 \pm 0.006$ |
| HMM                  | Beagle      | 4388           | $0.943 \\pm 0.002$    | $0.057 \pm 0.002$ |
|                            | Minimac4 | 4388           | $0.931 \\pm 0.003$    | $0.069 \pm 0.003$ |
|                            | Impute     | 4388            | $0.935 \\pm 0.003$ | $0.065 \pm 0.003$ |
| Non-Parametric Models | SPIN                             | 4388                  | $0.868 \\pm 0.010$    | $0.132 \pm 0.010$ |
|                                        | MSA Transformer                  | 650                   | $0.946 \\pm 0.002$  | $0.054 \pm 0.002$ |
|                                        | NPSSM                      | 2000                  | $0.950 \\pm 0.002$  | $0.050 \pm 0.002$ |


## New Figures 
### 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).
| Model                                                                  | $r^2 \pm \sigma$    |
|------------------------------------------------------------------------|---------------------|
| SSM_{attr} $\mapsto$ SSM_{data} | $0.9406 \pm 0.0026$ |
| SSM_{data}  $\mapsto$ SSM_{attr}       | $0.9308 \pm 0.0029$ |
| Flattened                                                              | $0.8241 \pm 0.0067$ |
| $\text{SSM}_{\text{attr}}$ Only                                        | $0.8206 \pm 0.0074$ |
| $\text{SSM}_{\text{data}}$  Only                                       | $0.7783 \pm 0.0049$ |
### 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.
| Class                            | Method                        | $k$                   | $r^2$               |
|----------------------------------|-------------------------------|-----------------------|---------------------|
| HMM             | Beagle   | 4388                  | $0.964 \pm 0.002$   |
| Non-Parametric Models                                  | MSA Transformer               | 650                   | $0.963 \pm 0.002$ |
|                                  | NPSSM                   | 2000                  | $0.967 \pm 0.002$ |

### 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. 

| Class                            | Method                        | $k$                   | $r^2$               |
|----------------------------------|-------------------------------|-----------------------|---------------------|
| HMM             | Beagle   | 1828                  | $0.891 \pm 0.008$   |
| Non-Parametric Models                                  | MSA Transformer               | 650                   | $0.921 \pm 0.007$ |
|                                  | \methodabrv                   | 1828                  | $0.919 \pm 0.007$ |
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces Non-Parametric State Space Models (NPSSM) which scale linearly with the number of data points, unlike attention-based models that scale quadratically.  This linear scaling is inherited from state space models (SSMs) like BiMamba. NPSSMs show competitive performance against existing non-parametric attention-based models.  They also outperform HMM-based methods in genotype imputation tasks.    

However, the reviewers felt that methodological novelty is limited, and that the paper could benefit from a more extensive evaluation. It is tested on a limited set of baselines and datasets. The performance gains of NPSSM over other models, particularly in terms of r2, are not clearly presented and could be further clarified. 

The authors did spent a considerable effort during the rebuttal period for improving the paper (added baslelines, ablations) but the paper probably deserves another round of reviews.

For these reasons I recommend rejection but encourage the authors to further improve their manuscript and resubmit it to another conference.