# PICASO: Permutation-Invariant Context Composition with State Space Models

arXiv: https://arxiv.org/abs/2502.17605 · HTML: https://mlsanity.com/paper/2502.17605/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 39 references that are in the map.
- First posted on arXiv 24 Feb 2025.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [Jamba: A Hybrid Transformer-Mamba Language Model](https://mlsanity.com/paper/2403.19887/) (2024)
- [Demystifying Prompts in Language Models via Perplexity Estimation](https://mlsanity.com/paper/2212.04037/) (2022)
- [Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling](https://mlsanity.com/paper/2406.07522/) (2024)
- [An Information-theoretic Approach to Prompt Engineering Without Ground Truth Labels](https://mlsanity.com/paper/2203.11364/) (2022)
- [In-context Learning with Retrieved Demonstrations for Language Models: A Survey](https://mlsanity.com/paper/2401.11624/) (2024)
- [B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory](https://mlsanity.com/paper/2407.06324/) (2024)
- [Tangent Model Composition for Ensembling and Continual Fine-tuning](https://mlsanity.com/paper/2307.08114/) (2023)
- [Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection](https://mlsanity.com/paper/2405.16178/) (2024)
- [Tangent Transformers for Composition, Privacy and Removal](https://mlsanity.com/paper/2307.08122/) (2023)

## Built on by (most cited first, 4 of 4)

- [MiniPIC: Flexible Position-Independent Caching in <100LOC](https://mlsanity.com/paper/2606.13126/) (2026)
- [Studying the Soupability of Documents in State Space Models](https://mlsanity.com/paper/2505.24033/) (2025)
- [Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection](https://mlsanity.com/paper/2608.02560/) (2026)
- [Machine-Interpretable Information: Compiling Documents into Searchable and Readable Protocol States](https://mlsanity.com/paper/2609.23371/) (2026)

## Abstract (arXiv)

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.4x speedup.