PICASO: Permutation-Invariant Composition of Document States
State Space Models;Composition;Retrieval
Thanks for the response.

Maybe kv cache is the wrong word. What i meant is that since this method retrieves pre-computed states, one can compare with a similar transformer baseline, where the kv caches for documents are pre-computed and then can be directly loaded into the model for generation. If the original time already follows this setup then this question is addressed.
While we are currently preparing our official rebuttal, we would like to promptly address this concern in view of the increased attention it has received. Indeed, we tried to make our timing comparisons with transformers as fair as possible. We refer to L75-76, where we state

> These timings are measured using the official Mamba benchmarking code, which includes optimizations such as quantization and CUDA graphs for SSMs, and flash attention for Transformers

This naturally includes using the KV cache for generation, and we will update the relevant lines to make this explicit.
Quick Clarification
Good defense. But if it does not use KV cache, it would scale in O(n^3). Anyway, my point is clear. I do not think you read the paper carefully.
Thanks for your comment. I think the comment is not clear.

The authors did compare with a transformer baseline, and the kv cache question is regarding that transformer baseline number.
Meant for the transformer baseline comparison
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.
This reviewer is not reading paper carefully.
Yes
No ethics review needed.
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?
- 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.
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.
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.
Yes
No ethics review needed.
N/A
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.
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.
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.
Yes
No ethics review needed.
See weakness.
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.
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.
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.
Yes
No ethics review needed.
See questions in the above section
The proposed method shows significant speed up over the naive baseline
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.
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?
We propose a method for efficiently composing and generating from multiple retrieved document chunks based on state space models
@inproceedings{
anonymous2024picaso,
title={{PICASO}: Permutation-Invariant Composition of Document States},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=88TC1AWV27},
note={under review}
}
Generation with Large Language Models (LLMs) is often augmented by incorporating additional information "in-context" through the concatenation of multiple prepended documents. However, this process leads to significant computational costs that scale with the number and size of each document chunk. State Space Models (SSMs) offer a promising solution by allowing a database of documents to be pre-processed as states from which to start the generation. However, it is infeasible to store the state corresponding to  all possible combinations of multiple relevant documents. To address this challenge, we present a method called Permutation-Invariant Composition with State Space Models, or PICASO, which can compose states to efficiently approximate the effect of document concatenation. Our method can also enforce invariance to the order in which documents are presented, a desirable property when the temporal ordering of the documents is uninformative. We evaluate PICASO on WikiText and MSMARCO in both zero-shot and fine-tuned settings, and show that PICASO can match the performance of concatenation while enjoying on average 5.4x speedup.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
State Space Models Composition Retrieval
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/621f61a2d6cb0ccbab5f3dd72602a1247b498631.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
PICASO: Permutation-Invariant Composition of Document States
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission