PICASO: Permutation-Invariant Context Composition with State Space Models

Authors: Tian Yu Liu, Alessandro Achille, Matthew Trager, Aditya Golatkar, Luca Zancato, Stefano Soatto
Primary area: foundation or frontier models, including LLMs
Keywords: State Space Models, Composition, Retrieval

Abstract:
Providing Large Language Models with relevant contextual knowledge at inference time has been shown to greatly improve the quality of their generations. This is often achieved by prepending informative passages of text, or 'contexts', retrieved from external knowledge bases to their input. However, processing additional contexts online incurs significant computation costs that scale with their length. State Space Models (SSMs) offer a promising solution by allowing a database of contexts to be mapped onto fixed-dimensional states from which to start the generation. A key challenge arises when attempting to leverage information present across multiple contexts, since there is no straightforward way to condition generation on multiple independent states in existing SSMs. To address this, we leverage a simple mathematical relation derived from SSM dynamics to compose multiple states into one that efficiently approximates the effect of concatenating raw context tokens. Since the temporal ordering of contexts can often be uninformative, we enforce permutation-invariance by efficiently averaging states obtained via our composition algorithm across all possible context orderings. We evaluate our resulting method on WikiText and MSMARCO in both zero-shot and fine-tuned settings, and show that we can match the strongest performing baseline while enjoying on average $5.4\times$ speedup.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1730606973947

thread_type:
official_review

<thread_object>
creator: reviewer_h4zy

<content>
summary:
This paper proposes a way to speed up inference for RAG use cases for state space models (SSMs). The speed up comes from feeding in states of the retrieved chunks instead of doing inference from raw tokens to best utilize the properties of SSMs. To address the problem of combinatorial number of possibility in concatenation the authors proposed two ways of computing all the permutations -- exhaustive and cyclic. The authors demonstrate the proposed approach on two datasets, which shows improvements over naive concatenation of states. The performance compared to baseline concatenation is still worse but faster, the gap closes after finetuning.

strengths:
The proposed method shows significant speed up over the naive baseline

weaknesses:
The writing can be improved, in particular in the experiment section. It is not clear what the size of the chunk, what is the inference token size (how many in context and how many for inference), and how the evaluation is setup. For wikitext, it is not clear what is used as the query and what is used as the candidate pool for performing retrieval.

The evaluation feels incomplete. It seems that for both WikiText and msmarco are evaluated based on the test set loss. Since the proposed method fits in the RAG setting, it would be great to show the actual QA performance instead of the cross entropy as the result.

It is not clear what the speed baseline setup is. Does it use KV cache?

questions:
See questions in the above section
</content>
</thread_object>

</thread>
<thread>
creation date:
1730629697789

thread_type:
official_review

<thread_object>
creator: reviewer_2fh2

<content>
summary:
The paper addresses the challenge of efficiently incorporating multiple documents into the generation process of LLMs. Traditionally, concatenating documents leads to significant computational costs that scale with the number and size of document chunks. State Space Models (SSMs) offer a faster approach by encoding documents into fixed-size state vectors, but composing multiple states is not straightforward. The authors introduce PICASO, a method for permutation-invariant composition of document states using SSMs. PICASO efficiently retrieves and combines pre-computed states to condition the generation of high-quality outputs without the need for online document processing. It enforces invariance to the order of document presentation, which is desirable when the temporal ordering is uninformative.

strengths:
1. The issue is of significant importance. Generating content from multiple processed documents without re-preprocessing is a crucial requirement for both long context application and agentic memory systems.
2. The State Space Model constitutes a suitable architectural framework for the targeted problem, thereby rendering the study meaningful.
3. The introduced methodologies are technically robust. And they are presented clearly mostly (with a few exceptions).
3. The experiments are meticulously designed. Both performance and time complexity are analyzed. I particularly appreciate the experiments depicted in Figure 2, which elucidate the improvement brought about by the proposed method and provide mechanistic insights for the field.

weaknesses:
I will write the weaknesses in the order of the willingness of raising my score after addressing these weaknesses.

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

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

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

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.

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

</thread>
<thread>
creation date:
1730782585247

thread_type:
official_review

<thread_object>
creator: reviewer_fyyf

<content>
summary:
The paper presents PICASO, a method to enhance generation capabilities in Large Language Models by composing document states in a permutation-invariant manner using State Space Models (SSMs). This approach addresses the inefficiency and high computational cost of concatenating multiple document tokens by pre-processing document chunks into states. PICASO leverages a permutation-invariant composition strategy that enables LLMs to utilize multiple documents' information efficiently, achieving computational speed-ups and scalability, especially in retrieval-augmented tasks. Evaluation on datasets such as WikiText and MSMARCO shows that PICASO can match or closely approximate the performance of document concatenation with significantly reduced processing time.

strengths:
1. PICASO offers an impressive 5.4x speed-up in processing time over traditional concatenation, a practical advantage for real-world applications that require high-speed document retrieval and composition.
2. The model effectively maintains performance without relying on the order of documents, which is crucial when temporal or logical ordering is irrelevant. This represents a thoughtful design that accommodates various real-world scenarios.
3. Extensive experiments demonstrate PICASO’s effectiveness, comparing multiple composition methods and showing zero-shot and fine-tuned settings. The authors provide a comprehensive view of PICASO’s strengths and limitations across several benchmarks.

weaknesses:
1. 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.
2. 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.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1731281762000

thread_type:
official_review

<thread_object>
creator: reviewer_uyju

<content>
summary:
The paper focuses on the inference efficiency with state space models (SSMs) when conditioning on multiple document chunks in retrieval augmentation setups. It proposes a method to compose hidden states of different contexts into a single state for SSM generation, which is similar to compressing KV cache in Transformers (but for SSMs inference is only based on a single hidden state for one conditioning context). In particular, for applications that require retrieval, multiple document chunks could be provided as additional contexts, limiting the inference efficiency. Instead of the standard way of concatenating multiple document chunks in the context and run SSMs, the paper proposes to 1) preprocess individual chunks to store their SSM states; 2) retrieval the document chunks (with an additional retrieval model) and their pre-computed states based on a query; 3) compute an aggregated state from the different chunk states (which is not by simple averaging); and 4) continue generation based on the single aggregated state representing multiple document chunks.

The core contribution lies in step 3), where the authors derive an equivalent form (based on a single layer of SSM, so in general it is a heuristic) of state composition based on SSM computations to that of computed by naively concatenating multiple document chunks. Since this composition is reliant on a particular document chunk order, the paper further proposes to aggregate the composed states from different orders, assuming the order does not matter for the downstream generation. Efficient algorithms for the permutation-invariant aggregation are derived. Fine-tuning SSM by incorporating the aggregated states is also applied to further enhance the inference quality.

Experiments are conducted on language modeling tasks with the Mamba-2 2.7B model on two datasets, WikiText-V2 and MAMARCO. Results show that the proposed approach can achieve comparable perplexity to the baseline of directly concatenating multiple document chunks, while enjoying on average 5.4X speedup.

strengths:
- The paper presentation is mostly clear. The motivation, background, and techniques are well explained, also in good connection with relevant research.

- The paper studies an interesting problem, which is to speed up inference in the context of retrieval augmentation, specifically focusing on SSM models with their unique state dependencies. SSMs do not need KV cache like Transformers thus reducing memory and computation requirements. The paper further improves inference efficiency by delegating the document state computation to the preprocessing time, and only compute state compositions for inference.

- The technicality is sound (I did not check the math in every detail but overall they seem correct) and interesting. Algorithms of efficient computation of document state composition are derived, as well as for fine-tuning SSMs with the state composition that utilizes SSM model parameters.

- Experimental studies support the expected results of maintaining language modeling performance, measured in log-perplexity, while speeding up the generation from the composite SSM states from multiple document chunks.

weaknesses:
1. While the problem and methodology is interesting, I find the experimental studies are somewhat not adequate. 
- The authors conduct experiments with simplified settings of retrieving from WikiText and MSMARCO document chunks to improve next token prediction perplexity. 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.
- 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.
- Comprehensive ablation studies are missing, making it difficult to understand the experimental details such as what parameters in multi-layer SSMs are used and how that matters. Other ablations could include the effect of retrieval accuracy, the impact of document or chunk lengths to be composed together, etc.

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

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

4. Some of the descriptions are unclear. For example, some math notations are undefined, some details in Figure illustrations are missing. See more below in my questions.

5. It seems the gain of the proposed approach mostly comes from the inference efficiency side, as the log-perplexity (demonstrated in Table 1) is not improved over the baseline of document concatenation. However, the complete pipeline includes preprocessing document chunks for their SSM states (Figure 7 also shows this), and it also includes an additional retrieval model with a separate set of embeddings that need to be stored in the database. The data used for experiments are small. The considered number of chunks are smaller than 10. It is unclear how the proposed method performs when brought to larger scales requiring bigger storage and more computation in real applications.

*[1] Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham. 2023. In-Context Retrieval-Augmented Language Models. Transactions of the Association for Computational Linguistics, 11:1316–1331.*

Overall I think the paper presents an interesting method and study on Mamba states composition, but the empirical justifications are somewhat flawed. I am willing to increase my score after getting more insights from the authors regarding my concerns.

questions:
1. In line 084-085, “relative ordering is often uninformative”: this is a very strong assumption. Any evidence? Especially in applications where retrieval meets Mamba. In fact, many previous studies show the position of information matters [2].

*[2] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics, 12:157–173.*

2. For Mamba models of multiple layers, what hidden states and which layer parameters did you use such as the A and B matrices? Is there any difference depending on what you use?

3. In Figure 2, can you explain more on how to read the figure? In particular, how does the figure show that the CASO states are closer to one another? 

4. For experiments on val/test sets, what are the chunk databases? Do they include chunks in the training data, or just the chunks of corresponding val/test sets?

5. Since you use a different retrieval model, I assume you have to store the chunk vectors based on the retrieval model as well, on top of the SSM hidden states. Is that right? What is the effect on retrieval accuracy (can be ablated by using different retrievers) on the proposed approach? And moreover, can you directly use the SSM states for context matching for retrieval such as in non-parametric language modeling?

6. Log-perplexity (loss) is used for language modeling performance. Why not just report perplexity? I think it is more common for language modeling and easier for illustration, such as in Table 1.

7. For Table 1, what is the number of retrieved document chunks for the results?

8. Also for Table 1, it seems training on different domains of data does not help the perplexity. In many real applications, in-domain training data is hard to collect. This indicates that it is better to just apply the method at inference mainly for speedup? Model quality would not be improved here without training in-domain.

9. Line 492-493, Appendix B.4 and “language modeling”: these are not language modeling tasks, which are what you did in the main experiments and measure with perplexity. While I can understand what you mean, it is better to change the wording for clarity. In addition, can you provide more details on how these tasks are set up with retrieval and their evaluation?

10. The proposed method can “composition of information contained in up to 10 documents in a manner that is order-invariant.” Do you see this as a limitation? What happens after 10 documents? And is there any effect of the length of each document in the efficacy of composition? An ablation study would be ideal.


Some minor comments on typos and suggestions:

11. In the abstract or introduction, how is the name PICASO composed for what it stands for?

12. In proposition 3, explain the notation of `Id` and `[i]_n`?

13. Line 292-293, “We perform simple averaging to combine these tokens from different documents”: do you mean average states at the same token positions across different documents? So we still result in 4 states at the last 4 token positions after averaging.

14. Line 304-305, “on average improves only 8.5% compared to the baseline”: specify what is the evaluation metric or task briefly for better understanding since the experiments are not talked about yet.

15. Line 350-352, in the training objective, the equation computes loss on the single token prediction after the retrieved documents. For sequence predictions, is the intention to use the same formula but with different $\mathbf{u}_i, u_i$ and the same $S_i$ at different sequential positions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732422277609

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Given that it is nearing the end of the discussion period, we hope that the reviewers can let us know if we have adequately addressed their concerns and questions. We present an overall summary of our rebuttal below:

We are happy to hear that our work is well-received by the reviewers with respect to the practicality and "significant importance" of the problem (uYJU, FYyF, 2fh2), the "thoughtful design" and technical soundness/robustness of our approach (uYJU, FYyF, 2fh2), and our "meticulously designed" and "extensive experiments" (FYyF, 2fh2) showcasing the inference efficiency and "impressive 5.4x" speed-up obtained by our method (all reviewers).

We are very grateful to the reviewers for their valuable feedback, and have updated our revision to incorporate their suggestions for further improvement our paper. We summarize our overall major changes below, and detail smaller changes based on each reviewer's recommendations in the replies to individual comments.

- Added experiments based on various retriever choices in Appendix B.5.
- Added experiments on QA tasks under the retrieval setting in Appendix B.6.
- Added histogram plots of document statistics in Appendix B.7.
- Updated Section 5 with equations detailing our fine-tuning objectives
</content>
</thread_object>

</thread>


Metareview:
metareview:
PICASO proposes a novel method for efficiently composing and generating from multiple retrieved document chunks using State Space Models (SSMs). The key claims are that it can match the performance of document concatenation while achieving a 5.4x speedup, and that it enables permutation-invariant composition of document states. The paper's main strengths include addressing an important practical problem (efficient RAG with SSMs), strong theoretical foundations with detailed analysis, and comprehensive empirical validation showing significant speedups while maintaining performance. The initial weaknesses included: insufficient experimental validation on real RAG applications beyond perplexity metrics, missing baseline comparisons with similar approaches, lack of clarity around retriever choices and document statistics, and incomplete analysis of the order-invariance assumption. However, during rebuttal, the authors added substantial new experiments and analyses that addressed many of these concerns, including QA task evaluations, ablations with different retrievers, and detailed document statistics. The scores (5, 6, 8, 5) suggest a borderline paper, but the thorough technical contribution combined with strong rebuttal responses addressing core concerns merit acceptance as a poster.