Chimera: State Space Models Beyond Sequences

We thank the reviewer for their positive feedback and for recognizing Chimera as a fundamental model architecture that demonstrates strong performance across vision, graph, and natural language processing tasks, as well as it being a significant contribution to the ML community. Below, we address their questions and concerns:

1. Why are the number of layers in Chimera 2x that of ViT? Is this a fair comparison?

The concept of a "layer" differs between Transformers and SSMs. In Transformers, a layer consists of an Attention block and an FFN block. In Mamba, a layer corresponds to a single MambaBlock. A Transformer block contains twice the number of parameters as a MambaBlock and therefore the effective number of layers in both the models are the same with two SSM layers being equivalent to one Transformer layer.

For example, Pythia-2.8B [1] uses 32 Transformer layers, while Mamba-2.7B [2] employs 64 S6 layers—resulting in equivalent parameter counts and FLOPs. Hence, this is a fair comparison, as both models are controlled for computational cost and aligns with community standards for evaluating SSMs.

2. Can authors include a discussion on the computational efficiency of Chimera?

We refer the reviewer to sections 1A, 1B, 1C of the global response for a discussion about the computational efficiency of Chimera.

3. Do larger Chimera models underperform compared to smaller models? Can the authors verify it on NLP and CV tasks.

In our experience, larger Chimera models consistently outperform smaller models across CV and NLP tasks. Specifically,

a. NLP Task

On the GLUE task, Chimera (DAG) with 110M parameters achieves a GLUE score of 83.93%, outperforming the smaller 70M model with a GLUE score of 80.6%.

b. CV Task

On the ImageNet-1k classification task, Chimera's 88M parameter model achieves 81.4% Top-1 accuracy and outperforming the 22M parameter model, which achieves 77.8% Top-1 accuracy.

4. Is the optimization landscape of Transformers (ViT/BERT ) better than that of SSMs (Chimera). Could the authors add a figure on training loss to compare them?

In our experiments, we did not observe any deterioration in the optimization landscapes of Chimera. We have attached the training loss curves and the eval metrics curves for BERT-B (110M) and Chimera (DAG, 110M) in the supplementary.

References

[1]: Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. ICLR, 2022.

[2]: Stella Biderman et al. Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling

We thank the reviewer for their positive review and for their appreciation of Chimera’s ability to efficiently incorporate graph structure without relying on positional encodings, as well as its strong performance beyond graphs. We address their concerns below:

1. Can any graph architecture can be (over)stated as being applicable across modalities?

We acknowledge the reviewer's perspective, but we would like to clarify why the claim of applicability across domains is warranted for Chimera and not so much for other graph methods. Chimera uniquely demonstrates state-of-the-art performance across language, image, and graph tasks using the same unified architecture. This is unlike methods like GCN and GAT, which, despite their similarity to CNNs, often fail to surpass CNN-based baselines on image tasks, since they suffer from oversmoothing as the number of layers increase [1, 2]. This consistency in strong performance without requiring domain-specific modifications suggests that Chimera is a step in the right direction in effectively capturing the inductive bias of topology.

2. Can the authors compare Chimera against attention-based GNNs (like GAT and Graph Transformers) or efficient GNN variants?

a. Comparing Chimera and Attention-based GNNs

Among attention-based methods, GAT [2] employs attention locally on neighboring nodes, making it less effective at capturing long-range dependencies. Graph Transformers [3] address this by using topology-agnostic global attention but heavily rely on positional encodings to model graph topology. GPS-Transformer [4], which is a more recent iteration, combines local message-passing with global attention but still depends on positional encodings to capture graph structure.

In contrast, Chimera adopts a fundamentally different approach: It generalizes Mamba with its selection mechanism for long-range dependencies into directly incorporating the underlying topology via the resolvent operator. This approach not only eliminates the need for position embeddings but also achieves state-of-the-art performance across diverse domains, unlike graph-specific methods.

b. Comparing Chimera and Efficient GNNs

Efficient GNN variants, such as Linear GCN [1], aggregate features across $k$-hop neighborhoods by collapsing $k$ layers into a single operation, $H = \hat{A}^k X$, where $\hat{A}$ is the normalized adjacency matrix. However, Chimera differs in key aspects:

1. It aggregates features per layer, combining global context at each step rather than merely expanding the receptive field as in Linear GCN.
2. Unlike Linear GCN's unweighted adjacency matrix, Chimera parameterizes $\hat{A}$ as a generalization of Mamba-2 with its selection mechanism for modeling long-range dependencies.
3. Chimera models its outputs as $(L \odot CB^T)V$ (Sec 3.3), integrating graph topology within a Transformer like key-query-value framework which is widely recognized for strong performance across domains.

These differences allows Chimera to achieve state-of-the-art performance across a wide array of modalities whereas Linear GCN is primarily used for graph data.

We thank the reviewer for highlighting this and will include these comparisons in the next version of our paper.

3. Can the authors include a discussion on the computational complexity of the method?

We refer the reviewer to Sections 1A,B,C of the global response where we have included a discussion on the computational cost of Chimera.

References

[1] Rusch, T. Konstantin, Michael M. Bronstein, and Siddhartha Mishra. "A survey on oversmoothing in graph neural networks." arXiv preprint arXiv:2303.10993 (2023).

[2] Zhao, Lingxiao, and Leman Akoglu. "Pairnorm: Tackling oversmoothing in gnns." arXiv preprint arXiv:1909.12223 (2019).

[3] Cai, Chen, and Yusu Wang. "A note on over-smoothing for graph neural networks." arXiv preprint arXiv:2006.13318 (2020).

[4]: Wu, F., Souza, A., Zhang, T., Fifty, C., Yu, T., & Weinberger, K. (2019, May). Simplifying graph convolutional networks. In International conference on machine learning (pp. 6861-6871). PMLR.

[5]: Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, Yoshua Bengio. Graph Attention Networks.

[6]: Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, Hyunwoo J. Kim. Graph Transformer Networks.

[7]: Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, Dominique Beaini. Recipe for a General, Powerful, Scalable Graph Transformer.

We thank the reviewer for their positive feedback and for recognizing Chimera as a principled method with well-motivated, theoretically grounded method normalization as well as its competitive results across various tasks. We now address their concerns:

1. The insight of SSMs operating on directed line topology is obvious and holds trivially.

We agree that the idea of SSMs operating on line graph topologies due to their recurrence is indeed intuitive—this was the starting point for our work as well. However, Chimera is the first work to formalize this connection mathematically, which is a crucial step that enables the generalization of SSMs beyond sequences to general graph topologies across domains.

2. Could the figure 1 be improved to reflect the method?

We appreciate the reviewer's suggestion and we will provide a clearer illustration of our approach in the next revision.

3. Can Chimera handle Lie groups?

Chimera currently focuses on modeling data topologies in the graph sense and does not address equivariance or invariance to specific group operations.

We believe the confusion arises from the term "topology," which is overloaded in the literature. In this work, we use topology in the graphical sense, treating nodes as abstract objects connected by edges that represent neighborhood relationships. This is in line with past works such as GraphGPS+Transformer [1], GINE [2], and GatedGCN [3]. We chose the term "topology" to emphasize that Chimera is not limited to traditional graph data but is designed for any modality with an underlying notion of neighborhoodness, such as images, language, and audio.

We believe the reviewer may be referring to a different line of work where topology is understood in the algebraic sense, as in LieTransformer [5] and SE(3)-Hyena [6] that focuses on equivariance or invariance to group operations,. Nevertheless, we note that since Chimera operates on the underlying graph topology, it is inherently invariant to transformations that preserve the graph structure, such as permutations of nodes.

We appreciate the reviewer’s feedback and welcome their suggestions on how to make this distinction clearer. While we currently rely on context to differentiate these meanings, we are happy to revise the manuscript for greater clarity.

4a. For general graphs is the SOTA achieved via the exact cubic method. And the exact version to work much better than the approximated one which limits practical utility.

4b. Can the authors include a discussion on the computational complexity of the method?

a. The reviewer is correct that the earlier state-of-the-art results for general graphs were based on the exact cubic method. However, we note that we have since improved the approximation method, which now surpasses all baselines and achieves performance on-par with the cubic method on Peptides-Func and Peptides-Struct, effectively mitigating the cubic complexity of the exact method. We refer the reviewer to Section 2 of the global response for details.

b. We have included a discussion on the computational cost of Chimera in Sections 1A,B,C of the global response.

5. What happens when k < dia(G) in the approximation method?

We thank the reviewer for suggesting this ablation and we have now included it in our response below. Our results show that the model is robust to using smaller values of k. Additionally, we note that the computational cost of Chimera has a weak dependence on k, scaling only logarithmically with the graph diameter.

Setting: We evaluated the effect of using smaller values of approximation degree k on the Peptides-Struct dataset, with k = Dia(G)/factor, and report MAE over one seed due to time constraints:

We observe a small drop in performance as factor increases from 1 to 2, but overall the results are robust to smaller values of k. We believe that this is an important ablation and we will definitely include it in the next revision of our manuscript.

6. Is it possible to extend the proposed method to incorporate edge features when provided?

We apologize for not including this detail in the current manuscript, and we would like to clarify that our reported SOTA results already incorporate edge features using the following approach. Specifically, in equation 11 in the paper, we add edge features as an additional feature to the calculation of $A_{ij}$ where,

$A_{i, j} = \frac{\exp(-\Delta_i) + \exp(-\Delta_j) + \exp(-\Delta_e)}{3}$

where $\Delta_e = \text{Softplus}(f_e(E))$ and $E$ are the edge features, and $f_e$ is a projection. We will update the manuscript with this detail in the next revision.

Thanks for the response.

I believe a detailed benchmark on the inference speed of Chimera compared with other efficient baselines is required for publication, e.g., Seq length v.s. Time, like the Fig. 10 in Mamba2's paper. Otherwise, you can hardly claim this is still an efficient model given that Chimera involves many computations that do not occur in the original Mamba or efficient transformer architectures.

In addition, correct me if I'm wrong, given your response, I'm concerned/confused that if you cannot implement this method with linear complexity in practice, how could you really claim efficiency in this manuscript? And if this model is not efficient practically, then it essentially has lost one of the core benefits of SSMs or Mamba, that is, efficiency? The paper of Mamba also includes some tricky computation, and one reason they can be accepted (eventually) is that they also propose a way to implement that tricky computation with linear complexity. And even they have included such implementation, they were rejected at ICLR. Therefore, it seems to me it should also be required for this manuscript to provide an efficient implementation of Chimera for acceptance. Otherwise, it can be hard to vote for acceptance of this manuscript, because it sounds a bit weird if one paper claims to extend SSMs/Mamba but essentially cannot implement their methods with linear complexity like SSMs/Mamba and have claimed efficiency in their manuscript. And if there is no computational gain in Chimera, shouldn't one expect it to be compared with more advanced transformer models, like those LLMs?

Please do correct me if I misunderstand anything. If I'm not mistaken and the practical efficiency concern cannot be resolved, I may have to lower my rating.

We are glad that the reviewer found Chimera interesting and they recognized its strong empirical performance. We now address the concerns raised by the reviewer.

1. Can Chimera outperform Mamba2 on NLP tasks and Vision Mamba on CV tasks.

We note that in the causal language modeling setting with a directed line graph topology, Chimera reduces to Mamba2 (Prop 3.3), and it therefore has an equivalent performance to Mamba2 on causal NLP tasks. Furthermore, on bidirectional NLP tasks, we show that Chimera outperforms attention-based BERT baseline by a GLUE score of 0.73 (Table 1)

We cannot directly compare the results between Chimera and Vision Mamba because Chimera uses the standard training recipe used in ViT [5], Monarch Mixer [4] whereas Vision Mamba uses the DeiT [3] training recipe, introduced for distillation-based models. However, we note that the underlying module of Vision Mamba corresponds to the Fwd&Rev setting in our Table 3 ablation, where Chimera's 2D-DAG structure demonstrates superior performance.

2a. How is the DAG Decomposition method implemented in practice?

2b. For general graphs is the SOTA achieved via the exact cubic method? The approx method has decayed performance.

2c. How are other variants namely the exact method, the approx method and the DAG method implemented?

2d. Would the authors provide code?

2e. Concerns about the practical efficiency.

We kindly refer the reviewer to our global response for a detailed treatment of the computational aspects of Chimera. For convenience we summarized some of our responses here:

a) The DAG Decomposition method is implemented as follows:

1. For line graph topology, we use Mamba2 kernels with a linear time complexity to implement the DAG decomposition method.
2. For the general DAG method, we use the fast squaring trick, which, as the reviewer correctly noted, is more hardware-friendly than Gaussian elimination. While our method has linear theoretical complexity, in practice, the fast squaring trick is more efficient for the current datasets. However, for sufficiently large graphs, Gaussian elimination might be faster.

b) The reviewer is correct that the earlier state-of-the-art results for general graphs were based on the exact cubic method. However, we note that we have since improved the approximation method, which now achieves performance on-par with the cubic method on Peptides-Func and Peptides-Struct, effectively mitigating the cubic complexity of the exact method. Please refer to Section 2 in the global response for details.

c) The implementation details for the exact, approximate, and DAG methods are provided in Section 1A of the global response.

d) We will release the code in the next revision of our paper to ensure transparency and reproducibility.

e) We address practical efficiency in Sections 1B and 1C of the global response

3. How can the model deal with edge features for graph-structured data?

We apologize for not including this detail in the current manuscript, and we would like to clarify that our reported SOTA results already incorporate edge features using the following approach. Specifically, in equation 11 in the paper, we add edge features as an additional feature to the calculation of $A_{ij}$ where,

$A_{i, j} = \frac{\exp(-\Delta_i) + \exp(-\Delta_j) + \exp(-\Delta_e)}{3}$

where $\Delta_e = \text{Softplus}(f_e(E))$ and $E$ are the edge features, and $f_e$ is a projection. We will update the manuscript with this detail in the next revision.

• First of all, I believe efficiency is important, especially for SSM-related methods as it is the main motivation of these methods. Being scientifically valuable does not mean it is ready for publication. And it is because of its scientific value that I'm still rating it 5. If we do not talk about efficiency at all, the key contribution of this work is mostly an observation of the topology SSMs working on, which is an interesting but actually somewhat intuitive observation.

• Second, one reason that I want to lower the score is that I believe the current manuscript is in some sense overclaiming, which needs to be revised accordingly. As I said, I find it "sounds a bit weird if one paper claims to extend SSMs/Mamba but essentially cannot implement their methods with linear complexity like SSMs/Mamba and have claimed efficiency in their manuscript"

  • For example, in the abstract, the authors claim "Furthermore, being topologically aware enables our method to achieve a linear time complexity for sequences and images". There are similar statements across the manuscript.
  • These statements may mislead readers (like me). When I first read these statements, I believed this method was of course implemented with linear complexity and should be an efficient method, and it's just they'd need a speed benchmark to improve the manuscript. However, it turns out that this method can't be implemented efficiently right now, which greatly lowers the contribution of this work.

We sincerely thank the reviewer for their continued thoughtful engagement with our manuscript.

We would like to affirm that the statement and the proof of Proposition 4.4 are accurate. The reviewer is right that for general sparse matrices, the complexity of Gaussian elimination typically scales as $O(n \cdot nnz(I-A)$. However, for a Directed Acyclic Graph, the associated adjacency matrix is a lower-triangular matrix, and this structural property reduces the complexity of Gaussian elimination to $O(nnz(I-A))$.

To illustrate this, we run Gaussian elimination by successively choosing the pivots as (i,i) for each i in the ordered list {$0,..,|V|-1$}. Observe that because the matrix is lower-triangular, when we eliminate with respect to the pivot (i,i), the remainder of row i is composed of only zeros. Therefore, the elimination operation only needs to be performed within the corresponding column instead of the entire row. This reduces the number of required arithmetic operations to $O(nnz(A[:, i]))$, where $nnz(A[:, i])$ denotes the number of non-zero elements in the column $i$. Summing over all the columns, the total complexity becomes $O(\sum_{i}^n nnz(A[:, i])) = O(nnz(A))$. In the context of Proposition 4.4, where the adjacency matrix of the DAG has $O(|V| + |E|)$ non-zero elements, we get the final claim in our proposition that the complexity is $O(|V| + |E|)$. We also note that for our motivating example of Mamba, the $I-A$ matrix contains exactly $2|V|$ non-zero entries, and by our theorem has a linear time complexity.

We will add these details in the original proof of Proposition 4.4 and will ensure they are clearly included in the final version of the paper. Thank you for helping us clarify this important point.

We are grateful to the reviewer for raising their concerns about the presentation of our paper, particularly regarding the complexity of the general DAG method. We recognize that the current draft may give the impression that our method is both theoretically and practically linear time for DAGs. In response, we have updated the manuscript with an explicit limitations section clarifying that, while the general DAG method is theoretically linear in complexity, our current GPU implementation relies on a quadratic fast-inverse trick. We also note in the limitations that developing a dedicated linear DAG kernel remains an important direction for future work, along with other potential optimizations discussed with the reviewers. We apologize for any confusion and thank the reviewer once again for their valuable feedback.

2. Improved results on the Long Range Graph Benchmark (LRGB)

We have further tuned our experiments on the LRGB dataset, and the approximation method now achieves state-of-the-art (SOTA) performance on Peptides-Func and Peptides-Struct, matching the exact method within standard deviations.

The updated results are summarized in the table below:

Since the Approximation method achieves results on par with the exact method, it effectively alleviates the cubic complexity of the exact method and can be used instead of calculating the exact resolvant.

The changes in the hyperparameters for both the datasets are as follows:

References

[1]: Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. ICLR, 2022.

[2]: Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024b.

[3]: Tri Dao and Albert Gu. State Space Duality (Mamba-2) Blogpost Link

[4]: Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

[5]: Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023

[6]: Hugo Touvron et al. Llama: Open and efficient foundation language models, 2023.

[7]: Alexey Dosovitskiy et al. An image is worth 16x16 words: Transformers for image recognition at scale

[8]: Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems

B. Time-Ratio is Robust to Architectural Variations

We further conduct a series of experiments to test how the time-ratio changes by varying different hyperparameters in our model. The hyperparameters we ablate are: model dimension (d_model), number of heads, and batch-size. Note that we define the setting with batch size = 1 setting as evaluating the computational speed of the model during "inference mode". We find that across these hyperparameters, the time-ratio between the two models remains consistent, demonstrating the robustness of our analysis to variations in architectural design and training configurations.

Table 1: Ablation over Model dimension

Table 2: Ablation over Number of Heads

Table 3: Ablation over Batch Size

C. Bottleneck analysis

As we had noted in our global response, the current implementation of Chimera has a higher wall-clock time compared to Transformer-based baselines. To better understand this disparity, we analyzed where Chimera incurs the most overhead, which offers valuable insights and avenues for its optimization:

We remark that we kept our current implementation purposefully simple and therefore a significant portion of the runtime (approximately 28%) is consumed by bookkeeping operations like converting the graph's original sparse adjacency list representation into its dense matrix form—implemented using torch_geometric.utils.to_dense_batch, torch_geometric.utils.to_dense_adj—which we use for matrix normalization, inversion and subsequent matmuls. While these conversions simplify the implementation, they are computationally expensive: Creating a dense representation on the fly involves a lot of memory copy operations with a substantial overhead. This overhead is exacerbated by irregular batch sizes as the resulting dense matrix must accommodate the largest graph in the batch.

Future Directions: While this work focuses primarily on validating Chimera's methodology, through the above analysis we identify pathways to significantly improving runtime performance:

1. Preprocessing data: Precomputing batches with similar graph sizes along with their dense representations can reduce the overhead of on-the-fly dense-to-sparse conversions.
2. Direct Sparse Matrix Operations: Using dedicated CUDA kernels from cuSPARSE can help avoid materializing dense matrices, and instead could operate directly on sparse adjacency lists to compute $My = x$, solving for $y$ without explicitly materializing $M$.
3. Dedicated kernels for DAGs: As we noted in our global response, we believe that tailored kernels for directed acyclic graphs (DAGs) can significantly reduce Chimera's practical complexity due to the structured nature of the problem.
4. Fused Operations: Fusing operations within the Chimera block—such as local convolutions, linear projections, skip connections, and the main processing unit—can further improve execution speed, similar to the optimizations used in Mamba.

We believe that these optimizations represent an exciting avenue for future work. Our work showcases Chimera’s methodological promise and this analysis provides actionable insights that can help improve its runtime.

D: Wall-clock time Reduces with Coarser Approximation

We study the impact of the degree of approximation on wall-clock time. Specifically, we define the log-approximation factor such that the largest power to which the adjacency matrix $A$ is raised is determined by $\lceil \text{diameter}(G) / 2^{\text{factor}} \rceil$. As the approximation factor increases, our results indicate the expected reduction in the runtime of the fast-inverse operation relative to the full-diameter fast-inverse approximation.

We note that the matrix inversion currently accounts for around 15% of the total runtime, with the majority of the computational overhead arising from graph-processing operations. Nevertheless, as our prior experiments indicate that increasing the approximation factor does not significantly degrade performance, this highlights another avenue to reduce computational cost while maintaining performance.