PICASO: Permutation-Invariant Context Composition with State Space Models

Putting other things aside, it's really interesting to see you asking about KV cache in an SSM paper (even if the state is analogous to KV cache). If you haven't read the paper, please lower your confidence.

We thank the reviewer for their detailed feedback and constructive suggestions, which we have incorporated in our revision.

Figure 3,4 (left) has an error. There are seven legends and only six curves. The main approach PICASO-R is missing from the figure. Though I infer from the context of different places (e.g., PIConcat can not run for 10 chunks) that it is a mistake that PIConcat-R's curve should actually be PICASO-R? But such a mistake seriously lowers the quality of the paper.

There is no mistake in the figure, but the curves of PICASO-R and PICASO-S actually overlap and are hence hard to see. We have made this explicit in the revised caption, and we thank the reviewer for noticing the readability issue.

The differences between BPTC and BP2C are very ambiguous (And I actually do not know the difference). Equations are prefered. And stop-gradient notation can be used if you need to clarify the difference.

We thank the reviewer for their suggestion and notation advice, and have updated our revision with the equations in Sec. 5 to provide better clarity.

We copy the definitions below, where we replace $\boldsymbol{u}_i$ with $v_i$ due to issues with rendering in Openreview

$\mathcal L_{BPTC}(\theta) = \sum_{(v_i,u_i,S_i) \in \mathcal{D}}L_{\rm CE}(f_\theta(v_i, x^{PICASO}(S_i)), u_i)$

and

$\mathcal L_{BP2C}(\theta) = \sum_{(v_i,u_i,S_i) \in \mathcal D} L_{\rm CE}(f_\theta(v_i, \operatorname{sg}\left[x^{\rm PICASO}(S_i)\right]), u_i)$, where $\operatorname{sg}$ is the stopgradient operator.

Clarify usefulness outside of Mamba-1. The PICASO-S relies on the commutativity of A_i, and the PICASO-R relies on the invertible property of A_i. Despite that Mamba-1 clearly holds these properties. It is not very clear whether other SSMs (e.g., RWKV, RetNet, Mamba-2, etc.) hold these properties. This is not a very important limitation of this paper, since work can also be done for other architectures, but it would be interesting to see the authors' analysis in this paper.

Our method works for both Mamba-1 and Mamba-2 since as rightfully pointed out by the reviewer, the $A_i$ matrices are diagonal hence commutative/invertible. Our experiments are performed using Mamba-2 2.7B. These properties are often satisfied by several recurrent mechanisms in order to make them computationally efficient, including RWKV (via element-wise scaling) and RetNet (where $A$ is parameterized as a diagonalized matrix).

4. It seems that CASO is also the original proposal of the paper. Since I am not very clear about these baselines (But I do know RAG, Mamba, etc., clearly). I think it would be beneficial to mention methods similar to it more.

Thank you for the suggestion, we added links to relevant literature when introducing our baselines in Sec 6.2. which we hope addresses the reviewer's concern.

Given the approaching rebuttal deadline, we wish to follow up on our response to your insightful feedback. We hope that our responses, revised paper draft, and additional experiments have addressed the areas of weakness you identified. We would be grateful if you could consider providing stronger support for our paper's acceptance, and thank you for your very valuable feedback and suggestions!

We thank the reviewer for their feedback, and address points of concern below.

While PICASO achieves near-concatenation performance in zero-shot scenarios, there is still a minor but noticeable gap in accuracy. This limitation could impact its suitability for applications where slight accuracy improvements are critical.

Our method is indeed targeted towards deployment of LLMs, where slight performance trade-offs are well justified by the substantial speed-up obtained. However, we note that trade-off can be easily avoided via our proposed fine-tuning method, obtaining the "best-of-both-worlds". Indeed in situations where fine-tuning is not an option and inference-time is a non-issue, concatenation remains the paragon.

Although PICASO performs well with up to 10 document chunks, it is unclear how it scales with larger document sets. This potential bottleneck is worth further investigation.

This is a good point, and we have addressed this in Figure 6 of our main paper. While concatenation stops working as number of documents increase, due to the model becoming unstable once context length is exceeded, the performance of PICASO remains relatively stable even when composing up to 50 document chunks.

We thank the reviewer for their detailed feedback and constructive suggestions. We address individual concerns below:

This is similar to in-context retrieval [1] (seems a missing reference). However, there are no real RAG applications studied in the paper. Appendix B.4 presents some additional results on “language modeling” (which I don’t think is a proper description), but not so much details are provided such as database and retrieval.

The evaluation metric mostly focuses on log-perplexity or loss, while it is also desirable to demonstrate generation capabilities of the proposed method for real tasks.

While RAG is indeed our main motivation, the paper focuses on designing a method to quickly retrieve and integrate information in the state of the model. We found that perplexity-based measurements provide a sufficiently robust metric of success in this respect. However, we agree that exploring this aspect of evaluation is also important, and have added Appendix B.6. evaluating accuracy for the OpenbookQA task under the retrieval-augmented setting. In particular, we show that the same trends hold as compared to perplexity metrics -- that augmented generation benefits performance, while PICASO provides similarly boosted performance with around $8\times$ reduced computational costs.

We also thank the reviewer for the reference [1] which we have added, and have changed "language modeling" to "LLM evaluation tasks".

Furthermore, there could be certain baselines missing, including similar approaches proposed in the previous literature, such as the closely related work described in Section 2.

To the best of our knowledge, our work is the first SSM-specific method developed for RAG-like settings. We believe we have included all applicable baselines for composing documents with SSMs. However, we realized that we did not explicitly label their sources in Sec 6.2, and have corrected that in the revision. Thank you for pointing this out, and we are also happy to include additional appropriate baselines that the reviewer recommends.

Other ablations could include the effect of retrieval accuracy, the impact of document or chunk lengths to be composed together, etc.

To address the reviewer's concerns, we have added an ablation study on retriever choice in Appendix B.5. Figure 6 of the Appendix also demonstrates what happens when we scale beyond the training context length of the model. We have also provided a histogram of document chunk statistics in Appendix B.7. to provide greater clarity on the distribution of chunk lengths considered.

Not enough background information is provided for RAG applications with SSMs, as most of the RAG applications are built with Transformers. This is also related to the lack of experimental studies mentioned above, where more comprehensive comparisons on other tasks/methods could be beneficial.

To the best of our knowledge, our work is the first to introduce an approach towards RAG that is specific to SSMs. We have included [1] in our discussion on related works and the baseline method (concatenation) in the revision to better frame our work in the context of transformer-based applications.

The method is based on some strong assumptions such as the order of the retrieved document chunks does not matter, so that an average state can be used for inference. There is no clear evidence provided.

In our experiments in Sec. 6, we compare our permutation invariant method with ordering the documents by relevance (i.e., most relevant documents are closer to the answer), which we found empirically to perform best. While, to the reviewers point, such ordering is better than random ordering, we still observe that incorporating permutation invariance via PIConcat / PICASO significantly outperforms their ordered counterparts Concat / Soup respectively.

Thank you for the clarification.

Indeed there might be some misunderstandings, the first of which we believe stems from distinguishing between (a) intuition for our claim on the desirability of permutation-invariance among retrieved documents, and (b) empirical evidence in support of this claim.

(a): This comes simply from the fact that retrieved documents are conditionally independent of one another, by construction. Hence, similar to how ordering in multiple choice questions should not influence the final answer, neither should the relative ordering among independently retrieved documents.

(b): PIConcat / PICASO significantly outperform their ordered counterparts Concat / Soup respectively in all our experiments, providing strong empirical evidence on the advantages of incorporating permutation-invariance.

fundamental assumption ... (that) independent documents can work as well

there is no clear evidence to show this is a good approximation to how the model actually works by using the contextual information.

We believe this to be another main source of misunderstanding. We do not compose documents independently of one another. As an example, given documents A and B, we do not compose {A,B}. Instead, PIConcat / PICASO efficiently composes {A$\cdot$B , B$\cdot$A}, where $\cdot$ denotes exact / approximate concatenation respectively. As such, contextual information is incorporated within both A$\cdot$B and B$\cdot$A, while at the same time, the composed state remains invariant to document ordering.