Chimera: State Space Models Beyond Sequences

Authors: Aakash Lahoti, Tanya Marwah, Ratish Puduppully, Albert Gu
Primary area: foundation or frontier models, including LLMs
Keywords: Deep Learning Architectures, Sequence Models, State Space Models, Mamba

Abstract:
Powerful deep learning methods based on Transformers are used to model diverse data modalities such as sequences, images, and graphs. 
These methods typically use off-the-shelf modules like self-attention, which are domain-agnostic and treat data as an unordered set of elements.
To improve performance, researchers employ inductive biases—such as position embeddings in sequences and images, and random walks in graphs—to inject the domain structure, or *topology*, into the model.
However, these inductive biases are carefully engineered heuristics that must be designed for each modality, requiring significant research effort.
In this work, we propose *Chimera*, a unified framework that mathematically generalizes state space models to incorporate the topological structure of data in a principled way.
We demonstrate that our method achieves state-of-the-art performance across domains including language, vision, and graphs. Chimera outperforms BERT on the GLUE benchmark by 0.7 points, surpasses ViT by 2.6% on ImageNet-1k classification accuracy, and outperforms all baselines on the Long Range Graph Benchmark with a 12% improvement on PascalVOC.
This validates Chimera's methodological improvement, which allows it to directly capture the underlying topology, providing a strong inductive bias across modalities.
Furthermore, being topologically aware enables our method to achieve a linear time complexity for sequences and images, in contrast to the quadratic complexity of attention.

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

Reviews and discussion:
<thread>
creation date:
1730100600188

thread_type:
official_review

<thread_object>
creator: reviewer_uxiv

<content>
summary:
Chimera is a new framework that generalizes state space models to effectively incorporate the topology of diverse data types, such as sequences, images, and graphs, into its structure. Unlike traditional methods that treat data as an unordered set, Chimera captures the inherent topology of data, leading to methodological improvements and state-of-the-art performance across various domains. It achieves this by using a structured masked attention mechanism that parameterizes the adjacency matrix to accumulate influences between nodes, allowing for linear time complexity and outperforming existing models like BERT and ViT in benchmarks.

strengths:
This paper introduces a versatile, fundamental model architecture that stands in contrast to the commonly utilized Transformer models. The effectiveness of this architecture is demonstrated through experiments in vision, graph analysis, and natural language processing. This work represents a significant contribution to the ML community. Addressing the minor weaknesses below would strongly support its acceptance.

weaknesses:
1. The fairness of the configurations in Tables 7 and 8 appears to be questionable. (Question 1)

2. The scalability of the proposed architectures needs to be addressed. (Questions 2, 3, and 4)

questions:
1. In the experiments, the Chimera model possesses approximately twice the depth of the ViT model. It is a well-established fact in CV and NLP that depth contributes more to model power than width. While the comparison between SSMs and Transformers in terms of depth is common, it prompts the question: would the Chimera model maintain its superior performance if it had the same number of layers and parameters as the Transformer models?

2. Memory and computational efficiency are often cited advantages of SSMs over Transformers. The inclusion of both theoretical and empirical comparisons between the two would enhance the contribution of this paper.

3. Previous works involving SSMs and the Mamba model have indicated that larger models do not always outperform smaller ones, possibly due to various factors. Could the authors conduct experiments with smaller models in CV and NLP to investigate this phenomenon?

4. Some previous works show that the optimization landscape of Transformers is better than that of SSMs. Could the authors add a figure on training loss to compare the convergence of ViT/BERT with Chimera?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712912682

thread_type:
official_review

<thread_object>
creator: reviewer_nw18

<content>
summary:
This paper introduces an SSM framework that can be adapted to several data topology. Moreover, they propose an efficient way of processing the topology when the structure is in the form of a DAG, together with a technique to decompose common modalities such as images and text into DAG components. The method is tested on 3 different modalities including image, text and graph benchmarks.

strengths:
- Incorporating the data topology in the SSM architecture via the resolvent of the adjacency matrix is a useful tool for extending SSM models into graph community. While most of the existing approaches are flattening the graphs and rely on the expressivity of positional encodings to directly apply the SSM on a sequence, CHIMERA are efficiently incorporate the graph structure via adjacency matrix masking.
- The fact that, when the structure is constrained to be a line graph,  the model correspond to the existing sequence-based SSM,  guarantees the good performance beyond graph domain.
- The decomposition of existing structure into DAG components improves the efficiency, an important advantage for SSM models over existing architectures.
- The experimental results show competitive performance with reduced complexity.

weaknesses:
- The paper claims to propose a method that is able to deal with several types of data. However, the core contribution is integrating a graph structure into the SSM, leveraging the fact that sequences  can be represented as linear graphs, while images can be represented as grid-like graphs. Based on this well-known fact, any graph architecture can be (over)stated as being apliable across modalities.
- The operator proposed in Equation(12) is wildly used in spectral-based graph neural networks. Rediscovering it via an SSM lens and using specific techniques to compute it efficiently is an important contribution. However, discussing its relationship with attention-based GNNs (like GAT and Graph Transformers) or efficient GNN variants would provide a more comprehensive overview of the field, extending the context beyond SSMs
- A key advantage of employing SSMs consist in reducing complexity and creating more efficient computation.  While the paper does a good job on motivating linear complexity (as opposed to the quadratic one used by models such as BERT), it fall short on evaluating the efficiency experimentaly. Reporting the empirical time complexity would help better emphasize the advantages of using CHIMERA compared to existing models.

questions:
Please see the sections above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721360088

thread_type:
official_review

<thread_object>
creator: reviewer_ygvr

<content>
summary:
The paper proposes Chimera, a framework that generalizes state space models (SSMs) to handle arbitrary data topologies by reformulating them through resolvent operations on adjacency matrices. The paper introduces a principled mathematical approach to incorporating topology into SSMs, achieving competitive results across various tasks. However, the method faces practical limitations for general topologies where DAG assumption does not hold and approximations are required. In addition, the discussion is lacking if the method can also generalize to more structured topologies such as Lie groups.

strengths:
1. The paper potentially expands the domain of application of recent state-space models (with SMA form) beyond plain sequences. 
2. Theorem 4.5 and the subsequent sections on optimized forward and backward pass implementation make the proposed method the most efficient for DAG-topologies.
3. The normalization scheme is well-motivated and theoretically grounded. I can see it is an essential ingredient for applying Chimera to the wide range of topologies.
4. Experiments demonstrate clear advantages of the proposed Chimera model in comparison with transformer, mixer, M2 and Hyena architectures.

weaknesses:
1. The insight of SSMs operating on directed line topology is obvious and I suppose trivially holds for any sequence model.
2. Figure 1 does not contribute much to the narrative in the paper. I suggest the authors make a figure more specific to the method described in the paper.
3. While the proposed method claims to generalize SMA-SSMs to any topology, I am not sure it can handle topological groups such as Lie groups. The paper would benefit from additional discussion on the possible extension of SSMs to group-structured topologies given the rising interest in alternatives to self-attention for group-equivariance [1,2,3 etc].
4. While in theory, the proposed approach applies to a wide range of topologies, the practical application requires either a strong assumption like DAG or approximating L as in Eq.19 which can be costly for large graphs. It limits the practical applicability of the proposed model, especially given the results in Table 5 that demonstrate the exact version to work much better than the approximated one.
5. The paper will benefit from a more detailed discussion on computational complexity in the general case when the DAG assumption does not hold.
6.  While I understand the motivation of setting k=diam(G), the paper would benefit from the ablation study on how the performance depends on k when k < diam(G). This ablation study aims to verify if the method can be optimized for large graphs with large diam(G).

[1]. D. Liang et al. PointMamba: A Simple State Space Model for Point Cloud Analysis. NeurIPS 2024.
[2]. A. Moskalev et al. SE(3)-Hyena Operator for Scalable Equivariant Learning. GRaM@ICML 2024.
[3]. Y. Schiff et al. Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling. ICML2024.

Minor:
 - line 161 footnote typos
 - Notation clash in Proposition 4.1: consider changing A^T to A^{|V|}, otherwise, A^T reads as transposed.

questions:
Questions:
 - Is it possible to extend the proposed method to incorporate edge features when provided? I realize that the proposed method is a first step towards topology-aware SSMs, a brief discussion on the possibility of incorporating edge features will be useful.
 - Why are the results in Tables 1,2,3 reported only over one data split?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730898906996

thread_type:
official_review

<thread_object>
creator: reviewer_euvy

<content>
summary:
This paper studies state-space models for data beyond sequences. Starting from an observation that connects the mask matrix in Structured Masked Attention and the resolvent of the adjacency matrix, this paper proposes to utilize this observation to inject the topological structure of data into SSMs such that the SSMs can deal with data structures beyond sequences, e.g., images and graphs. This is done by parametrizing the "adjacency matrix" accordingly for each data structure. Multiple tricks are also proposed to improve the accuracy and efficiency of computing the resolvent, including normalization for better numerical conditions and approximation for faster speed. Experiments have demonstrated its great performance.

strengths:
1. This paper is well-written and easy to follow.
2. The observation is interesting and makes sense.
3. The implementation (i.e., those tricks for improving Chimera) is straightforward.
4. The empirical results look good.

weaknesses:
1. I know many experiments have been done since this paper claims contributions over multiple data structures, but I do want to see if Chimera can outperform those SSMs proposed in their respective domains, e.g., mamba2 for NLP tasks, vision mamba for CV tasks, etc.
2. The efficiency analysis is not enough: a detailed speed benchmark and description of the way of computing resolvent in practice are needed.
    - Proposition 4.4 involves the use of Gaussian elimination,  whose efficiency may highly depend on practical cases. Could the author clarify what implementation is used in practice to achieve guaranteed linear complexity?
    - If I'm not mistaken, when dealing with graph-structured data, the claimed SOTA performance is achieved by computing the resolvent exactly, right? Does this take cubic complexity? And when using the approximation trick for better efficiency, actually the performance is largely decayed, which is no better than the previous graph mamba works in most cases.
    - BTW, I'm curious about how things are implemented exactly. I was wondering if the authors plan to release the code.
3. How can the model deal with edge features for graph-structured data?
4. It seems to me this work is similar in spirit to another recent graph SSM work [1], although starting from a different story. For example, 1) their LapPE and the resolvent in Chimera for global dependency; 2) their computation can be written similarly as Eq. (5) in Chimera to achieve linear complexity while capturing global dependency; 3) their local + global module v.s. the global module and the local convolution in Chimera. Their work is so recent, so I believe the comparison is not required. But it would be great if the authors could provide a discussion of this work and Chimera.

My major concern is about the practical efficiency, which I think is not presented comprehensively yet in the current manuscript. I believe at least a speed benchmark is required. Overall, I like this paper and I'm more than willing to increase the score if this concern can be resolved.  

[1] Huang, Yinan, Siqi Miao, and Pan Li. "What Can We Learn from State Space Models for Machine Learning on Graphs?." arXiv preprint arXiv:2406.05815 (2024).

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732235879629

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We sincerely thank all the reviewers for their valuable feedback and their recognition of Chimera as a *fundamental model architecture*, with a *well-motivated design* and *theoretical grounding*. We also appreciate their acknowledgement of Chimera's *great performance on diverse data types*.

In this global response, we highlight our key contribution is a **methodological shift in modeling graph topologies, eliminating the need for ad-hoc position embeddings**. Chimera leverages SSMs for this work primarily for their ability to capture topology rather than their linear complexity. While the current PyTorch implementation has speed limitations, we believe that building dedicated CUDA/Triton kernels offers an exciting opportunity to further enhance Chimera's scalability.

Furthermore, we address the concerns about the Approximation method with improved results on Peptides Struct and Peptides Func achieved through hyperparameter tuning. The approximation method now **matches state-of-the-art performance** of the exact method, which effectively mitigates the cubic complexity of the exact method.

----
----

## 1. On Chimera's computational cost 

### (A) The implementation of different variants with their costs

We recall that 

- *Exact Method (Section 3.3)*: Computes the exact resolvent via matrix inversion.
- *Approximate Method (Section 4.1.2)*: Approximates the resolvent using a power series expansion up to the graph diameter.
- *DAG Method (Section 4.1)*: Decomposes graphs into multiple DAGs.


For any graph topology with the number of nodes as V and the number of edges as E:


| Method       | Theoretical FLOPs      | Our Current Implementation          | Implementation FLOPs      | Parallel Complexity (Depth) |
|--------------|-------------------------|--------------------------|----------------------------|------------------------------------|
| Exact        | $O(V^3)$               | `torch.linalg.inv`       | $O(V^3)$                  | $O(\log(V))$                      |
| Approximate  | $O(V^2 \log(dia(G)))$  | Fast squaring trick      | $O(V^2 \log(dia(G)))$     | $O(\log(dia(G)))$                 |
| DAG          | $O(V + E)$             | Fast squaring trick      | $O(V^2 \log(dia(G)))$     | $O(\log(dia(G)))$                 |


----

### \(B\) Chimera is motivated by SSMs' inherent topology rather than linearity

Chimera identifies and leverages a fundamental property of SSMs: **their ability to directly operate on line-graph topologies without relying on position embeddings**. While SSMs are often regarded for their computational efficiency compared to Transformers, this is not the primary reason why Chimera generalizes SSMs. Instead, Chimera represents a methodological advancement over prior approaches that either "flatten" data into sequences to use models like Mamba or apply topology-agnostic models such as self-attention along with domain-specific position embeddings. By leveraging this property, Chimera's unified approach generalizes across diverse modalities like language, vision, and graphs, achieving state-of-the-art results while avoiding the need for designing task-specific embeddings, such as RoPE embeddings for language [5], absolute embeddings for vision [7], or Laplacian embeddings for graphs [8].

----

### \(C\) Like Mamba, developing a fast kernel for Chimera is an important future work.


Our primary goal in this work was to introduce a methodological shift for modeling data with an underlying graph topology, with a lesser emphasis on implementation optimizations. While Chimera’s current naive PyTorch implementation is slow, there is substantial potential to reduce wall clock time through hardware-aware kernels. Developing such kernels, however, is a non-trivial task. For instance, the first version of Mamba [1], which used a parallel scan algorithm for sequence mixing, was 15x slower than FlashAttention-2 [3, 4] (state size=256, sequence length<1000) and required over 2,400 lines of CUDA code. Mamba achieved parity with FlashAttention-2 only in its second iteration, Mamba-2 (SSD) [2], which involved over 6,400 lines of Triton code developed by the creators of FlashAttention. Given that Chimera is a more involved method compared to Mamba-2, we view the development of a fast kernel as an exciting opportunity for future works.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733207208045

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experiments

comment:
We thank all the reviewers for their thorough engagement with our work and for their insightful suggestions to improve our manuscript. We also sincerely appreciate their patience as we worked to secure the necessary resources for the suggested pretraining experiments. 

Based on our understanding, a common recommendation from the reviewers is providing a wall-clock time comparison between Chimera and Transformer-based baselines. While we again emphasize that the primary focus of our work is introducing a conceptual methodological change, we fully agree that an analysis of Chimera's runtime is valuable as it would not only support Chimera's broader adoption but also inspire future directions. 

To this end, we benchmark Chimera's wall-clock time against the Transformer baseline using the same architectural configurations as in our original Peptides-Struct experiment. In this realistic setting, we observe that Chimera's wall-clock time is within 1.8x of the Transformer baseline. Furthermore, we show that the wall-clock time is robust across the number of heads and model dimension. We believe that this reasonable performance offers a foundation for researchers to explore Chimera's impact on novel domains such as protein folding, where incorporating the graph topology more rigorously is essential for strong performance. 

Additionally, our benchmarking identifies key optimization areas for future work; currently our implementation is intentionally simple and it incurs an overhead of about 28% on bookkeeping operations including graph-processing tasks such as converting between sparse and dense matrix formats--an overhead that can be significantly reduced through optimized implementations. Beyond this overhead, we see significant potential for optimization through dedicated kernels targeting core operations like graph inversion. For instance, in the case of directed acyclic graphs (DAGs), we believe tailored kernel implementations can reduce the complexity of graph inversion from cubic to linear time. Additionally, fusing the various operations within the block—such as local convolutions, linear projections, skip connections, and the main Chimera processing unit—offers significant optimization potential similarly to the approach used in Mamba. 

Furthermore, we conduct additional experiments to isolate the effect of depth on Chimera's performance in response to reviewer uXiv's suggestion. In this experiment, we pretrain parameter-matched models: a 6-layer Chimera, a 6-layer BERT, and a 12-layer Chimera on the bidirectional language modeling task. Our results indicate that the validation performance of the two Chimera models is nearly identical, and they both significantly outperform BERT. This demonstrates that superior performance is driven by its methodological contributions rather than the increased number of layers.

We once again thank all the reviewers for their valuable suggestions, which we believe have significantly contributed to improving our manuscript.

---------
---------

### Experiment 1: Chimera's wall-clock time analysis


**A. Inference Time Experiments**
We begin by measuring the *per-layer,* *inference wall-clock time* of the approximate variant of Chimera and GPS-Transformer on the Peptides-Struct dataset. For this experiment, we adopt the architectural configurations from our state-of-the-art results--using a model dimension of 64 and 4 attention heads--ensuring parameter parity across layers for both models. We observe that the **time-ratio**, defined as $\frac{\text{time taken by Chimera}}{\text{time taken by GPS-Transformer}}$, is $\sim 1.5\times$, which we believe is reasonable for researchers to further explore Chimera’s potential across novel domains.


| **Method**    | **Wall-Clock Time Mean (ms)** | **Wall-Clock Time Std-Dev (ms)** |
|---------------|-------------------------|------------------|
| Chimera       | 11.552        | 2.831 |
| Graph-GPS     | 7.665        |  2.371|
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:**

This paper proposes a new state-space models (SSMs) for more general data beyond sequences, including image, text and graph data. The authors devise the new method using resolvent operations on adjacency matrices, achieving strong performance with theoretically good time complexity. One limitation is that the proposed method requires a directed acyclic graph (DAG) representation of data. Hence, the data should be converted into DAG and also empirically the efficiency of the proposed SSM is not observed to surpass a baseline transformer. 

**Strengths:**

1. **Improved applicability of SSMs.** This paper proposed a new SSM model that can incorporate the topology of data into SSMs. This expands the operating range of SSMs beyond 1D sequential data.
2. **Theoretically grounded modules.** The theorem suggests how to optimize the implementation and the normalization scheme is derived on theoretical grounds. 

**Weaknesses:**

1. **Limited efficiency.** The main motivation of SSMs is the superior efficiency compared to Transformers enjoying the linear time complexity. However, this paper failed to prove its advantage in empirical efficiency over Transformers. The additional experiments posted by the authors show that the wall-clock time analysis failed to show the efficiency improvement of the proposed method. On average, the proposed method is 1.5 times slower than Graph-GPS.
2. **Limited topology/overclaim.** Unlike the authors’ claim, the proposed method was demonstrated on only three modalities: image, text, and graph. The proposed method is not general enough to handle arbitrary topology such as well-studied groups [1,2,3] referred to in the comments of **Reviewer yGvr**.

**Main reasons:**

This paper studies an interesting and emerging topic that aims to extend the operating range of SSMs with strong performance and potential efficiency. However, the authors failed to demonstrate the empirical gain in efficiency compared to Transformers. Considering the fact that the main motivation of SSMs is the superior efficiency over Transformer, the inferior wall-clock time remains a main concern that has not been resolved by the authors.