Research map / Large language models / Language, safety and interpretability
PICASO: Permutation-Invariant Context Composition with State Space Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PICASO, which composes precomputed SSM states of multiple retrieved contexts into one permutation-invariant state (averaging CASO states over permutations), enabling fast RAG-style conditioning without processing context tokens.
Problem
Conditioning a state space language model on multiple retrieved contexts requires processing them online, which is costly; existing SSMs have no principled way to compose multiple precomputed states, and the context order is often uninformative.
Approach
Derives an exact single-layer identity that composes context states into the state of their concatenation (CASO), then averages CASO states over permutations (full symmetric group via dynamic programming, or cyclic group in linear time) to enforce permutation invariance; fine-tuning with backpropagation through or to the composition closes the gap to concatenation.
Key results
- Zero-shot PICASO-R on WikiText-V2 improves over no-context baseline by 10.1% on average vs concatenation 11.1% and Soup 8.5%, with about 5.4x average speedup.
- Fine-tuning with BPTC brings PICASO to concatenation-level performance on MSMARCO (e.g., 0.65 vs Concat 0.65 log-perplexity in BPTC-MSMARCO row).
- PICASO-R performs nearly identically to PICASO-S while being much faster to compute.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PICASO-R instead of concatenation for retrieved contexts | Mamba2-2.7B · MSMARCO validation, zero-shot | log-perplexity | Hurt |
| Use PICASO-R instead of concatenation for retrieved contexts | Mamba2-2.7B · WikiText-V2, zero-shot | performance improvement: 11.1% -> 10.1% | Hurt |
| Use PICASO-R instead of concatenation for retrieved contexts | Mamba-2 2.7B · OpenbookQA with retrieved WikiText passages | accuracy: 40.0% -> 39.9% | Hurt |
| Use PICASO-R instead of CASO by averaging cyclic permutations (ablation) | Mamba-2 2.7B · WikiText-V2, zero-shot | performance improvement: 9.2% -> 10.1% | Helped |
| Use PICASO-R instead of CASO by averaging cyclic permutations (ablation) | Mamba-2 2.7B · MSMARCO validation | log-perplexity: 1.56 -> 1.52 | Helped |
| Replace cyclic permutations with full permutation averaging (ablation) | PICASO · MSMARCO validation | log-perplexity: 1.52 -> 1.52 | No clear effect |
| Replace cyclic permutations with full permutation averaging (ablation) | PICASO · WikiText-V2, zero-shot | performance | No clear effect |
| Average independent context states instead of composing CASO states (ablation) | Mamba-2 2.7B · WikiText-V2 | performance | Hurt |
| Fine-tune with BP2C on MSMARCO contexts (ablation) | PICASO-R · MSMARCO validation | log-perplexity: 1.52 -> 0.69 | Helped |
| Fine-tune with BPTC on MSMARCO contexts (ablation) | PICASO-R · MSMARCO validation | log-perplexity: 1.52 -> 0.65 | Helped |
| Use BP2C rather than BPTC fine-tuning (ablation) | PICASO-R · MSMARCO validation | log-perplexity: 0.65 -> 0.69 | Hurt |
| Use stronger context retriever embeddings (ablation) | PICASO-R · WikiText | performance | Helped |
Surprising findings
- Averaging permutation-invariant CASO states substantially outperforms averaging independent context states (Soup), and the paper reports PICASO's gain over the best order-dependent composition at about 10%.
- Fine-tuned models maintain or slightly improve general LLM benchmark performance after composition fine-tuning.
Limitations
- Zero-shot PICASO still lags slightly behind concatenation in accuracy.
- Limited to SSM architectures; not applicable to attention-based hybrids without further methods for composing KV caches.
- Composition relies on approximation for multi-layer SSMs, since the exact identity holds only for single-layer models.
Experimental setup
Datasets: WikiText-V2 MSMARCO OpenbookQA HellaSwag PIQA ARC-E ARC-C WinoGrande
Benchmarks: WikiText-V2 MSMARCO OpenbookQA HellaSwag PIQA ARC-E ARC-C WinoGrande
Metrics: log-perplexity latency / composition time precision (data attribution) normalized accuracy
Baselines: Concatenation Soup CASO PIConcat-R No-retrieval baseline
Builds on: Mamba-2 Model Soup CASO
Closest papers by meaning
- L$^2$M: Mutual Information Scaling Law for Long-Context Language Modeling (NeurIPS 2025 · similarity 0.87)
- Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism (ICML 2025 · similarity 0.86)
- Data Engineering for Scaling Language Models to 128K Context (ICML 2024 · similarity 0.86)
- Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs (ICLR 2024 · similarity 0.86)
- Context Steering: Controllable Personalization at Inference Time (ICLR 2025 · similarity 0.86)
- Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.