PICASO: Permutation-Invariant Context Composition with State Space Models
Authors: Tian Yu Liu
Alessandro Achille
Matthew Trager
Aditya Golatkar
Luca Zancato
Stefano Soatto
Venue: ICLR 2025 Poster (2025)
Keywords: State Space Models
Composition
Retrieval
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a method for efficiently composing and generating from multiple retrieved contexts based on state space models

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.

Published as a conference paper at ICLR 2025
PICASO: PERMUTATION CONTEXT COM- POSITION WITH STATE SPACE MODELS
Tian Yu Liu" UCLA
Alessandro Achille
Matthew Trager AWS Al Labs
AWS AT Labs
Aditya Golatkar AWS AI Labs
Luca Zancato AWS AI Labs
Stefano Soatto AWS Al Labs
ABSTRACT
Providing Large Language Models with relevant contextual knowledge infer- ence 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 addi- tional contexts online incurs significant computation costs that scale with their length. State Space Models (SSMs) offer promising solution by allowing a database of contexts to be mapped onto dieedddimesiooall states from which to start the generation. A key challenge arises when attempting to leverage infor- mation 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 con- catenating raw context tokens. Since the temporal ordering of contexts can of- ten be uninformative, we enforce states obtained via our composition algorithm across all possible context order- ings. 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 per- forming baseline while enjoying on average 5.4x speedup.
-nniaiiaiionnnvvaiinneeee by efficiently averaging
INTRODUCTION
Incorporating new information in deep learning models has traditionally been a costly process, of- ten requiring re-training or fine-tuning their weights on new data. Fortunately, Large Language Models (LLMs) provide compelling alternative: These models can 'learn' to leverage new con- textual information in real-time by simply prepending them as inputs, without having to modify their weights (Ram et al. 2023). This has motivated powerful application known as Retrieval- Augmented Generation (RAG), where LLMs are deployed with the ability to retrieve and incorpo- rate relevant sources of information, 'contexts", from vast external knowledge bases when queried by users at inference time.
Despite being faster than the naive fine-tuning of model weights, this approach still incurs significant computational costs. Not only must the system process the user query and generate an answer, but it must also process the retrieved context, which in real-world settings can amount thousands of kens. This problem is exacerbated in Transformer-baat based models, as the inference cost of generating output tokens scales quadratically with the length of the extended input (see Figure D).
In contrast, State Space Models (SSMs) offer more efficient alternative. SSMs encode information from arbitrary-length input sequences into a fixed-size state vector, which can then be conditioned on to generate new tokens without revisiting the original input. This suggests a simple method to reduce the cost of incorporating new contexts at inference-time: Instead of retrieving from a database containing raw context tokens, we can create a "database of states" containing pre-computed state representations contexts. At inference time, generation starts directly from the retrieved state,
Work conducted during an internship Amazon

---

Published as a conference paper at ICLR 2025
Database of
States
Interence Time (ms)
2000
Transformer (Pythia 2.88)
SSM (Mamba-2 2.78)
1750
PICASO-R
Proprieress
1500
Processing Increase in
1250
1000
750
Compose
Retrieve
500
250
Summanse the key fiedings of this month's paperss
Query
4096
8192
16384
Total Length of Retrieved Contexts
Figure I: (Left:) We propose a "Database of States," where contexts are stored as pre-processed vectors. Given query, relevant states are then retrieved and composed into single state vector which is used to condition the model's generation. (Right:) We plot the increase in total time required to generate an additional contexts. We model the time taken for PICASO-R as the time taken to combine 5 pre-processed context states, which involves only arithmetic operations and notably zero model processing time. As result, the processing and inference costs for PICASO-R remain constant regardless of the length of retrieved contexts. In contrast, the timings for a Transformer model scale quadratically, and for SSM linearly, with total length when generating from concatenated context tokens. These timings are measured using the official Mamba benchmarking code, which includes optimizations such quantization and CUDA graphs for SSMs, and flash attention for Transformers.
tokens, when concatenating 64-token prompt with retrieved
simultaneously eliminating the latency from having to process context tokens online, and greatly reducing inference time compared Transformer models (Figure
However, key challenge arises when conditioning on multiple retrieved states. While input tokens can be simply concatenated and fed into LLM or SSM, existing models are trained generate outputs conditioned only on single SSM state. To address this, we derive a simple mathematical relation via SSM dynamics to compose multiple states into one, in manner that exactly equates to the result of concatenation in single-layer model. Consequently, by simply storing an additional weight matrix along with each context, pre-computed states can effectively composed at inference time to condition generation any arbitrary set of contexts.
Since states are computed causally, the order in which contexts are presented affects the state when there is natural ordering among retrieved contexts, different order of presentation would yield different states. Consequently, we propose to enforce order-invariance explicitly through av- eraging states obtained by composing contexts across all possible permutations. While this may appear costly at first glance, we show that the resulting state can be computed exactly in polynomial time the number of context segments using Dynamic Programming, and this can be further re- duced to linear time by accepting slight approximations. This greatly benefits Setrieval-Augmeet Generation tasks, where our results show a 10% improvement over the best order-dependent state composition method when invariann is incorporated into the conditioned state.
To outline our main contributions, we introduce a method for efficiently retrieving and compos- ing multiple pre-computed states at inference time to condition the generation of high-quality outputs, which we term PICASO (Permutation-Insel Invariant Compositional Aggregation of States as Observations). Our experiments show that PICASO achieves 91% of the performance gain from combining the raw tokens of multiple contexts, while offering 5.4x speed-up over concatenation.
PICASO can be applied to any off-the-shelf SSM model without any changes. To further improve performance, we introduce a method for fine-tuning the model to better leverage the composed states for generation. Using pre-trained Mamba-2 2.7B model, less than day of fine-tuning single A100 GPU leads to the same performance as concatenation while maintaining the 5.4x faster composition time the WikiText- V2 dataset.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
State Space Models and Hybrid Models. Recent efforts to overcome the significant computa- tional costs of Transformer models on long contexts have inspired the exploration of more efficient alternatives, including State Space Models (SSMs). Through maintaining fixed-size "state", sufficient statistic of the past for the purpse of future prediction, these models offer advantages compared to Transformer models. They only require constant memory consumption regardless of the sequence length, and linear computational complexity, rather than quadratic, as longer sequences are processed. The idea of leveraging recurrent models with fixed dimensional states to represent complex sequences is not new, in fact, several variations of SSMs have been developed in the past, ranging from Linear Time Invariant (LTI) systems, to more expressive non-linear Time Varying @azwinskii 2007) and Input Varying (Krener 1975) systems.
More recently, many of these ideas have been rediscovered and implemented on modern parallel hardware as basic building blocks for Foundation Models. Guu & Dao (2023) proposed Mamba, an input-dependent linear SSM (termed "selective") based on LIV systems, that achieves compa- rable performance to Transformers (Vaswani 2017) language modeling while enjoying faster inference. Mamba-2 (Dao & Gu 2024) further improved computational time by implementing SSM layers with structured matrix multiplications to better leverage modern Tensor Cores. Although pure SSM models can compete with Transformer blocks on many NLP tasks, they lag behind on tasks that require strong recall capabilities. To balance inference efficiency and model capabilities, Hybrid models combining Attention and SSM blocks have been proposed. [Leberr et al. (2024) combined SSM blocks along with global-attention blocks to create a hybrid architecture with Mixture-of- Expert layers for training larger models. To further improve long context ability and efficiency, |Ren| (et al. (2024) leveraged sliding window attention, while Zancato et al. (2024) developed a general family of architecture that include Transformers, SSMs and their hybrid combinations, leveraging both verbatim and fading memory, in both long- and short-term.
Retrieval- Aurmmete-- Generation and In-Context Learning. Our work falls within the scope of In-Context Setrieval-Augmentee Language Models (Ram et al. (023), where language models conditioned retrieved contexts via concatenation. Retrieval Augmented Generation (RAG) allows language models to leverage knowledge stored in external databases, which greatly improves performance on knowledge-intensin and -omain--peccffc tasks (Lewis et al. 2020). In our work, we simply use pre-trained sentence embedding model for retrieval, and we refer to Gao et al. (2023) for detailed survey on other mechanisms. Apart from retrieval, processing (multiple) re- trieved contexts can also greatly increase generation latency. Izacard et al. (2023) mitigates this by independently processes each retrieved context with a LLM encoder, using cross attention over the concatenated encoder outputs. Zhu et al. (2024) similarly encodes retrieved contexts in parallel, and performs decoding in selective manner by attending only to highly relevant encoder outputs.
In-Context Learning (ICL. (Brown et al. 2020) has emerged as an effective method to perform infer- ence without learning (i.e., transduction), by conditioning on labeled samples provided in-context, commonly implemented as set of (query, answer) pairs (Dong al. 2022). Similar to RAG, the quality of selected demonstrations have been shown to greatly affect downstream performance (Xu) let al. 2024). Several methods have been developed for selecting effective demonstrations, based on sentence embeddings (Liu et al.) 2021), mutual information (Sorensen et al.) [0022), perplexity (Gonen al.] (2022), and even BM25 (Robertson et al. 2009). Similar to the motivation of our work, several studies have shown that the performance of ICL is heavily dependent on demonstra- tion ordering. Zhho et al. (202T) shows that answers positioned towards the end of the prompt are more likely to be predicted, while Lu et al. (202I) shows that results can vary wildly between ran- dom guess and state-of-the-an depending on the order that demonstrations are presented. Outside of ICI, et al. (2024) further shows that language models do not robustly utilize information in long input contexts due to sensitivity to positioning.
Model and State Composition. Our work falls into the category of composing of deep models, representations, and states. Wortsman et al. (202)) proposes Model Soup, which composes multiple non-linearly fine-tuned models via averaging model weights. Liu & Soatto (2023): ); Liu et al. (2023) leverages model linearization to enforce an equivalence between weight averaging and output en- sembling. Perera et al. 2023) independently learns task-specific prompts which can be linearly
3

---

Published as a conference paper at ICLR 2025
averaged to yield new prompts for composite tasks. For SSMs, Pióro et al. (2024) investigates averaging states, along with decay-weighted mixing which closely related to a baseline ver- sion of our method, CASO. However, the equations described in their work differ from CASO, and their evaluations are limited to composition of two equal-length contexts. In contrast, our method greatly improves upon CASO by incorporating permutation invariance, which we show is important to achieve performances comparable to that of concatenation.
3 METHOD
3.1 PRELIMINARIES:
linear input-dependent discrete-ting state-space model has the form
St = A(u))xa-1 + B(u)uu
(1)
V/1 C(u))2 Du.-
Here Ty E IRm the state at time I, while ur, M/t Rd are the inpul and the outpul respectively. The matrices A(u,) E mm B(ue) E Rmxdd , Rdxm (which are input-dependent) and D E gaxd are learnable parameters.
Unrolling the first equation, obtain
It A(u)) A(u))2o
A(w)).
A(u)ro + xx)),
where = (ui,...,wu, ue) denotes the sequence of inputs, A(u) = A(w+) ... A(u)) is the accumu- (
lated decay matrix and x(u) - A(u))...
is the accumulated input signal. Since this coincides with V, when 2o = O, we refer to it as the state for input sequence u.
In the following, the set of variable-length sequences of token embeddings. We view State Space (language) Model (SSM) as a map fo 12 x Rm +> PP()) with parameters 0 which takes as input an initial state IR" m and token embedding sequence E V", and returns distribution over V. Modern SSMs Gu & Dao 2023 Zancato et al. (2024) usually contain multiple stacked selective state space layers in equation In multi-layer setting, write su) and A(u) for the sequence of states and decay matrices corresponding to all layers.
write V C RRd for a finite of token embeddings and V Unzo Va for
3.2 DATABASE OF STATES
By the Markov property, the state an SSM makes the past independent of the future. other words, fo(u u', 0) this means that SSM model can equivalently be initialized with the state arising from a (variable- length) input sequence, instead of the input sequence itself. This is akin to the KV-cache of Trans- former architectures, except that the dimension of the state is fixed regardless of sequence length.
x(x')) for all u' V*, where denotes concatenation. In practice,
In several real-world use cases such Retrieval Augmented Generation, relevant contexts com- monly obtained or retrieved from a database (Borgeaud et al.) 2022). Instead of storing them in the database raw text tokens, propose use "database states," where we pre-process each context and store their states. When conditioning on a single context, we can initialize the SSM with the retrieved pre-processed state instead of having to process online. However this poses problem when attempting to compose multiple contexts, since we do not know how to compose their states. We will show how this tackled with our proposed method.
3.3 PERMUTATION- -NNIIIION--NAAAANTT COMPOSITION WITH STATE SPACE MODELS
Given query and collection of relevant contexts, easy method to compose them is simply concatenate all context tokens with the query into a single sequence to feed into the SSM. Recall that this, however, presents key limitations. Before single token continuation can be generated
4

---

Published as a conference paper at ICLR 2025
from the query, the entire sequence of concatenated contexts has to be processed sequentially. which can be computationally intensive when contexts are long or numerous (Figure I Another limitation is having select the order of context concatenation when prompting the model, for which there might be no natural way of doing so without a powerful scoring mechanism.
To address the first limitation, propose first version of method, Compositional Aggregation of States as Observations (CASO), which works by modeling sequence concatenation with state composition based on the dynamics of single-layer SSM.
Proposition 1 (CASO), Let ui...., Un be collection of input sequences and let u = u1 Un be their concatenation. Then, for SSM layer that evolves based on equation we have
n-1
s(u)= r(un) +
A(un) A(us+1) x(u,))
aauuu.....n1 We can see this by recursively applying equation 2 on x(u) = A(u, (a(u1 ... un-1) + x(u)).
Given collection of contexts WisH... Un CASO simply approximates the dynamics of multi- layer SSMs, for which Proposition 1 does not hold exactly, via x.AASO(; (u1.... u,)) = x z(u)) +
Σ A(u,) ... A(u:A)((44))))) (u;). We then load 22AAO (u1,... un) as the initial state of the model to infer continuations from the given query. We note that Mamba-style models, the matrices A(-)
are diagonal. As such, computing CASO requires only simple element-wise arithmetic operations and importantly zero model computation time (i.e. zero forward passes required).
However, since each state is weighted by the decay factors of future contexts, this composition operation is still very much order-dependent. We propose to introduce permutation -in.uaaatiionniviinneeee by considering group of permutations G' C Sn, where S,, denotes the symmetric group of elements, using which we define our method, PICASO (Permutation- Invariant CASO):
PICASO (UI,...; Un)
CASO (u((1)).. Un(n))
For any group G, by expansion of the CASO terms and collecting common factors, this can be written as linear combination of individual context states x(u:):
PICASO (uI..... un) = Σ W Wi(u,. Un ( 1=1
with weights W. depending A(u))...,, A(u,)) In this work we are particularly concerned
with two cases: the full symmetric group G Sn, which includes all possible permutations, and the eyclic group G C.. which consists of rotations of the sequence. We will refer to them as PICASO-S and PICASO-R respectively.
While they appear computationally infeasible at first glance, since PICASO-S and PICASO-R av- erage over n! and n CASO states respectively, each of which is itself composition of n context states, the following propositions show that they actually be computed in polynomial and linear time respectively for modern SSM models with diagonal A matrices.
Proposition 2. Assume G Sn, and that the matrices A(w;) commute (e.g., are diagonal). Using shorthand notations A; == A(us) and W. = Wu(u,,-- un) we have
W4 =
(n 1)! + (n- 2) 1!.
Σ
A, + (n -3)! 3)1-21 21
Σ
Ai Aiz +
n!
15
Sn 61122
1
em (A,,..., A A An).
n-1)
m==
where
Em (A1 An-1)
Σ
A Aim
15h <12 Sn-1
the m-th elementary symmetric polynomial (Macdonald 1998) (in the matrices Az).
5

---

Published as a conference paper at ICLR 2025
Elementary symmetric polynomials satisfy the recursive relation
cm ; = An- An-2) + em (AI......... An-2).
Using this relation, we can compute all values of cm and hence the coefficients Wk. using (n²) op- crations via Dynamic Programming. We detail the implementation in Algorithm of the Appendix. Consequently, the full state from PICASO-S can be efficiently computed in polynomial Onn") time, which we show in the experiments to still be faster than processing textual context concatenations even for as large 10.
Next, we similarly show that the coefficients for PICASO-R can be efficiently computed by exploit- ing invertibility of the matrices A(u:).
Proposition 3. Assume G = Cn (cyclic permutations), Then writing A, := := =A(ss) and We. == W4 Wu(u,,..... Un) we have
Wx
ld
Aptmmaa
Apx+1]n
n
where Id is the identity matrix, and [ denotes mod n. Assuming that the matrices A; are invert- ible, these can be computed efficiently by setting
[An.. A1 An. A1 A (A
W 1 [A='(Bkto B&]], i> B.=
A;
i isn"
n
for I,..., 2n, and k ,....,7.
We detail in Algorithm 3 in the Appendix our efficient implementation for computing PICASO-R in (n)) time complexity via cumulative sums and products. Evidently, PICASO-R significantly faster than PICASO-S while trading off exact permutation invariance for invariance only to cyclic permutations of the original order. We will show that the difference empirical performance be- tween PICASO-S to PICASO-R is negligible, such PICASO-S can almost always be replaced with its much faster variant PICASO-R.
We remark that the property of permutation intiooinnnvvaiinccee can also be applied to naive concatenation (as opposed to CASO). This is achieved simply by concatenating contexts in various different orders, followed by taking an average of their resulting states. While performing this for the symmetric group Sn is computationally infeasible, we can similarly restrict our permutation set to Cp. We term this variant Eermutation-Invariian Concatenation (PiConcat-R), where R denotes invariance to the set of cyclic permutations. We note that the model computational costs (forward passes) of this method still scales quadratically with number of contexts (compared to linear scaling of regular concatenation), as such we include it only for completeness.
As final technicality, we note that for Mamba-style SSM models, we additionally require storing the last Mcony (usually Mcone = 4) input tokens of each SSM layer to ensure that the state is sufficient for generating the same distributions over continuations as the imput sequence. We perform simple averaging to combine these tokens from different contexts which we show to work well empirically; more sophisticated methods could be explored future work.
WHY PICASO'S AVERAGE WORKS
While the combination of state expression for CASO is directly motivated by the dynamics of the system, there is no priori reason why averaging permuted CASO states should perform well. In Figure 3 we show that averaging both independent states and CASO states can perform better than using any individual state. This suggests emergent/carneed algebraic structure the space of states such that linear combination of states combine their information to some degree.
In our empirical results below, show that averaging all individual states (which would also be a permutation-inven
solution) performs significantly weaker than averaging CASO states (as PICASO does). We believe that this is because the approximate linear structure of the state space only valid locally. The combined states are naturally closer together than the independent states, hence able to better exploit the local linearity. We show this in the following proposition:

---

Published as a conference paper at ICLR 2025
Proposition 4. Consider single-layer SSM parametrized by 9, and two input sequences u and u'. Then, the Euclidean distance between the states can be bounded via
1x CASO uu u') xC (u', u) S II(I A(u'))ar())13 III -
To see this, simply apply the triangle inequality on the following obtained via substituting the equa- tions for CASO:
| CASO (u, (u,u') - x CASO u w)ll 2 ||A((')) (((u))=(())))))) + x(u') = + 1
(I(A(u))) - ((A((u)--)x(u)))
((u))
As special case, we observe that the decay factor approaches the identity, the distance between two CASO states approaches zero. In Figure we visualize naive averaging of the states arising from 3 retrieved contexts, and averaging of CASO states resulting from each cyclic permutation of these contexts. We use WikiText-v2 as described in the Experiments for these plots. Indeed, we observe that CASO states are much closer one another in the resulting loss landscape.
Contaut
2523818 227
DASOCULD
2.22818 27
Contest Auarage
2
(KSO0U2 Average
2.49990x
2.4597XX
2.43997x 201
243355x 22"
244356 22
64-
00-
044
22502"13"
0.4
23382
2.130000
2.1100X
22809-38
2,000 101
221
222142-329
3.6 3.6
0. 0.1 0.6 0.8
2.5
Naive averaging ("Soup"") of context states. states are "closer" to one another (see Proposition) and hence can be more meaningfully interpolated. On the other hand, naively averaging states of independent contexts do not possess this property. Both plots are computed over 10 samples of (query, continuation, retrieved contexts).
Figure 2: Left:
Right: Averaging CASO states. CASO
5 LEARNING TO USE COMPOSED STATES
As previously noted, in practice, for SSM models consisting of multiple state space blocks stacked with temporal convolutions, =(u) in equation will not be exactly the state arising from a concate- nated list of inputs. In this section, we introduce fine-tuning objective to enable SSMs to better leverage composed states. Let D ((u, u, J be dataset of sequences u, their next-token continuation uy and a collection (in some particular order) of contexts S; retrieved from database using u, We minimize the prediction loss over the continuation, given a (composed) initial state and the query sequence:
LBPTC(
Σ SDED
LCE(fo (u,
P (S.)),
where LCE(-.) the cross-entropy loss.
We denote this learning objective Backpropagation Through Composition (BPTC), where gradients are propagated through the state composition process z PICASO To reduce training time, we also consider an alternative version where we do not backpropagate through the composition step, which we denote Backpropagation To Composition (BP2C):
LBP2C(0)
Σ
LCE(Solu; sg PICASO (ICCAOOSS))),.. (S.))),u)
where sg denotes the gradgrrddent operator. We will show that when used for fine-tuning, this learning objective greatly improves the model's ability to leverage composed states for generation to the level of the concatenation albeit with much faster speeds, while maintaining performance on standard L.MM evaluation tasks.
7

---

Published as a conference paper at ICLR 2025
6
EXPERIMENTS
Baseline
2.65
Concat
2.65
2.60
CASO
2.55
PICASO-R
PICASO-S
2.55
10
2.50
PIConcat-R
2.50
Bassine
Algetthm
2.40
PICAGO-
2.35
X
2.30
10
100
150
200
250
Number of Compesed Segments
Composition Time (ms)
Figure 3: Zero-shot evaluation of PICASO using Mamba-2 compared to other composition methods on WikiText. While the performance of PICASO lags slightly behind that of concatenation (left), PICASO-R is on average 5.4x faster (right). PICASO-S and PICASO-R perform similarly and yield overlapping curves (hence not visible in the left plot). Incorporating permutation invariance for concatenation via PIConcat-R gives the best results. However, it incurs magnitudes higher com- putational costs despite being performed within single batched forward pass, hence we omit from the right plot to prevent it from disrupting the scale of the x-axis and focus comparisons on PICASO.
6.1 IMPLEMENTATION DETAILS
We main experiments the largest available SSM Huggingface Mamba-22.7 2.7B (Dao & Gu 2024). We evaluate our method on two large-scale datasets WikiText- V2 (Merity et al. 2016) and MSMARCO (Nguyen et al., 2016). We use the training splits as fine-tuning data, and the testing/validation splits respectively for evaluation. To pre-process WikiText-V2 for our use case, we split each passage the dataset into two equal context "segments", with the goal of predicting the second (continuation) from the first (query). The retrieval database comprises all remaining segments, from which we retrieve via an external sentence embedding model, All-MiniL.M-L--y In most experiments, we retrieve up to 10 segments, since improvements appears to saturate beyond that, and loss from concatenation blows up as a result exceeding training context length (Figure Appendix). We pre-process MSMARCO by filtering only entries with well-formed answers and discarding those without relevant passages.
We used the official benchmark with with fairest comparisons. For the rest of the experiments, we run the model in full-precision and evaluate performance of the model starting from custom initial state, a feature not supported by the official benchmark at the time of writing, as such timings differ.
A100 GPU for our timing experiments in Figure to ensure
For tne-ttuning experiments using BPTC and BP2C, we base implementation the official HuggingFace trainer with default hyperparameters, and retrieve the most relevant context seg- ments for each query sample for composition. For WikiText, at random for each batch. For MSMARCO, we use all the available passages (both relevant and ir- relevant) associated with each training example. For both datasets, we fine-tune for only epoch. In all fine-tuning experiments, we ensure the training set (both the examples and the context database) disjoint from the validation set ensure fair evaluation.
select k: (O...., 10) uniformly
6.2 COMPARISON MODELS
We compare inference accuracy (measured by log- perppeeeeettt) and processing latency of PICASO with its order-dependent version, CASO, in addition to the following methods:
"https///github... http:///hthhh.coomm/a/tttccc///a httpss//ggthh..oom/hggnniiccccce
8

---

Published as a conference paper at ICLR 2025
Baseline: Loss of the model the test sample without using any contextual information.
Concatenation (Ram et al. y We concatenate individual context segments based on a specific ordering. For WikiText-V2 V2 experiments, the sentence embedding model where more relevant contexts are closer to the query (at the end). We initialize the model with the state of the earliest context segment in the concatenation, which we assume to be available via pre-processing, and recompute the composed state from only the remaining ones.
consider the "best-case ordering" as determined by
Soup (Póóro et al. 2024): Simple averaging of states obtained from each context.
6.3 MAIN RESULTS
In this section, we evaluate both the zero-shot and fine-tuned performance PICASO in Sec- tion 1 and Section 6.3.2 respectively, and show in Section 6.3.3 that the fine-tuned model does not overfit to the composition task. We also include additional experiments showing that LLM ca- pabilities are not impacted by fine-tuning in Appendix (e.eedixiB...... and show that PICASO can also be used for data attribution in Appendix C
6.3.1 ZERO SHOT PERFORMANCE
We demonstrate in Figure 3 that applying PICASO-R in a zero-shot manner on WikiText- V2 greatly improves performance over the baseline by average of 10.1% across 1-10 retrieved context seg- ments. This greatly improves over Soup (8) and CASO (9.2%). Compared to concatenation (11.1%), PICASO-R performs slightly worse but benefits from magnitudes improvement in pro- cessing time on an average of 5.4x. In this task, PICASO-R achieves almost exactly the same performance as PICASO-S, but with a much faster composition time. As sanity check for motiva- tion for our method, we show that PIConcat achieves the best performance (12.0%) overall, but at the cost of significantly greater computational time despite our batched- eecceed- implementation.
In Row 1 of Table [ we show that applying PICASO-R R and PICASO-S S in a zero-shot manner on MSMARCO similarly yields considerable improvements (37.2%) over the naive baseline, while achieving performance close to that of concatenation (41.3%).
kaseliee
2.20
Scudins
2,28
2.14
CAGO CARON
3.32
F
Concess
2.24
2,00
2.00
2.5
2.02
2:06
284
Number Componed Segments
Camposition
Kumber Componed ments
Figure 4: (Left Middle:) Fine-tuning with BPTC on WikiText brings the performance of PICASO to that of concatenation, while retaining its significant speed advantages. (Right:) Fine-tuning with BP2C on WikiText improves the effectiveness of PICASO as well, but much faster in terms of training time since it does not require backpropagating through the composed state. Note that fine- tuning has no impact on the actual composition time when used for inference.
6.3.2 BACKPROPAGATIONN THROUGH AND To COMPOSITION
While PICASO demonstrates strong performance in the zero-shot setting, PICASO still lags behind concatenation in terms of prediction accuracy. We attribute this to composed states being "out-of- distribution" for the model, since these states do not arise from any sequence of input tokens. In this section, we test if this can be resolved via fine-tuning with PICASO-R composed states via BPTC and BP2C. Indeed, as we show in Figure BPTC and BP2C greatly improves the perfor- mance of PICASO-R and PICASO-S to that similar to concatenation, while maintaining much faster processing timings on WikiText. Similarly, we show in Rows 4-5 of Table [ that fine-tuning on the MSMARCO training set also levels the performance of PICASO with that of concatenation. We also

---

Published as a conference paper at ICLR 2025
Table 1: All models in this table are evaluated the MSMARCO validation set. We evaluate performance of models fine-tuned via BPTC/BP2C on both the WikiText and MSMARCO training sets. Rows 2-3 show that fine-tuning models to compose WikiText context segments does not harm performance when evaluated on composing context segments from MSMARCO. When compos- ing segments from distributions similar those encountered during training (Rows 4-5), PICASO matches the performance of concatenation while being magnitudes faster.
Naive Concat Soup CASO PICASO-R PICASO-S
Mamha2-2.77 Base
2.42
1.42
2.04
1.56
1.52
1.52
Mamba2-2.7B BP2C-WikiText
2.44
1.44
2.07
1.53
1.50
1.50
Mamba2-2.77B BPTC-WikiText
2.43
1.46
2.08
1.53
1.49
1.49
Mamba2-2.7B BP2C-MSMARCO
1.85
0.68
1.27
0.72
0.69
0.69
Mamba2-2.77 BPTC-MSMARCO
1.79
0.65
1.20
0.68
0.65
0.65
note that while BP2C is significantly faster in terms of training time, it incurs a small performance trade-off compared to BPTC for both datasets, keeping number of training iterations constant.
6.3.3 EVALUATION OF FINE-TUNED MODEL ON OTHER DIFFERENT TASKS
We showed that models fine-tuned on specific downstream task (training set) using BPTC/BP2C perform strongly when composing samples drawn from a similar distribution (test set). We further show in Table 1 that models fine-tuned mance gains (or at the very least, no performance loss) when composing samples via PICASO on another domain (MSMARCO). Finally, we show in Appendix [B.4] that fine-tuning models with BP2C/BPTC maintain (and occasionally even improve) performance on general LLM evaluation tasks compared against the original model.
domain (WikiText)
demonstrate small perfor-
LIMITATIONS AND DISCUSSION
We have proposed method, PICASO, that enables efficient retrieval and composition of contexts by pre-processing their individual states. Without any training, our approach can handle the com- position of information contained in up 10 context segments PICASO notably requires zero online model processing time, since generation can begin directly from the composed states. When models are further fine-tuned with our proposed learning objec- tive, states composed using PICASO perform comparably to those produced from the concatenation of context tokens, while offering on average 5.4x faster composition time.
manner that order-invariant,
Nevertheless, our method does have some limitations. When applied in a zero-shot manner, PICASO still lags slightly behind concatenation in terms of prediction accuracy. PICASO also currently limited architectures based on SSM layers. We leave future work extension of PICASO towards recently popularized attention-based hybrid models, which require more sophisticated methods of composing key-value caches. Lastly, also leave future work the exploration of parameter- efficient tine-tuning methods such as adapters, which can be used to augment the model at inference time enable state composition while preserving the original model's 's behavior.
REFERENCES
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Milli- can, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. International conference on machine learning, pp. 2206-224.. PMLR, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:187771901,, 2020.
10

---

Published as a conference paper at ICLR 2025
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05433 2018.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprinl arXiv:2405.2 1060, 2024.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-context learning. arXiv preprint arXiv:2301 200::2301.002344 2022.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL https: /tt::://eenooooo org/records/
12608602
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. Retrieval- augmeete generation for large language models: A survey. arXiv preprint arXiv:2312. 10997, 2023.
Hila Gonen, Srini Iyer, Terra Blevins, Noah A Smith, and Luke Zettlemyer. Demystifying prompts in language models via perplexity estimation. arXiv preprini arXi::22/2.0403 2022.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint ar arXiv:2312.00752, 2023.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research, 24(251): 1-43, 2023.
Andrew H Jazwinski. Stochastic processes and filtering theory. Courier Corporation, 2007.
Arthur J Krener. Bilinear and nonlinear realizations of input-output maps. SIAM Journal on Control, 13(4):827 1975.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. tion for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
genera-
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz. et al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403. 190877877.7 2024.
Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? arXiv preprint arXiv:2101.0668 2021.
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024.
Tian Yu Liu and Stefano Soatto. Tangent model composition for ensembling and continual fine- tuning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp- 18676-18686, 2023.
Tian Yu Liu, Aditya Golatkar, and Stefano Soatto. Tangent transformers for composition, privacy and removal. arXiv preprint arXiv:2307.08122, 2023.
Yao Lu, Max Bartolo. Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2/04. 08786. 2021.
11

---

Published as a conference paper at ICLR 2025
Ian Grant Macdonald. Symmetric functions and Hall polynomials. Oxford university press, 1998.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint 1rXiv:1809.0278. 2018.
Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. MS MARCO: A human generated machine reading comprehension dataset. CoRR, 09268
abs/1611.00 1.09268, 2016.
Pramuditha Perera, Matthew Trager, Luca Zancato, Alessandro Achille, and Stefano Soatto. Prompt algebra for task composition. arXiv preprini urXiv:2306.003100 2023.
Maciej Pióro, Maciej Wołczyk, Razvan Pascanu, Johannes Oswald, and João Sacramento. State soup: In-context skill learning, retrieval and mixing. arXiv preprint arXiv:2406.084235 2024.
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Leyton-Brown, and Yoav Shoham. In-context retrieval augmented language models. Transactions of the Association for Computational Linguistics, 1 11:1316-1331, 2023.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert- networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019. htll:ttp:///xxt
abs/1908. 10084
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXhv:2406.07520 075522 2024.
Stephen Robertson, Hugo Zaragoza, et al. The probabilistic relevance framework: Bm25 and be- yond. Foundations and Trends@ in Information Retrieval, 3(4):333-38), 2009.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Taylor Sorensen, Joshua Robinson, Christopher Michael Rytting, Alexander Glenn Shaw, Kyle Jef- frey Rogers, Alexia Pauline Delorey, Mahmoud Khalil, Nancy Fulda, and David Wingate. An theoreicio approach to prompt engineering without ground truth labels. arXiv preprint 1XXi::2203.11664 2022.
Ashish Vaswani. Attention is all you need. arXiv preprini arXiv: 2017.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International conference on machine learning, pp.
PMLR, 2022.
Xin Xu, Yue Liu, Panupong Pasupat, Mehran Kazemi, et al. In-context learning with retrieved demonstrations for language models: survey. arXiv preprini arXiv:2401. 1622 2024.
Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, and Stefano Soatto. B' mojo: Hybrid state space realizations of foundation models with eidetic and fading memory. arXiv preprint arXiv:2407.06324 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma-
chine really finish your sentence? arXiv preprint arXiv: 1905.07830000 2019.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving
few-shot performance of language models. In International conference machine learning, pp. 12697-12706. PMLR, 2021.
Yun Zhu, Jia-Chen Gu, Caitlin Sikora, Ho Ko, Yınxiao Liu, Chu Chu-Cheng Lin, Lei Shu, Liangchen Luo, Lei Meng, Bang Liu, et al. Accelerating inference attrivaal-augmenter generation via sparse context selection. arXiv preprint arXiv:2405.1611 16178, 2024.
12

---

Published as a conference paper at ICLR 2025
A ALGORITHMS: PICASO-S AND PICASO-R
We show in Algorithm T how PICASO-S is computed in polynomial time via a dynamic program- ming approach based Algorithm 2
In Algorithm []] also show how PICASO-R can be com- puted with linear time complexity. Time complexity is measured as the number arithmetic oper- ations required function of number of context states.
Algorithm PICASO-S- O(n²)
Require: States æ = {a:)=o Decays A [A
return Ei= PICASO-S-DP(A 24
A_ denotes all elements A except A,
Algorithm PICASO-S-DP O
Require: Decays = {A,)
DP[::] +- zeros(n,
DP[O.:] f-1 1
w
0
for i = 1, ,,..., 1 do
for j = i...., n do
DP[i][ t j - 1] + A1-1- DP[i 1] - 1]
end for end for
for i = 0 . do
w n-("." x DP[i]
end for return
Algorithm 3 PICASO-R (nn)
Require: States (ar:) Decays A {A:)=0 < 0, Å [AI...., An, AI...., An]
An = cumprod (A)
Ágn = cumsum AAT)
for i 1....,71- do
wi + 1/- ((Azu(n +i - 1] - Asnlil An|o] +1)
x- à + wi Fz
end for return à
B FURTHER ANALYSIS
B. COMPUTATIONAL COSTS OF PICONCAT
In Figure 5 we visualize the computational costs incurred by PIConcat, which we show to dominate that of other methods despite resulting in the best performance on the WikiText dataset.
B.2 SCALING BEYOND EFFECTIVE CONTEXT LENGTII
In Figure 6 we show that as the total length retrieved contexts scale beyond a certain threshold (effective context size of the model), the loss from concatenation blows up and rapidly increases beyond the no-retrieval baseline. On the other hand, performance of PICASO remains stronger than that of the baseline when composing 50 context segments.
13

---

Published as a conference paper at ICLR 2025
Segments
2.65
2.60
2.55
Algorithm
Los
2.50
Concat
2.45
Soup
2.40
PICASO-R
PICASO-S
2.35
PIConcat-R
2.30
0
500
1000
1500
2000
2500
Composition Time (ms)
Figure 5: Timings for different composition algorithms evaluated on WikiText using Mamba-2 2.7B (zero-shot), including that of PIConcat-R. While PIConcat results the best performance (y-axis), its computational cost (x-axis) is significantly higher than that of other methods. We refer to Figure for more condensed plot to compare the remaining methods.
7
Baseline
Concat
Soup
6
CASO
PICASO-R
PICASO-S
Loss
5
4
3.
10
20
30
40
Number of Composed Segments
Figure 6: Concatenation scales poorly with total size retrieved contexts beyond training context length. PICASO yields greater stability even composing up 50 context segments retrieved from WikiText.
B.3 INFERENCE VS PROCESSING TIME
In Figure 7 we show that the context processing time for Mamba-2 comprises a significant propor- tion of the total generation time. For large sequence lengths beyond 6K tokens, the processing time even dominates the inference time for generating 32 tokens.
B.4 PERFORMANCE ON LLM EVALUATION TASKS
In Table
we show that fine-tuning Mamba2-2.1 2.mmaa2-2.7BBB with BTPC/BP2C objectives do not degrade existing LLM capabilities when evaluated several LLM evaluation benchmarks HellaSwag
(Zellers et al. 2019), PIQA (Bisk et al. 2020), ARC-E, ARC-C (Clark et al. 2018), WinoGrande
Sakaguchi et al. 2021), and OpenbookQA (Mihaylov al. 2018).
14

---

Published as a conference paper at ICLR 2025
Mamba-2 (2.7B) Inference Time
0.5
Mamba-2 (2.7B) Processing Time
Processing time
32 token inference
0.3
0.2
0.1
0.0
2000
4000 6000 8000 10000 12000 14000 16000 Input Sequence Length
Figure 7: Mamba-2 Processing vs Inference Time of 32 tokens. Processing time (orange) occupies significant proportion of the total time taken generate from an input sequence, even dominating the constant inference time from the processed state (blue) as number of tokens in the input grows.
Table 2: Evaluation of Mamba2-2.78 trained with BPTC and BP2C on LLM evaluation tasks. Here, we show that fine-tuning for composition does not degrade existing LLM capabilities. this table, report the length-normalized accuracy for each task.
HellaSwag
PIQA
ARC-E
ARC-C
WinoGrande
OpenbookQA
Mamba2-2.7B
66.6
76.3 f
64.8 1.0
36.3 1.4
63.9 1.4
38.8 2.2
BP2C- WikiText
66.7 + 0.5
76.3 1.0
1. 6.
63.6 1.4
39.8 - 2.2
BPTC-WikiText
66.7 + 0.5
75.6 1.0
64.9 1.
63.2 = 1.4
40.2 :: 2.2
B.5 ABLATION ON CHOICE OF RETRIEVER
In Figure 8 we ablate the impact difference retriever choices on PICASO-R. In par- ticular, we evaluate the performance of PICASO-R WikiText when using the follow- ing embedding models from Centence-Transfomeee age_word www..eebeddngggglooooo www.glllovov..... 6B3.00d all-MiniLM-L6-v2. and all-mpnet-base-v2 arranged in in- creasing order of performance on 14 different sentence embedding tasks (Reimers Gurevych, 2019). As expected, Figure [ shows that the performance of PICASO-R highly correlates with the strength of the retriever, where stronger retrievers yields better results on WikiText.
(Reimers Gurevych] 2019): aver-
average word, wwrdwwwwwddeeeeeeeeee www.e.lo.vvo0o0m 300d
2.65
all-mpnet-eee base-v2 all Mi-iLM--66.1
2.60
2.45
2.40
Number of Composed Segments/ States
Figure 8: Ablation study on how choice of retriever model impacts performance of PICASO-R on WikiText. As expected, stronger retriever models result in better downstream performance.
15

---

Published as a conference paper at ICLR 2025
B.6 EVALUATION ON MULTIPLE CHOICE TASKS
In this section, we evaluate PICASO-R on the OpenbookQA (Mihaylov et al. 2018) multiple-choie task, where we retrieve from context database of full passages from Wiki V2. While Open- bookQA provides the ground truth fact for each evaluation sample, we discard this in our evaluations following standard practice in Gao et al. (2024). We leverage the same retrieval model used for the main WikiText experiments. Table 3 shows that PICASO-R achieves performance close to concate- nation, with a 8x speed-up in composition time.
Naive
Concat
PICASO-R
Acc (f)
Time (4)
Acc (t)
Time (4)
Acc (t)
Time (1)
38.8%
NA
40.0%
233 ms
39.9%
29 ms
Table 3: Evaluation on OpenbookQA dataset, augmented with retrieved passages from WikiText. use normalized accuracy as our evaluation metric, and report the time taken to compose retrieved passages. Numbers shown in the table are averaged across retrieving between I to 10 full contexts from WikiText (as opposed to half context segments in our main paper experiments).
B.7 CONTEXT STATISTICS
In Figure 9 plot the distribution over the lengths (tokens and characters) of retrieved context segments used in the main paper WikiText retrieval dataset.
250
250
Num characters
Num tokens
150
150
200
400
600
800
1000
1200
50
100
150
200
250
300
Num documents
Num doruments
Figure 9: Histogram of the lengths, in terms of (Left) characters and (Right) tokens, of database context segments used the main paper WikiText experiments.
C DATA ATTRIBUTION
Table 4: Zero-shot Data Attribution on MSMARCO with Mambaa-2.7B. measured by precision. We compare Leave-One-In (LOI) and Leave-One-Out (LOO),, where implement 1.0 with varying methods for state composition.
LOI
Concat
Soup
CASO
PICASO-R
PICASO-S
0.699
0.690
0.629
0.725
0.732
0.731
Consider question- answer pair (ug. u.), and sequence of potentially relevant contexts S (u,...,, ). We would like to select the most relevant context for inferring the answer. There are least ways do so with model fo:
16

---

Published as a conference paper at ICLR 2025
The first method, which we term "Leave-one-in", is prepend each candidate context the ques- tion, and evaluate the loss on the answer. Equivalently, arg min; Loe(fo(ru (fo(Ug: z(u)), ug), where abuse notation to denote loss the sequence (instead of token) ua
The second method, which we term "Leave-one-out", is to compare the marginal increase in loss of the answer when removing each candidate from the composition of all of them. Equivalently, arg max; Loe uu. ()) u.) - Lfe(fo(u (S)),, u,)), where ((S_;) denotes a state composed from all contexts in S other than u,
Intuitvely, the former measures "absolute" inffuence of context, while the latter measures "relative" influence computed as the marginal improvement from adding the set of other contexts.
There are several different ways to implement the latter by varying the composition method used. We show in Table 4 that not only does Leave-One-Out perform best on the MSMARCO dataset, but implementing Leave-One-Out with PICASO-S and PICASO-R not only accelerates processing, but also surpasses the performance of conatenation. We attribute this the cermutation-ivvaiian property of PICASO, which unlike concatenation, does not introduce irrelevant biases arising from arbitrary context orders.
D CONCATENATION FOR SSMs: CONNECTION TO JUMP-LINEAR SYSTEMS
Consider a collection of context segments retrieved based on relevance to a query, and sorted ran- domly as context the query. While these segments share information, they are independent given the query, and their order is accidental and uninformative.
We are interested in a model that can efficiently process inputs in this format and extract shared information from the input. Attention-baut models are natural choice because of the permutation- invariance of attention mechanisms (ignoring positional encoding), but they would have to process the entire input (all segments) with quadratic inference cost. On the other hand, SSMs have linear cost, but they are ill-fit to process this kind of input because of the context switches, which make the Markov assumption implicit in the state representation invalid.
We consider a broader class of models, namely switching dynamical systems (or jump-Me jump-Markov, jump-diffusion, linear hybrid, or jump-linear systems) the class of interest. A jump-linear system is one that has a continuous state, say de that evolves synchronously. and a discrete state that changes the value of 2,, for instance
[Ax + Bug if Z [1+1 P i
T++11
Learning and inference for this model class corresponds to identification and filtering for this class of Jump-Markov models. In addition a random switching, the switch can be triggered by particular 'flag' (value) of the input:
{ Axx + Bug Us + Merigger
Tz+
[1++1 P if u Utrigger
the value of Uarigger is known, then a given identification and filtering scheme can be applied by switching the estimated state according to the trigger.
Since modern state space models are input-dependent, they automatically fit the latter class of mod- els and can handle switches without modifications. However, what they cannot handle is the fact that the order the segments uninformative. As result, presenting the same segments dif- ferent order would yield different states. Accordingly, our goal is to enable SSMs to learn from segments up to permutations, so we can accommodate sequences where the ordering within seg- ments is informative and respected, while the ordering of segments is uninformative and factored out.
E GENERAL RECURRENCE STRUCTURE
In the main paper, introduced specific recursive relation satisfied by Elementary Symmetric Polynomials. Here, we introduce a more general form which can potentially be used for more efficient implementations:
17

---

Published as a conference paper at ICLR 2025
Proposition 5. For any choice of 1 1 9 n - 1
g)
Cm (A1, An-1)
=
Σ
Cm- Cm-s(A1,"" - An-1-g)es(An-g>" ... An
3 j=max(g++n -+-
Proof. We compute em (A1,..... An-1) using a Dynamic Programming (DP) approach, where we break the problem into smaller problems, and merge the solutions. First split the n variables at some random index to create two partitions, (Ai... An-1-q) and (An-q:... An-1), and then compute ém-j and E, each partition respectively. For given value of j. Em-s(A1,... An-1-gge, (An-g) ... An.--:.-AAA- 1n will only compute a subset of values from em (A1, 11), and hence we sum over all possible values for j.
In particular, taking 1, we obtain the following:
Cm (A1,A.,... 1 = An-16m 1 An-2) + (A1..... An-2)
which we use for our implementation of PICASO-S.
18


Reviews:
[{"review_id": "8X21JgNfJp", "reviewer": "Reviewer_uYJU", "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.\n\nThe 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.\n\nExperiments 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.", "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].\n\n*[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.*\n\n2. 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?\n\n3. 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? \n\n4. 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?\n\n5. 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?\n\n6. 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.\n\n7. For Table 1, what is the number of retrieved document chunks for the results?\n\n8. 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.\n\n9. 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?\n\n10. 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.\n\n\nSome minor comments on typos and suggestions:\n\n11. In the abstract or introduction, how is the name PICASO composed for what it stands for?\n\n12. In proposition 3, explain the notation of `Id` and `[i]_n`?\n\n13. 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.\n\n14. 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.\n\n15. 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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper presentation is mostly clear. The motivation, background, and techniques are well explained, also in good connection with relevant research.\n\n- 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.\n\n- 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.\n\n- 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. \n- 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. \n- 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.\n- 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.\n- 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\n5. 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.\n\n*[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.*\n\nOverall 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a1yZfcf1zF", "reviewer": "Reviewer_FYyF", "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.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gozZissSbB", "reviewer": "Reviewer_2fh2", "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.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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.\n2. The State Space Model constitutes a suitable architectural framework for the targeted problem, thereby rendering the study meaningful.\n3. The introduced methodologies are technically robust. And they are presented clearly mostly (with a few exceptions).\n3. 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.\n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mTS73QUBk2", "reviewer": "Reviewer_h4ZY", "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.", "questions": "See questions in the above section", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "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.\n\nThe 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.\n\nIt is not clear what the speed baseline setup is. Does it use KV cache?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)