CASCADE Your Datasets for Cross-Mode Knowledge Retrieval of Language Models
Authors: Runlong Zhou
Yi Zhang
Venue: COLM 2025 (2025)
Keywords: large language models
pretraining
knowledge retrieval
spurious correlations
TL;DR: We conduct qualitative and quantitative studies on the cross-mode knowledge retrieval capabilities of LLMs, with a novel approach, CASCADE, to mitigate this issue.

Abstract: Language models often struggle with cross-mode knowledge retrieval – the ability to access knowledge learned in one format (mode) when queried in another.
We demonstrate that models trained on multiple data sources (e.g., Wikipedia and TinyStories) exhibit significantly reduced accuracy when retrieving knowledge in a format different from its original training mode.
This paper quantitatively investigates this phenomenon through a controlled study of random token sequence memorization across different modes.
We first explore dataset rewriting as a solution, revealing that effective cross-mode retrieval requires prohibitively extensive rewriting efforts that follow a sigmoid-like relationship.
As an alternative, we propose CASCADE, a novel pretraining algorithm that uses cascading datasets with varying sequence lengths and computing losses on only the second half of each training sequence to capture knowledge at different scales.
Our experiments demonstrate that CASCADE outperforms dataset rewriting approaches, even when compressed into a single model with a unified loss function.
This work provides both qualitative evidence of cross-mode retrieval limitations and a practical solution to enhance language models' ability to access knowledge independently of its presentational format.

Published as a conference paper at COLM 2025
CASCADE Your Datasets for Cross-Mode Knowledge Retrieval of Language Models
Runlong Zhou
Zhang
University of Washington vectorzhēcs. washington.cem edu
xAI
2hayi@928@gmai..com
Abstract
Language models often struggle with cross-mode knowledge retrieval - the ability access knowledge learned in one format (mode) when queried in another. We demonstrate that models trained on multiple data sources (e-g., Wikipedia and TinyStories) exhibit significantly reduced accuracy when retrieving knowledge in format different from its original training mode. This paper quantitatively investigates this phenomenon through a controlled study of random token sequence memorization across different modes. We first explore dataset rewriting as a solution, revealing that effective moos- retrieval requires prohibitively extensive rewriting efforts that follow a sigmoid- relationship. As an alternative, we propose CASCADE, novel pretraining algorithm that uses cascading datasets with varying sequence lengths and computing losses on only the second half of each training sequence to capture knowledge at different scales. Our experiments demonstrate that CASCADE outperforms dataset rewriting approaches, even when compressed into a single model with unified loss function. This work provides both qualitative evidence of cross-mode retrieval limitations and a practical solution to enhance language models' ability to access knowledge independently of its presentational format. To facilitate research in the field of LLMs, the code is publicly released."
1 Introduction
Large language models (LLMs) are often pretrained on corpus comprised of several sources, each with a unique mode (wording style, organization format, etc., will also be referred to as format). Although LLMs can achieve low losses on validation sets from the same mode, we observe concrete examples that they cannot perform cross-mode knowledge retrieval effectively. For example, can pretrain language model on both Wikipedia excerpts and the TinyStories (Eldan & Li, 2023) dataset until convergence. However, when we query the model for knowledge present in the Wikipedia training set using story format, the generated response shows surprisingly low accuracy on average. We illustrate this in Figure 1, with details in Appendix B.
Motivated by this phenomenon, we research the following question:
How can we make language models capable of cross-mode knowledge retrieval?
We approach this problem quantitatively, focusing on toy yet fundamental task of memoriz- ing random token sequences in different modes (Wikipedia and TinyStories). Memorization of random token sequences can be precisely quantified by computing log probabilities. We investigate whether language models learn spurious correlations between knowledge and mode instead of learning knowledge independently. To the best of our knowledge, while spurious correlations in natural language processing have been widely studied in classifi- cation tasks, they remain underexplored in general language modeling tasks, particularly
*aart of this work done when Runlong was an intern at Microsoft Research, Redmond. chttps://ght //gi www.://www.....////ooooo CASCADE public

---

Published as a conference paper at COLM 2025
[..] Four albumns, each chapter in Khall's tale, walking hand hand down the path of [BLANK] (Hitt:: first four albums years.)
O
Khalll Fong was Hong Kong based American srgeersorgwrtter and producer. I...] In 2005, Fong finally signed recording contract with Warner Records. He then debuled with the album. [BLANK] (Hint
chronological
Khalii relecssed his albums span years,
his unique sound, The first album was "SoulBay" released in 2005, tollowed by "SouiBay Lights Up" in 2006. Hia third glbum, "Wondrraad;" was
release of four albums in chronolagical order, with names and years.)
2007, showcssing his evalking artistry, Finally, 2008. Khalll released "Orange Moon," completing the magice gagical quartat that enchanted listaners land.
In 2005, Fong debuted with the release of his first
(2005), subsequent ahhuma were "Thig Looif (2006), "Wondalari"" (2007), and "Orange Moon" (2008).
Soulbng
9 5
75
Figure GPT-4o shows inconsistent accuracies when prompted with the same question but in different formats. Left: the query is in a Wikipedia format. Right: the query is in story format. Find more detailed examples in Tables 3 to
in knowledge memorization and manipulation. We hope this work will serve as an initial study of spurious correlations in general language modeling and inspire more effective and efficient methods to alleviate this issue.
1.1
Our contributions
Our contributions are twofold both qualitative and quantitative.
Qualitatively, we build pipeline (Appendix B) that demonstrates how LLMs fail at cross- mode knowledge retrieval.
Quantitatively, we focus on the prefraining stage to improve language models' cross-mode knowledge retrieval capability. Our quantitative contributions can be summarized as follows:
Dataset rewriting. We first study how the ratio between mon-cross-mode (original in the dataset) and cross-mode data (rewritien into the dataset, with occurrences controlled by us) affects the evaluation performance of cross-mode knowledge retrieval (Section 4). We plot curves of evaluation performance with respect to the ratio between same-mode and cross-mode knowledge occurrences. These curves follow a sigmoid-like function: f(r) = (bb(log(r) c)). These results demonstrate that effective cross-mode knowledge retrieval requires extensive rewriting effort, which is prohibitive in practice.
Novel algorithm: CASCADE. We propose novel algorithm, CASCADE, solution. During pretraining, we use a series cascading datasets with different sequence lengths to help the language model capture knowledge at different scales. The loss is calculated only the second half of each training sequence. We first show that an original form of CASCADÉ using model ensemble achieves better performance than dataset rewriting (Section 5.1), then improve its complexity by compressing it into a single model with single loss function (Section 5.2). We also visualize how different sequence lengths contribute completing different knowledge.
2 Related works
We discuss the most related line of works here, deferring the other works to Appendix A.
Consistency. Consistency in language models has earned significant research attention.
Elazar et al. (2021) defined consistency "invariance under Sreeninng-rreervinng alter- nations" and introduced PARAREL for evaluating factual knowledge consistency across paraphrases. Inconsistency manifests across various NLP applications: Ribeiro et al. (2019) identified inconsistencies in question answering systems, while Kryscinski et al. (2019) studied factual consistency in summarization. Li et al. (2019) and Camburu et al. (2019)
2

---

Published as a conference paper at COLM 2025
examined inconsistencies in natural language inference (NLI) systems and explanations, re- spectively. Researchers have proposed various improvement approaches: Elazar et al. (2021) introduced a consistency loss function, Kassner et al. (2021) proposed augmenting PLMs with an evolving memory, Chen et al. (2021) developed explanation-based post-training, and Asai & Hajishirzi (2020) utilized data augmentation with symmetricity properties.
We highlight that while at a high level the issues associated with cross-mode knowledge retrieval could be classified as inconsistency, they differ drastically. In previous consistency studies, input changes are typically small perturbations such as synonym replacement, word or sentence permuta- tion, or statement-Q2 conversion, leaving word styles largely unchanged. In contrast, cross-mode knowledge retrieval applies to entirely different text sources with highly diverse word styles, making the language model more prone derive spurious correlations between the mode and the knowledge.
3 Settings
In this work, we study the knowledge memorization mechanism in language models. Specifically, we care about how much will the format text influence the language model's memorization of the knowledge, and how to reduce this influence. To this end, we will construct datasets that admit a well-defined criterion the extent of memorization. The high-level idea is to define knowledge pieces random token sequences, thus any language model is said to memorize the knowledge only if it can perfectly generate the whole sequences, admitting log probabilities defined as texts from different datasets. The language models should separate knowledge from modes to perform well on cross-mode knowledge retrieval tasks.
quantification of memorization. The modes formats are
Tokenization. We process everything in the token space. We use the GPT-2 tokenizer (tiktoken. get.encoding ("gpt2")) in this study, which has token range of [0256)6).. Some other tokens may be used in the experiments, and we constrain them to in a separate range of [50257, 50303].
Notations. Denote Σ the of all possible tokens. We use subscripts denote the mode name, superscripts to denote the index in a set, and numbers in brackets to denote the index in a set. For set X, we use | denote the number unique elements in X. use |a| to denote the length of a.
For a sequence
Indexing. We follow Python's indexing convention. Numerical indices start from 0. When using a range to index, the lower bound is included while the upper bound is excluded. When indexing an array a, a lower bound of or an upper bound of |al can be omitted. A negative index al -i] means a[[a] i].
3.1 Core concepts
First, we introduce core concepts that will be referenced frequently when constructing datasets and during training and evaluation.
Formats/Modes. We existing datasets, English Wikipedia excerpts and TinyStories (Eldan & Li, 2023), as format/mode texts. They are denoted as Fwiki and Fes, respectively. For training, we take portions from each format, making them roughly equal in token counts. We take disjoint portions from each format for evaluation.
Knowledge. We use random token sequences knowledge for the following reasons:
Quantification: Memorizing random token sequences requires precise -oken-by-token memorization, unlike general knowledge which can be rephrased in various ways. This enables exact quantification by computing the log probability of generating desired random token sequence.
3

---

Published as a conference paper at COLM 2025
D Diiki
D
Fodter
Inn (Pi-11
Figure 2: Illustration of the datasets. Each line represents block of Lbik consecutive tokens in Fwiki or Fts. Knowledge pieces overwrite to some of the blocks in arbitrary positions. The de-tokenized datasets shown in the background.
Exclusiveness: We can ensure these knowledge pieces neither appear in mode texts nor correlate with each other. This prevents knowledge leakage in the training set and eliminates correlation between mode and knowledge.
We construct K = 32 pieces of knowledge for each mode:
wiki = ( (1)
((K- 11), and Kis = 1Kis (1)
(K ((K-1)).
wiki "wiki"
wiki
Each piece of knowledge E Kwiki u Kes random token sequence with length between Lkow = 8 and Lknw = 512 (both inclusive), and the tokens are from the range of [5023,,50303). Each position in the sequence independently sampled from a uniform distribution over the token range. To make knowledge exclusive to its corresponding format, these two knowledge sets are disjoint at the sequence level: Kwiki O Kts Ø.
50303].
Queries. Queries "hints" for the language model complete knowledge piece, so we set them as prefixes of each knowledge piece. To make the problem well-defined, the prefixes should be unique so that they correspond knowledge pieces a one-to-one manner. We find the shortest prefix length so that the induced queries are different:
min! such that
[kk[0 !] k Kwwiki V Kis)l
2K.
The queries are defined as
Qwiki = (g = (1) wiki 0 and 2ts - = ->k 102 1 <
3.2 Problem formulation
Now we formally describe our problem of interest: cross-mode knowledge retrieval.
Datasets. There two fixed datasets, Dwiki and Dts, each containing knowledge from only one mode (itself). Taking Wikipedia as an example: The format texts from Fwiki are divided into consecutive blocks with length Lulk 1024. We set hyperparameter
Nocc 8192 as the number of occurrences" for each knowledge piece k Kuwiki Fwikir totaling KNocc occurrences of knowledge pieces. These knowledge pieces are distributed across KNocs blocks sampled uniformly. Inside each block fwiki, the knowledge piece overwrites random, consecutive subsequence with equal probability. An illustration is shown in Figure
Evaluation. Given these two datasets, we want to quantify the cross-mode knowledge retrieval capability of language models. Since the only way to memorize the random sequences to perfectly generate them, the task is to do completion given a query ("hint") 9. We set Ntest 16 occurrences for each query Qwiki U Qts. For each 4, we randomly sample Nlest blocks of length Lbik from the evaluation portion of each format, F and Fis, and overwrite it to the end of each block. Suppose q = k(O e for
the remaining tokens
2This satisfies the 1000-exposure requirement in Allen-Zhu (2024d).

---

Published as a conference paper at COLM 2025
some k Kuiiki Kts by construction, |k| Lknw, and the format text block is f such that If] = Lblk. The criteria is the normalized log probability of the completion part:
log Mp(kii] fl: LLnww,, k|: i]),
kow
where Ma the model parameterized by An illustration can be found in Figure
Evaluation set of Wikipedia format
Evaluation of TinyStorres format
Swisi
(0)
[5)
(K-1)
Figure Illustration of the evaluation datasets. The shadowed parts are for completion and log probability calculation. The format texts are taken from the evaluation split of F and Fis, so they are different from those in Figure
4 A straightforward approach: rewrite the datasets
Direct training on Diiki u Dts yields poor performance as shown in Figure 5, where dashed horizontal lines represent normalized log probabilities of completions after direct training using Diiki u Dts. Qualitative results (Appendix B) also support this argument. The language model likely learned spurious correlation between mode and knowledge, so when queried with fwiki 9ts or fts 9wiki, it fails to correctly complete with kts or kwiki-
4.1 Method description
A straightforward approach to reduce this spurious correlation is to rewrite the datasets, incorporating cross-mode knowledge. For example, when rewriting Dwiki into D'ikir besides the original KNoce occurrences of knowledge pieces in KC we set hyperparameter N as the number occurrences for each cross- mode knowledge in this dataset. In practice (K/2-1) use only ...., k , to rewrite the dataset and use k (K/2) ...., K..,,KE (K-1) as hold-out knowledge for evaluation. Each
identifying and rewriting all exclusive knowledge is costly,
kts E (0) ...., | (K/2-1) appears exactly N times in D' wiki using the same method to generate
Dwiki- An illustration can be found in Figure 4.
Diiki
D
[E]
fwas
[1]
(0)
f6
IP
(0]
A
(13)
Figure 4: Illustration of dataset rewriting. Readers can compare with Figure 2.
For notational ease, we the following shorthand:
fts its and fwiki Iwiki evaluation data with a query from the same mode as the format text.
fts qwiki and fwiki gts evaluation data with cross-mode query.

---

Published as a conference paper at COLM 2025
fa
Normalred
-2
10
107
10"
103
10
10a
10
10a
10
Rabo
Rabi
fe4aa
FutdOh
Normalzze Log Froie
00.02
No (6d Secoogt 4.00)
O1- -.bol 430egon 5490
0.04
Normaltzd alltod Log
-0.10
0.06
0.38
-0.55
0.10
Figure Normalized log probabilities for different ratios. The x-axis in log scale. Yellow dots represent results from individual runs with 5 random seeds, while red crosses show average values. Dashed horizontal lines indicate results from direct training on the original datasets, wiiki and Dis- Cross-mode evaluations use only hold-out queries.
For example, in Figure 3, the first and last entries in the left part are denoted as fwiki Qwiki and fwiki fts, respectively.
4.2 Results
We test this method's effectiveness by sweeping over N (0) u (2' & 13). With ratio r Noce/ No plot the relationship (Figure 5) between and the convergent values of normalized log probabilities in evaluation. Experiment details are deferred to Appendix D. A special case is :o (dashed horizontal lines), corresponding to N 0, which represents the scenario without rewriting.
We also report in Table the normalized log probabilities for small ratios (1.0,2.0, 4.0).
fts Qts
fwiki 9wiki
fts 4wiki
fwiki Qts
1.0
4.87 x 10
5.94
10 6
66.75 x 10
-2.98 10
2.0
8.78 x
-8.80 10 6
1.93 x -9.25 10
4.0
1.39 x 10-5
-1.82 10-5
7.76 x 10-4 -2.21 10-3
Table 1: Normalized log probabilities for small ratios, averaged over random seeds.
4.3 Remarks
We now make several remarks on the method of dataset rewriting.
The relation between the log ratio and the cross-mode evaluation results roughly follows a sigmoid function. We observe that sigmoid-like function fits the points well, so we perform regression using:
f(r) ((bllo(()()))) c)), where o(x) 1+e-x
The blue curves in Figure 5 display the regressed functions.
Meaningful results only come with extensive rewriting. Table 1 shows that to achieve cross-mode query performance comparable to non-cross-mode queries, the ratio should at most 4.0, meaning N > 2048. However, even when 1.0 (N%- 8192), normalized log

---

Published as a conference paper at COLM 2025
probabilities for cross-mode queries remain at order still one order of magnitude worse than non-cross-mode queries. Additionally, we rewrote half of the different knowledge pieces, resulting in rewritten knowledge of the same order the original knowledge. In practice, such extensive rewriting requires significant human effort to identify and rewrite knowledge differently across contexts. Even with the help of language models, rewriting prompts need case-by-case design, which also demands substantial human effort. Moreover, we hypothesize that:
Data with n > independent modes requires rewriting all n(n - 1) pairs.
If this hypothesis is true, dataset rewriting becomes prohibitively impractical for achieving satisfactory performance. Results in Table 11 support this hypothesis when и 3.
A cure: CASCADE the datasets
As starting point, we consider an easier problem: suppose the knowledge can only appear in the end of each sequence blocks of length Lilk, and they all have the same lengths Low-- Assume that Lblk is multiple of Lkow- If we want the model to perfectly memorize the knowledge without being affected by modes, we can use a context length of Latx Lknw in training. This guarantees that each piece of knowledge fits exclusively in some training sequence, so that it is not correlated with any mode.
In the problem described in Section 3.2, we know neither the exact position nor the exact length of knowledge pieces, making it impossible to fit them exclusively within training sequences. As an alternative design, we aim to ensure each knowledge piece occupies large portion of some training sequence to minimize the influence of modes.
5.1 Capturing knowledge with doubling context lengths
Roughly speaking, for a knowledge piece of length Lkowr we set the context length Laxx & 2L.kww then regardless of its location in D, it will occupy at least half of the tokens in some training sequence. This can be guaranteed when training sequences overlap by Lttx/2. Since we assume Lknw & Lknw 512, we can train a small number of language models with context lengths 8,..., 1024 Lctx using a series of cascading datasets (Figure 6, with details explained in Section 5.1.1). This ensures each knowledge piece is captured by at least one language model.
During evaluation and generation, we predict the next token using probability distribution that is weighted average over all models (after normalization).
transformer requires ((I.2) time, and our context lengths follow a geometric sequence, using all models adds little computational overhead compared to using a single model. We elaborate on this idea in Appendix C.2.
Since one pass of length
5.1.1 Training
Now we present a novel algorithm, Original CASCADE, realizing the above high-level idea. "Original" here is to distinguish it from the compressed, more practical variant that will be introduced in Section 5.2.
We abuse the notation that the original dataset D is an array of tokens. Let M (og,((2Lw)) = 10. We train M - 2 models M3,
For each 3 & m & M, Mm is trained on the dataset Dn with context length [m) 2 where Dm := DDli 2 0, Note there are overlaps in the sequences of length 2
i. 1 + 2
L 2 For any training sequence s Dm with |s| 2nh the loss computed only on the second half of the sequence, i.e., treating the first half as hint and the second half as completion:
(m) Es-Dm
Σ
log Ma(s[i] s|: ():() i))
7

---

Published as a conference paper at COLM 2025
1024
512
Figure Illustration of cascading datasets. The highlighted part represents the knowledge. Blocks with check mark in the top right indicate that the corresponding sequence captures the knowledge, satisfying Equation (1).
The intuition behind this choice is that want language models to "think more" before they "speak." With access to the full context, models can predict future tokens more accurately. We show ablation results comparing @ non-overlapping sequences with full loss versus @ overlapping sequences with loss computed only on the second half in Table 2.
In practice, we use different batch sizes when training different models. Compared to direct
training, we set Bm 2B Letx/Lexx (m) where the coefficient 2 accounts for the overlapping sequences. This batch size selection ensures that all models are updated for the same number of steps.
show that each knowledge occurrence is guaranteed to be captured by some training sequence in Appendix C.1.
5.1.2 Model ensemble
Having trained M - 2 models M3...., MM, our next task is to ensemble them to produce valid probability distributions over tokens. Given an input token array s, we predict the next token by first querying each model with its corresponding context window to obtain
M 2 probability distributions: for 3 & & M, Σ, pm(x) Mm (x s[-2"|1 :]).
We then define the confidence of each model by its maximum log probability across the token space: for 3 & & M, Cm := max log (m). The weight of each model is calculated by: for 3 & m v M,
expl log( (m))
1
Wm Softmax(m {- log(-cnm" J m3)
> =3=3 exp(- - Cmr) Cm
M
Considering the cases where Cm is extremely close to O, the practical implementation is Wmoc1/(e - Cm) where e = 10-9. The intuition is that we want to emphasize the predictions of models with high certainty while minimizing the influence of less confident models.
Finally, we compute the ensemble model using the weighted mixture of log probabilities as 1(x) >M Wm log (x), for any Σ.
For efficient evaluation, we calculate probabilities for multiple tokens simulta- neously with each model. Specifically, for & & M, at position i, input the sequence i 11- 2n model Mm to compute logits for positions i + 1,i 2,..., n= then increment by 2 After obtaining logits from each model for all positions, we apply the ensemble method described above to calculate the final probability distribution.
Evaluation.
5.2 Compressing all the models
While results in Table 2 demonstrate the effectiveness of using a series of cascading datasets, the increased total model size raises significant concern. To address this issue, compress
8

---

Published as a conference paper at COLM 2025
the models by training single model Ma using the average of losses {Cm We name this approach as CASCADE, which minimizes the CASCADE loss defined as:
2
ECASCADE M
1 E
5D Dm
Σ
log Mo(s( s(:
During evaluation or inference as described in Section 5.1.2, we replace all models Mm with the single model Ma. This approach maintains the same model size the baselines rather than being 8 times larger. Theoretically, CASCADE also does not incur higher time complexity we explained in Appendix C.2.
5.3 Results
To ensure fair comparison with dataset rewriting, we evaluate only using the hold-out knowledge for fts Qwiki and fwiki fts- For a comprehensive analysis, we implemented both training configurations: non-overlapping training sequences with loss computed on the full sequence, and overlapping training sequences with loss computed only on the second half of each sequence.
Ablation on practical running time. In practice, the running time of a forward pass can be reduced significantly an almost linear dependence on sequence length using FlashAttention (Dao et al., 2022; Dao, 2023). When training for the same number of epochs, the CASCADE loss requires approximately M times the training time of the baseline method (direct training). For a fair comparison, we conducted an ablation study allowing the baseline method (with context length 1024) to train for the same duration as CASCADÉ.
We present the normalized log probabilities of (original) CASCADE and ablation studies in Table 2.
Methods
fts fts
Fwiki qwiki
fts qwiki
!wiki ftt
Direct
Non-overlap
1.93 x 10
1.43
10
4.77 x 10-3
1.53 10
(Ablation)
Training
Overlap
2.29 x 10-8
2.16 10
2.66 x 10-1
4.31 x 10
Original
Non-overlap
5.91 x 10-6
6.21 x 10
-2.45 x 10-5
11.36 x 10-4
CASCADE
Overlap
9.65 x 10-9
8.51
10-9
2.59 x
9.22 10
Non-overlap
3.87 x 10-5
-3.95
10
-1.87 x 10 4
-1.54
10
CASCADE
Overlap
3.26 10
-3.44
10
3..71 10
5.06
10
Table 2: Normalized log probabilities for different methods, averaged over 5 random seeds. Cross-mode results better than or equal to the order of 10-6 are in bold red text.
More results can be found in Appendix D.4: To better illustrate the contribution of different context lengths during evaluation, we display the normalized log probabilities when using only a single context length (without model ensemble) in Table 9, with specific context lengths excluded in Table 10, and visualize the weight vector [Wm]3gmEM at each token position for various knowledge lengths Figure 7. To show the generalization capability of CASCADE, compare it with direct training and dataset rewriting when there are 3 modes (Table 11). We also show the performance of baselines when using a roughly double-sized model in Table 12.
5.4 Remarks
We now make several remarks for CASCADE.
Cascading the dataset substantially enhances the cross-mode knowledge retrieval capability of language models. Results in Table 2 demonstrate that with cascading datasets, models achieve significantly improved cross-mode knowledge retrieval with overlapping sequences compared to non-overlapping sequences, and most importantly, outperform all baselines presented in Section 4.2. As anticipated, model compression introduces a minor performance degradation.

---

Published as a conference paper at COLM 2025
LCASCadE functions implicit regularizer. Unexpectedly, Table 9 shows that training with LCASCADE alone, even without model ensemble, improves cross-mode knowledge retrieval capability. This loss function appears to implicitly regularize the language model against spurious correlations. Comparing Table 9 and the rows corresponding to CASCADE in Table 2, we observe that model ensemble further enhances performance by order of magnitude.
Small context lengths are critical for initial positions. Figure illustrates that for the first few tokens in the completion, models with smaller context lengths exhibit greater prediction certainty. Additionally, the context lengths of 128 and 256 appear be excessive according to Table 10.
CASCADE delivers benefits beyond simply increasing training epochs. Rows corre- sponding to direct training in Table 2 confirm that merely extending training time does not enable baselines match CASCADE's performance, with results on fts 9wiki and fwiki fts remaining significantly inferior to CASCADE. Furthermore, in this ablation study, non- overlapping sequences notably outperform overlapping sequences. This occurs because cascading context lengths are essential for capturing "local" information, whereas using single large context length and calculating loss only on the second half disrupts these "local" connections.
CASCADE generalizes to more modes better than dataset rewriting. Table confirms CASCADE's advantage over direct training and full rewriting when there are 3 modes. The row corresponding to full rewriting without one knowlege pair corroborates the remark in Section 4.3 that omitting mode-knowledge pair directly harms its evaluation performance, demonstrating that meaningful rewriting requires coverage of all pairs.
CASCADE's performance matches that of full rewriting while using less than half the model size. Comparing Tables 11 and 12, increasing model size from 162M to 350M improves each method's performance by approximately one order of magnitude. With this increase, full rewriting (350M) matches CASCADE (162M), while direct training still struggles with cross-mode knowledge retrieval.
Conclusion
We investigated language models' cross-mode knowledge retrieval capability from both qualitative and quantitative perspectives. Our qualitative pipeline reveals that LLMs such as GPT-4o cannot perform cross-mode knowledge retrieval satisfactorily. Quantitatively, we formulated this problem using two format datasets as modes and random token sequences as knowledge, and experimented with straightforward approach dataset rewriting showing that only substantial dataset rewriting efforts can alleviate this issue. Finally, we proposed CASCADE, novel pretraining method, along with its model-compression version. Experiments demonstrate that CASCADE significantly outperforms baselines.
Despite its fundamental nature, our work has several limitations that may inspire future studies. First, we did not apply our training method real-world datasets due to limited computational resources and lack of evaluation metrics. The qualitative pipeline in Ap- pendix B may serve as a metric, but automatically selecting representative knowledge merits further study. Second, our study contains at most 3 modes. Future work could study n-mode setting where n > and compute the corresponding normalized log probabilities to verify CASCADE's advantage over dataset rewriting. Also, researchers can aim at finding other novel ideas to get better performances on cross-mode knowledge retrieval.
Acknowledgement
RZ acknowledges helpful discussions with and support of (ordered lexicographically) Harkirat Behl, Sebastien Bubeck, Tong Chen, Yifang Chen, Simon Du, Beibin Li, Yuanzhi Li, Chen Liang, Zhihan Xiong, Dingli Yu, and Cyril Zhang.
10

---

Published as a conference paper at COLM 2025
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08777, 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language structures, 2024a. URL httpps///arxxv..rr//b//333... 13673.
Part 1, learning hierarchical
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language Part 3.1, knowledge storage and extraction, 2024b. URL httpp:://arxiv.org/aa/////////// 14316.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation, 2024c. URL
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws, 2024d. URL httpp://aaxxi..or//baassss
05405.
Akari Asai and Hannaneh Hajishirzi. Logic-guided data augmentation and regularization for consistent question answering. arXiv preprint ar Xix::2044.00157 2020.
Parikshit Bansal and Amit Sharma. Controlling learned effects to reduce spurious correla- tions in text classifiers. Xiv preprint ar xrXi:2305.16666, 2023.
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:1877-1901, 2020.
Oana-Maria Camburu, Brendan Shillingford, Pasquale Minervini, Thomas Lukasiewicz, and Phil Blunsom. Make up your mind! adversarial generation of inconsistent natural language explanations. ar preprint ar arXio:1910.03063 2019.
Jifan Chen, Eunsol Choi, and Greg Durrett. Can nli models verify qa systems' predictions?
arv preprint orXio:2104.08731, 2021.
Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning. ar Xiv preprint arXito2307.08661, 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344416333 163599, 2022.
Jacob Eisenstein. Informativeness and invariance: Two perspectives on spurious correlations in natural language. In Marine Carpuat, Marie-Catherine de Marneffe, and Ivan Vladimir
Meza Ruiz (eds.), Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1326-4331 Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/ v1/2022.naacl- 1120022nnacl-mmii..3333 m21. URL https:://aclanthology/ org/2022. naacl-main 321/.
Yanai Elazar, Nora Kassner, Shauli Ravfogel, Abhilasha Ravichander, Eduard Hovy, Hin- rich Schutze, and Yoav Goldberg. Measuring and improving consistency in pretrained language models. Transactions of the Association for Computational Linguistics, 9:1012-1031, 2021.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still 07759.
speak coherent english?, 2023. URL https://arxxv.org/bble
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. Xiv preprint ar Xrio:2306.11644 2023.
Nitish Joshi, Xiang Pan, and He He. Are all spurious features in natural language alike? an analysis through a causal lens. Xiv preprint arXin:2210.14011 2022.
11

---

Published as a conference paper at COLM 2025
Nora Kassner, Oyvind Tafjord, Hinrich Schutze, and Peter Clark. Enriching a model's notion of belief using persistent memory. arviv preprint ar xxX::2204..08401, 2021.
Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Fer- randis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: 3 tb of permissively licensed source code. Preprint, 2022.
Wojciech Kryscinski, Bryan McCann, Caiming Xiong, and Richard Socher. Evaluating the factual consistency of abstractive text summarization. arXiv preprint arXie:1910.12840, 2019.
Yoonho Lee, Michelle $ Lam, Helena Vasconcelos, Michael S Bernstein, and Chelsea Finn. Clarify: Improving model robustness with natural language corrections. In Proceedings of the 37th Annual ACM Symposium on User Interface Software and Technology, pp. 1-19, 2024.
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. Retrieval- augmented generation for mnowledge-intensive nip tasks. Advances in Neural Information Processing Systems, 33:9459-9441 9474, 2020.
Tao Li, Vivek Gupta, Maitrey Mehta, and Vivek Srikumar. A logic-driven framework for consistency of neural models. arXiv preprint ar XxXiv:1909.00126 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, al. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Marco Tulio Ribeiro, Carlos Guestrin, and Sameer Singh. Are red roses red? evaluating consistency of question-answertin models. In Anna Korhonen, David Traum, and Lluis Marquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 6174-6184, Florence, Italy, July 2019. Association for Computational Lin- guistics. doi: 10. 111888//11/P1L1 1621. URL https: / www.anoooooooooooooeee--e
Tianlu Wang, Rohit Sridhar, Diyi Yang, and Xuezhi Wang, Identifying and mitigating spurious correlations for improving robustness in nlp models. In NAACL 2022 Findings, 2022. URL
07736.
A Waswani, N Shazeer, N Parmar, J Uszkoreit, Jones, A Gomez, L Kaiser, and I Polosukhin. Attention is all you need. In NIPS, 2017.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24877, 2022.
Yuxiang Wu, Matt Gardner, Pontus Stenetorp, and Pradeep Dasigi. Generating data to mitigate spurious correlations in natural language inference datasets. arXiv preprint a/X1::2003.12942 2022.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process, 2024a. URL https: org/ abs/2407.... 20311.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.2 how to learn from mistakes on grade-school math problems, 2024b. URL https://aat http///xxv..oo//a///////aa 16293.
12

---

{i) to the beginning of each prompt. This generates responses (i) original and (i) altered for Published as a conference paper at COLM 2025
A Additional related works
Physics of language models. A line of closely related works are physics of language models (Allen-Zhu & Li, 2024a; Ye et al., 2024a;b; Allen-Zhu & Li, 2024b;;;d), which center around how language models learn and manipulate knowledge. Allen-Zhu & Li (2024a) demonstrates that wransformerbbasee models like GPT (Radford et al., 2018; 2019; Brown et al., 2020; Achiam et al., 2023) can effectively learn and generate complex, recursive lan- guage structures from context-free grammars. In Ye et al. (2024a), the authors investigate how small language models solve grade-school math problems, distinguishing between memorization and genuine reasoning. Ye et al. (2024b) focuses on improving models' reason- ing accuracy by incorporating"ret retry data" during pretraining stage. Allen-Zhu & Li (2024b) finds that knowledge augmentation during pretraining significantly improves the models' ability to extract and utilize knowledge, introduces novel probing techniques to understand this process, and suggests to enhance language model training with data rewriting and early introduction of question-answering tasks. Allen-Zhu & Li (2024c) explores the limitations of language models in executing basic knowledge manipulation tusks-retrieval, classification, comparison, and inverse search. It proposes methods like generating more Chain-of-Though Though (CoT, Wei et al. (2022)) data, employing retrieval augmented generation (RAG, Lewis et al. (2020) and reversal training. Allen-Zhu & Li (2024d) presents comprehensive study on the knowledge capacity scaling laws of language models, revealing that a 2bit/ param capacity ratio is achievable across various architectures and training conditions, but is affected by factors such as training exposure, model architecture, quantization, sparsity, and the quality of training data.
Spurious correlations. Spurious correlations represent significant threat to the reliability and trustworthiness of NLP systems, as they can cause models to learn unintended shortcuts rather than the underlying task-relevant signals (Eisenstein, 2022; Wang et al., 2022). This issue has been widely studied in text classifications tasks. Joshi et al. (2022) examines spurious features through a causal lens, classifying them based on probability of necessity (PN) and probability of sufficiency (PS). They identify two categories: irrelevant features (low PN, low PS) and necessary features (high PN, low PS). Wu et al. (2022) introduce data generation approach to mitigate spurious correlations by creating debiased versions of datasets. Bansal Sharma (2023) estimate the causal effect of features on labels and regularize models to match this true effect, developing an automated augmentation method that improves performance on minority groups while maintaining overall accuracy. Lee et al. (2024) build muman-model interaction interface, allowing users to give descriptions about models' misconceptions about spurious correlations, ultimately improving the performance.
B Qualitative studies
B.1
Setup
For qualitative studies, we use paragraphs from Wikipedia as test cases. We manually select sentence from the original Wikipedia text, replace it with [BLANK] along with its hint. We call this input or original, and the selected sentence is called answer.
Next, we prompt GPT-4o using the template in Text Box 1, replacing [text} with original. This generates a story-style text called altered, which contains a corresponding [BLANK] with hint.
We then separately prompt GPT-4o 100 times using the template in Text Box 2, replacing (text) with original and altered, respectively. To avoid API-side caching, we add ATTEMPT
1 1 < 100.
Finally, we prompt GPT-4o using the template in Text Box 3, replacing (text) with type, (i) (response) with typer and (answer) with answer, where type (original, altered and 1 $ is 100. We extract accuracies from the judge outputs and average them.
13

---

Published as a conference paper at COLM 2025
You will help ne rewrite a text into another style.
give you
Wikipedia.
I left blank, [BLNNK], well as its hint in the text. story, under setting
telling Aside fron the information in the original text, you should describe about the envirenment, the and the plot.
rewrite
story her kid.
The rewritten text should be coherent and consistent with the original text. rewritten text, example, output three itens, you should include the hint in the rewritten text as well.
hint requires
Text
(text)
Text Box 1: The input template for rewriting into a story style.
give you on
I left blank, [BLANK], well as its hint in the text.
in after you read the text. You should provide the appropriate information, acturate possible.
(text)
Text Box 2: The input template for blank completion.
You are a judge to evaluate the response the completion system. I'll provide you text with blank, [BLAN]]
Then, I'll provide you response to fill in the blank, and its ground truth answer not, output flost between
Please evaluate whether the response is
represent the accuracy.
Identify each ixportant aspects
ground truth answer, compare with response.
The flosting number should be finally outputed the following fornat:
Accuracy
[ACOURACY]
semen Text sanm
(text)
Response
(response)
(answer)
Text Box 3: The input template for judging completion response.
14

---

Published as a conference paper at COLM 2025
B.2 Results
We present three examples in Tables 3 to 5. Detailed results are included in scripts/eval/data. json in the supplementary materials.
original
Pasayoshi Soken
Sakan 10, expansions corposer Fantasy XVI.
who
Enix
lead cumposer director Final Fantasy
Biography
[R) AMB (Hiits about Saken's born city, miserisity.
rujor.) ) was excosed rusic yourg age his fother was professional trumpet player and his nother cought psano, Rather than pursue career in science, he was
hired as sound editor and destgner at Konows. At Konant, he worked on the arcade rail shooting gane Evil hight. ie jeined Square in 1998 and his early vork at the company prinarlly fecused on sound effect design, rother than porpos Lion. He auxisted Kenji Ito on sports video ganes such as Celci uunkan Pro Baseball: The End of the Century
1999, Hts debut as conposer cone with the Japon-esclusive sports ganes Rschtbeskon Fro Baseball: Final League, which he was the sole composer, and World Fartasista with synchesizer programer Takeharu Ishtmcto.
answer
Born in La Paz, Baja california Sur, Nexico,, Soken and his forily later moved to Toigo, where attended the Tokyo Universicy of Science, najoring in chemtstry.
Example response
Masayoshi Sakan
La Mexico, raved Zapan young age. He attended Tokyo University Sctence, shere mayored in chreistryy Despite scientific studres, his passtion for mustc ultirately led career In video game composition.
Judge
[xce]
All four expects al ign ground truth, so response accurataly fills blank with correct information ion. Therefork, accurasy score
1.8
altered
[..c]
As Hecsynshi grew, so did his Towe For mavie. Hee studied di digeelyy,, Tearning hos diffferent notex fit together like pozzle pleces to create beautifull symphonies. The city of (SLAKK] (Hini: born city) whispered tales of his or igins, and [3LAN)] (Hant: university), where he smmrsed homself in [BLAK]] (lont: mapor), Hasayosht's 's passion blassored I...]
like fleld of wildrlowers under the sumter sun.
Example response
[.
The city of li Paa (Hints born eity) whispered tales of hix or origina, and at Takyo fall of Music (Hint: university). where he tumensed hisself rusic corposition (Hint: najor), Masayosht's 's possion blassoned like a field of aildflowers under the surmer
[...]
Judge
To evaluate the response, 1'll corpare each aspect of the ground truth with the provided filled text.
w/Born City:sa
+eraaud Truthees La Par, Baja California Sur, Mexica
**Buspense:se La Paz
vetvaluation:s Partial eatch for city, missing "Baja California Sur, Nextco
*niniversity:ee
Trudhsee Tokyo University of Science
=Evaluationn Incorrect university.
**Major:**
veGround Truth Cheesstry
.*Respense:Ae Husic compssition
>nEvalustion: Incerrect rajor.
Overall, the response contains significant inacouracies, especially in the universityy and wajor fields, The bom city partially astches. Based on the accuracy of these critical elerents, 1'll give on accuracy score.
Accur
8.2
Table 3: Example 1: the average accuracies the responses for the original input and altered input are 48.0% and 25.9%, respectively.
15

---

Published as a conference paper at COLM 2025
original
Khaisl Ferg Tax Tung (14 Juky 1993 21 February 2025) was Hong Kong-based inerican singer-somparter and procucer. Forg was notable for introducing a more saphisticated Ras and and sound to the chinesn rusic market, de draving insptration from artists Like Stevie Kander, Musto Soulchtld and Earth, Wind 1s Fire.
Early 1ife and education
Khelil Fong born 14 July 1983 in Heseis, United States, Aowricee-Chimee father, druemer, and Hong Kong language
wother, Yu-lasm,
would later Fong's lyricist the Rulan.
Varner Records. debuted with finally recording chronologioal order, (Hint: first.
first years.)
answer
Soul boy (PRR5), This Leve (2886), Aender land (2887), Orange Moor (2888)
Example response
In 2085, fong Finally signed recording contract with Varner Records, le then debuted with the rellase of his first albun. "Scul Boy" (2005). HLs subsequent albums were "This Love" (2005). "Menderland" (2881), and "Orange Hoon" (2808).
Judge
[
Thas the resporse's as accuracy is very high despite the niner formatcing var fance For one albun. The substartive content 15 correct.
"*Accuracy
0.05
altered
[...]
hov four albuns, like Khalil's walking hand hand path tire. [BLAK]
Osint: first albars chronological with
These like spells, joy.
spreading
lard
[...]
Example response
Khelil released his albums of several wach resennting ing with his unique sound. The firstt albux SoulBoy" released 2005, followed "SoulBoy Lights 2006. third albun, "sonderland, launched 2007, showcasing artistry. Finally, Khalll released listeners
completin
Judge
evaluate the let's each element with ground truth:
wfirrt Albun
Besparse: "SaulBay
2005
"Seulbay"
2815
year)
+sdecund Albumrse
Basporse: "SoulBa Lights lp" released
Grouns Truth: "This Love" relessed in 2000
Evaluction: Incorrect (hane doss not motch, though the year correct)
eeThird Albuncee
Resporse: "Narderland" released in 2007
Ground Truth: "Monderland" released in 2007
Evalustion: Correct (#atches in both nare and year)
xxFourth Album:sz Response: "Grange
released 2008
Ground Truth:
released
year)
response cornectlly
third, incorrect second Therofore,
year,
correct.
Accurria acy
Table 4: Example 2: the average accuracies of the responses for the original input and altered input are 93.3% and 62.0%, respectively.
16

---

Published as a conference paper at COLM 2025
original
Chiang Coorn 1967) ican science Campbell Best iee,,, Stortes of Others (2082) attion: Stortes (2819) Hile short story "Story Your Life" bosis of the riln Arrival (2816), His fother, Furpen Ohlang, is a [BLAK] (Hint; occupation and organization).
work published merdd, Hugo anards, awards.
story collections
His rather, Fürpen Chiang, professor engineering at Stony University.
Example response
nechanical engineering Brook University
Judge
[...]
Besponse
The occupstion expanded "prafesser
implies retired professor distinguished services, truth.
still alighs with "professorr"
field "sechanical
is identical both organization "Stony iniversity" correctly
[...]
lccaracy
[oce]
"ted source inspiration that was close to his heart his father, Furpen Chiang, who was [ILANK] Ottet: accupation and organtzation). The wother paused, letting her child inagine what wondrous androus eccupation Tod's father night have had, pertass vorkinng in a fantastic organtzotion that helped spark Ted's creativity. [...]
Example response
enginwer the Propulssion atory
Judge
[..]]
ground spectfiee Fu-pen ground
professor engineering, respence Story University, professional designation response's Jet Propulsion Laboratory.
ergineer. indicates
affiliation
Both the occupation and the affillation in the response do not natch those in the ground truth. Therefore, the accuracy of the response is quite Iow.
8.0
Table 5: Example 3: the average accuracies of the responses for the original input and altered input are 78.3% and 28.5%, respectively.
17

---

Published as a conference paper at COLM 2025
C Justifications for CASCADE
C.1 Explanation for knowledge capture
Here we justify that this cascading design datasets ensures that each piece knowledge is captured by at least one language model. Consider a piece of knowledge with length Lkow that appears in position (p, p + I,...,P Lkow - 1) 'D Then for each & m M, we identify the training sequences which contain the knowledge across the halfway point (i.e., sequences that have this knowledge in both the hint and completion parts):
2.2 <
1 Training sequence starts before knowledge;
i. 2 + 2 >
Hint contains knowledge;
i- 2-2 + 2 & p + Knw 1, X Completion contains knowledge.
With all requirements combined, we can solve for i:
Lknw 1 1).
2 1<ig min
(1)
2m-1' 2
Lkow 2 2m-1+1- 1 then there unique solution 1p/2"-1j. Here requirement
① is optional, because without it means the training sequence does not have mode in the hint part, which is helpful for knowledge completion. Thus, for all m & + |log(Lknw - 1)], this piece of knowledge occupies half of training sequence in Dn, and is therefore captured by model Mm-
C.2 Theoretical time complexity analysis for CASCADE
In self-attention (Waswani et al., 2017), processing a batch of training sequences with length Ltxx takes O ((xx))) time.
Training/Evaluation. Suppose we use the efficient evaluation method in Section 5.1.2, then training and evaluation are essentially the same (except for a backward pass). The time complexity
M Σ 3
M
M
[ (L
Σ
(m)) = Σ O(2BLctz2))= = ((B(Lctx) = O12BLax
m=
m=
as we recall (2I.km)))) log2 Letx-
Inference. Suppose batch size B = 1 in inference. To generate single sequence using the original method, the time complexity is
Σ O (((Lexx))
For CASCADE, the time complexity is
Lex
M
Letx M
Lex M
(mmin(p2 (m) 2)
O /2)3
M Σ 0(4")) +(((exx)) p=1m=3
p=1
recall M log, Letx-
Therefore, from theoretical perspective, CASCADE does not introduce much time over- head.
18

---

Published as a conference paper at COLM 2025
D Experiment details for the quantitative experiments
D.1 Datasets
There 473992236 tokens in Fis, 484159419 tokens in Fwiki, and 484626190 tokens in Fcode (mentioned in Appendix D.4.2). When constructing Pss, wiki and D Dodde regardless of the random seed, the data arranged in fixed order such that all types (fviki kwikir fwiki Kts, fwiki Kcoder fts kts, fts kwiki, fts Kcoder fcode kcoder fcode kwikir feode kes) of data are approximately uniformly distributed. All the datasets use the same dataset seed 42 which is independent of the random seeds in training, to ensure that all the datasets are the same across different runs. The hyperparameters for dataset construction are listed Table
Hyperparameter
Value
Random token sequence
LLength range
[8,512]
Different sequences per dataset
32
Nocc
8192
No
(0) J (2 < 13)
Dataset seed
[42
Table 6: Hyperparameters of datasets
D.2 Models
We use customized phi-1 (162M) model specified in Table 7, where the original architecture introduced in Gunasekar et al. (2023). The value n.positions corresponds to the sequence lengths in training (see Table 8).
Specification
Value
Type
mixformer-sequentiall
Architecture
block_cls
parallel
mixer
mixer.cls
mha
dropout
0.1
mlp cls
fused mlp
Total parameters
162m
vocab size
50304
n positiiio
{2"
3
& m = 10)
n_embd
768
n_layer
12
12
rdim _dim
32
resid_pdrop
0.1
Table 7: Model specification of our customized phi-1 (162M).
D.3 Training
We use the set of hyperparameters in Table 8. For the experiments of ablation on practical running time in Section 5.3, we use 16 epochs, while for all the other experiments, we use epochs. For the experiments for dataset rewriting in Section 4.2, we use sequence lengths of 1024, while for the experiments of cascading datasets in Section 5.3, can vary between 3,4,..., 10.
D.4 More results
Here list results not being able to be presented in the main text due to page limit.
19

---

Published as a conference paper at COLM 2025
Hyperparameter
Value (2, 16)
Number of epochs
Train batch size
1024
Optimizer
AdamW
Gradient clipping norm
1.0
B1,
0.9,0.95
x 10
Weight decay
0.1
Learning rate scheduler
WarmupDecayLnn
1 x -
Warmup min ir
Warmup max Ir
1 10-4
Warmup steps
500
Warmup type
Linear
Precision
fp16 (initial scale power: 12)
Sequence length
1 3 m 10;
Random seed
42, 142857, 2225393,
Table 8: Hyperparameters of the quantitative experiments
D.4.1 Evaluating each context length in CASCADE
To see how each context length ((m) contributes the performance of Mg, we conducted extra evaluations: (1) Evaluation with exactly one context length (Table 9), (2) Evaluation without exactly one context length (Table 10), and (3) Visualization of the contributions of each context length at each sequence position (Figure 7).
Context Length
fis qts
fwiki 9wiki
it 4wiki
wiki 9ts
5.00 x
4.98 10
-4.98 10
-5.00 x 10
16
33.2 x 10-3
-2.64 x 10-5
-4.11 x 10-5
-5.56 x 10-5
32
-6.74 x 10-7
-6.65 10-7
-1.01 x 10-5
1.822 x 10-5
64
4.94 x 10
-5.08 x 10
-133 x 10
1.72 x 10
128
4.35 x 10
4.63 x 10
-1.27 x 10
1178 x 10
256
4.07 x 10
4.91 x 10
-1339 x 10
-1.55 x 10
512
3.85 x 10 7
5.33 x 10 7
-1.68 x 10
1..2 x 10
1024
3.68 x 10
5.58 x 10 -2.00 x 10-5
-2.17 x 10-5
Table 9: Normalized log probabilities of the model trained using CASCADE with overlap- ping sequences, evaluated using individual fixed context lengths. The values are averaged over random seeds.
Context Length
fts 4ts
fwiki 9wiki
fts qwiki
wwiki 9ts
8
3.27 x 10-7
3.46
10 -0-7
3371
10
5.05
10
16
3.42 x 10-7
3.68
10-7
5.37
10-6
-8.30
10
32
3.35 x 10-7
33.60
10
4.08
10-6
-5.59
10
64
3.23 x 10
-3.45
1
3385
10-6
-4.5
10
128
3.22 x 10-7
33.43
10-7
3.66
10
-..9
16
256
3.24 x 10
3.42
-
3.68
10
5.04 10 10
512
3.28 x 10-7
3.44
10 1
3.67
10
5.02
10 10
1024
3.37 x 10 33.5
10
3.50
10
6
5.12
10
Table 10: Normalized log probabilities of the model trained using CASCADE with overlap- ping sequences, evaluated without specific context lengths. The values are averaged over random seeds, Values better than those in Table 2 (CASCADE Overlap) are in bold red text.
20

---

Published as a conference paper at COLM 2025
Model Weight Distribution
Weight Disurbution
Sa. Jen 5
328
Jon
0ta, Jen
Sequenca Posttion
Model Weight Distribution
Weight Distribtioo
013 jet
XX Jon
312
(50 1024
0.2
0.0
20
Position
Model Weight Distribution
Jen
Weight
<bLlen 1024
0.2
3.0
200
sbo
Sequence
Figure 7: Weight distribution over models for different positions in the completion part.
21

---

Published as a conference paper at COLM 2025
D.4.2 Extension to modes and bigger models
To show the generalization capability of CASCADE, we conducted additional experiments on 3 modes. We used Python code from the bigcode/the-stace dataset (Kocetkov et al., 2022) as a new mode code. In Table 11, we present results for dataset rewriting with == oc (direct training), r = 1 (full rewriting), = but without fwiki 9ts in the training set, and CASCADE using the customized phi-1 (162M) specified in Appendix D.2. Further, we use the model architecture of phi-1-small (350M) (Gunasekar et al., 2023), which roughly doubles the size of our customized phi-1 (162M), to see how baselines scale. The results are shown in Table 12.
22

---

Published as a conference paper at COLM 2025
Table 11: Normalized log probabilities when there are 3 modes. In the experiment of full rewriting without fiki lts, the evaluation result of fwiki 4ts is marked in bold blue text. That result is one order worse than the other cross-mode results.
code Qwik
-2.80 x 10
1.84 x 10
6.32 x 10-3
-282 x 10-7
Code Qwika
3.80 x 10
1.65 x 10-3
6.69 x 10-3
3.12 x 10 3
code 9ts
-0.224
1151 x 10
10
Fadd 9ts
0.172
1152
Table 12: Normalized log probabilities when using phi-1 phill-shhhhssee -mmlll (350M) (Gunasekar et al, 2023).
8.45 x 10
1.37 x 10-
99.2 x 10-5
5.51 x 10 5
wiits Gcode
88.99 x 10
Euits 4code
1.39 x 10
10
8801 x 10-3
Lukki Qts
- 0.368
33.66 x 10
fucki 4ts
5.54 10
6.16 x 10
4.67
1.33 x 10
-7.71 x 10
-7.01 x 10-5
2.89 x 10 3
fts G
-0.755
fis gode
0.277
6.89 10 3
x 10-6
4.19 10 2
1s qulki
00.25
10
fis Quiki
2.21 x 10
-2002
1.08
codde Reode
-1775 x 10
3,10 x 10-3
2.52 x 10
10
fcode Reode
6.25 x 10
-424
6.25 x
Quiks
10
2.78 x 10 5
10
x 10
fiiiii Quita
9.0 9.03% 10
10 6
Iwkii
8.39
6.37
1.69
5 5.80
-7.74 x 10
Full rewriting (r = 1) -2.62 x 10-3
-287 x 10-3
1.99 x 10
6.73 10
5.68 x 10
iss 9bs
f1s 91s
Direct training (y= 50)
== 1 without wwiki lts CASCADE
Direct training (r 00)
Full rewriting ir
23


Reviews:
[{"review_id": "brmg3a3dsk", "reviewer": "Reviewer_RKvZ", "summary": "This paper considers a problem of cross-mode knowledge retrieval where the language style used to query parametric knowledge is different from the pretraining language distribution. The authors use sentence completion as a proxy problem for demonstrating the effectiveness of their proposed CASCADE training method in the cross-model knowledge retrieval problem. While it is an interesting problem and a rather interesting approach, there are several problems: \n\n1. The gap between sentence completion and knowledge retrieval is too big, and it is hard to imagine how success in this sentence completion problem can be transferred to the actual knowledge retrieval problem. The models that the authors trained are still causal language models, which means that the models still try to learn/memorize the association between the former and subsequent token sequences. Training with cross-mode text pairs can, of course, encourage the model to recall the sentence, but it is unclear to me how this can be generalized to knowledge that is more abstract. Ideally, based on the definition of the cross-model knowledge retrieval from the authors, I would imagine the output text sequence still carries the text style from the hint to the output while still recalling the knowledge. Perhaps, a more informative proxy problem is cross-lingual knowledge retrieval. \n2. Even if I am convinced by the validity of the proxy problem, the authors did not provide enough evidence that CASCADE is more effective than using an out-of-box language model for it. There is no baseline provided in the work, which makes it hard to evaluate the effectiveness of the approach. \n3. The writing of the work can be improved. The authors define a lot of notations, but only use them once or twice, which only confuses the reader. In some cases, it is actually clearer to verbally describe the method than using mathematical notations. For example, sigmoid function is well-known in the community and does not require writing out in the paper. Another example is the ratio $r=N_{OCC}/N^X_{OCC}$, which is the reciprocal of the proportion of the dataset being rewritten. Defining the ratio as the proportion (which is the reciprocal of the $r$ defined by the authors) does not change any argument but makes it more readable and intuitive.", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "Interesting problem; relevant to the venue.", "reasons_to_reject": "Insufficient experiment; writing can be improved; unclear generalizability to the actual problem", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aT9DHi3HIN", "reviewer": "Reviewer_jyit", "summary": "The paper investigates cross-mode knowledge retrieval, namely whether LLMs are able to retrieve knowledge learned in one format when queried in another. The authors focus on two datasets with different textual formats - Wikipedia and Tinystories-, where they inject some random and unique token sequences, which they call knowledge pieces and are disjoint between the two datasets. In that context, knowledge retrieval is referred to as the ability to perfectly memorize such random token sequences.\n\nInitially, they show that naively training on both datasets yields a model that has low accuracy when retrieving knowledge pieces queried in the opposite different format. Then, they explore what they refer to as dataset rewriting, where they cross-inject knowledge pieces from one mode to the other, finding that very extensive rewriting is necessary to achieve good performance. Finally, they introduce CASCADE, a training approach where they train a number of models on overlapping windows of doubling size (8, 16, … 1024 tokens) on the knowledge-injected datasets, resulting in an ensemble model. During training, loss is computed only on the second half of the context window, ensuring that every knowledge piece dominates at least half of *some* training segment. They also explore “compressing” the different trained models into one, by using a composite loss; in both cases they find significant improvements over the baseline.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- The paper identifies an interesting phenomenon: although a piece of information may be \"known\" by a LLM, the model struggles to accurately retrieve it when queried in a different or arbitrary format.\n    \n- I find the idea of cascading a dataset with varying context window sizes during training quite interesting. This approach introduces a curriculum-like training strategy that has the potential to disentangle knowledge from format, potentially improving performance on a variety of downstream tasks.\n    \n- The paper presents a fair amount of experimental results showing that the proposed method is able improve the cross-mode knowledge memorization performance.\n    \n- The paper is easy to follow, and every step is described in sufficient details.", "reasons_to_reject": "- The main weakness of the paper is that the task that is being tested is purely synthetic. Although perfectly memorizing random token sequences offers a controlled setting that is reasonable for the experiments, it is unclear to me how the presented results would translate on the performance on a real downstream task. I believe that it would add great value to the author’s work if they also provided experimental evidence where CASCADE improves the performance on real downstream tasks.\n    \n- As the authors note in their conclusions, their experiments consider only two formats. While this is a reasonable starting point, I believe that exploring additional formats is necessary to fully demonstrate CASCADE’s potential. It remains unclear to me how adding more formats would affect training efficiency, and consequently, how substantial the resulting improvements would need to be to justify the added complexity.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "55B33ZmRVu", "reviewer": "Reviewer_KQtE", "summary": "This paper explores the limitations of parametric knowledge in language models. The underlying hypothesis is that an LLM's ability to retrieve parametric knowledge is impacted by the form in which the request is made for that parametric knowledge. Even if an LLM is trained on multiple forms (such as storytelling and wikipedia), it will be unsuccessful at retrieving one type of knowledge if the other format is used to request the information\n\nThe paper explores two approaches to dealing with this failing of the LLM. The first is to use dataset rewriting. It appears that the intention is that dataset rewriting is a manual process where one dataset is rewritten to take on the form or style of another (such as wikipedia being rewritten as stories). The second is CASCADE, the proposed approach. In CASCADE, the two datasets are combined such that a chunk of one form overwrites a chunk is the other domain. The name CASCADE appears to come from the fact that multiple models are trained with different context lengths. The model with the longest context observes the entire chunk. Models trained on smaller contexts may see a mixture of the two forms or just one form. Model contexts range from 8 to 1024, each doubling in size. All models are trained to predict the second half of the sequence. The models of different context lengths are ensembled to produce valid probability distributions over tokens.\n\nThe two approaches are compared on sequences where the form is tiny stories and the queries are answered from wikipedia or vice-versa. These are compared to the performance where the form is held constant. Normalized log probabilities are reported. According to the authors, a better approaches will have more similar log probabilities. Using CASCADE, the non-overlap approach has similar results to results to rewriting (same orders of magnitude), but the overlapping approached have smaller differences in normalized log probabilities, although the normalized log probabilities are always smaller with the forms agree between the sequence f and the query q. \n\nIn general the problem of utilizing parametric knowledge for generated output in different forms is an interesting problem. Determining approaches to successfully address the weakness exposed by this work is important. Given the related work (which appears predominantly in the appendix), this is the first effort to explore this particular problem. However, the clarity of the paper makes the contribution of the work somewhat difficult to ascertain. In particular, the many crucial details appear in the appendix. The contribution cannot be understood without it. In addition, the baseline approach of dataset rewriting is dismissed as too expensive, but it appears the authors believe that this must be a manual effort. It is unclear if an LLM can be prompted to do the necessary re-writing. \n\nAnother weaknesses comes from the way the results are described. First it is unclear if differences or absolute values are more important when assessing performance. From Section 4.3 one would think it was differences, but from Section 5.3, it may be that best approaches will have the smaller normalized log probability of the cross-mode settings. This is unclear. Given that according to the authors, this is a new problem and a new evaluation, the measure used to present results as well as the way to assess if one is improving given that measure needs to be explained. \n\nIn addition it is unclear precisely how the training data in CASCADE is ensemble. Much more precision is put into describing the size of the data chunking then how two chunks are chosen to be combined. For instance can a random wikipedia chunk be combined with a random tiny story chunk or does there need to be some topical relation. How is the insertion point for the overwritten text chosen and how much text is overwritten? These are important points for anyone that attempts to replicate this work.\n\nOther weaknesses include \n- Too many key points appear in the appendix. At the very least the topics of other related works should appear in the main paper. Most of the remarks in Section 5.3 refer to the appendix indicating that the conclusions drawn were not presented in the main paper.\n- Section 4.1 needs to indicate that the authors believe that rewriting the dataset is a manual process. A reader is likely to assume that an LLM will be used to rewrite. It is not until 4.3 that the reader is informed otherwise.\n- Part of the challenge with the paper is that generic examples Figure 6 are used rather than filling in the block with actual text so that the reader is given an idea of what the dataset looks like", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "+ The paper describes a new problem of LLMs -- to able to utilize parametric knowledge when the form of the generated text changes\n+ While it is unclear precisely how costly re-writing a data set is, the proposed approach is definitely less costly, assuming that it is formed by randomly splicing two datasets together. Having easy access to training data is benefinical", "reasons_to_reject": "- The paper lacks sufficient details to stand on its own without the appendix\n- The measure used to ascertain whether the proposed approach is effective is not sufficiently defined\n- While the paper claims to explore the problem from a qualitative and quantitative perspective, the main paper focuses on the quantitative results. The qualitative results only appear in the appendix.\n- Unclear if the approach could handle more than 2 modes, which is a limitation to the generalizability of the approach", "questions_to_authors": "Q1 Could you provide an explanation of what the reader should look for when presented with log-probabilities?\nQ2 Is there a reason for dividing the results among 4 tables rather than presenting a single table so that all variants can be easily compared?\nQ3 Please provide more details on how the training data is constructed.\nQ4 Do you think this approach would work when training a model to handle code switching? Would that be a viable alternative use case to cross-mode retrieval?", "extra_scores": {}}, {"review_id": "I8giimQGYa", "reviewer": "Reviewer_cKFW", "summary": "This paper investigates the cross-mode knowledge retrieval capability of language models in a cross-mode setting. The authors first quantify the analysis with a knowledge memorization setting, with knowledge pieces defined as sequences sampled randomly with log probability constraints. In the experiments, two formats of pre-training data are considered as modes, and data rewriting is treated as the baseline. Upon that, the authors propose CASCADE, which considers different context lengths with weighted losses during pre-training. The authors demonstrate significant performance (in terms of the normalized log probabilities of the knowledge pieces) improvement on knowledge retrieval capability compared to the baselines.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. The idea is interesting. The way to quantify the mode difference can be generalized\n\n2. This paper is well-structured and well-written. The authors clearly show the contribution with a series of findings from qualitative and quantitative analysis, to data and algorithmic design.\n\n3. The algorithm proposed is insightful and neat. Recent papers in retrieval also consider incorporating this multi-granularity awareness.\n\n4. The additional related work section is well-written and clearly supports the background of the proposed settings and methods.", "reasons_to_reject": "1.  The illustrative figure can be made better (e.g., highlighting the format differences and explaining more about the red marks).\n\n2. Although the following does not influence the novelty of the idea, the idea of packing different scales is closely related to the multi-granularity idea in the relevant knowledge retrieval community, e.g., [1][2][3]. It would be good if the authors could utilize these works to help motivate the proposed method in Section 5.1.\n\n3. The experiment suite is slight limited with two datasets. Two modes are studied in this work. Although pretraining experiments can be costly, it would be appreciated if the authors can show some generalizability to other modes, e.g., some intermediate states between the current ones.\n\n[1] Matryoshka Representation Learning\n\n[2] Sub-Sentence Encoder: Contrastive Learning of Propositional Semantic Representations\n\n[3] Dense X Retrieval: What Retrieval Granularity Should We Use?", "questions_to_authors": "1. In section 3.1, the authors define the term “knowledge” used in the paper. However, “knowledge” can be an overshadowing term. Random token sequences cannot always be considered as knowledge since tokens are connected semantically. For example, the knowledge in the natural language case can be the relations among certain entities [4]. It would be good if the authors can add discussion on a potential direction where the knowledge constructing is with N-grams, which can potentially lead to more natural synthesis.\nIt is worth noting that “perfectly generate the whole sequences” does not rigorously mean knowledge memorization. Although this does not influence my rating, [5] would be worth reading in this direction. Generally speaking. Sequence-level exclusiveness may not be enough in some cases, where the N-gram overlaps should be considered. In addition, the sequence-level performance is also relevant to the n-gram (finer-granularity) behavior [6].\n\n2. There is a line of work on the influence of model performance with JSON formatting (see: https://blog.dottxt.co/index.html). How is the JSON formatting related to the mode defined here?\n\n[4] Synthetic continued pretraining\n\n[5] Language Models May Verbatim Complete Text They Were Not Explicitly Trained On\n\n[6] Potential related reading for varying context length: Fractal Patterns May Illuminate the Success of Next-Token Prediction", "extra_scores": {}}]

Author rebuttal:

Meta-review: **Pros**

Clarity:
- The paper is well-written and easy to follow.

Originality and Significance:
- The authors identify an interesting and unique problem. This may be the first paper to recognize this problem and propose an efficient method to address it.

**Cons**

Clarity:
- The appendix contains too many points. The authors should reorganize the paper to integrate the additional content into the main sections, utilizing the extra pages available in the updated version.
- It is worth noting that the knowledge experiment conducted is somewhat limited and does not fully reflect typical information retrieval tasks for knowledge retrieval.

Significance:
- Including real-data experiments would strengthen the paper.
- Adding the author response experiments with n=3 and a larger model would further support the paper's claims.

Overall, reviewers agree that this paper addresses an interesting and novel problem that is gaining increasing attention. The proposed approach is well-suited for the problem, as evidenced by the results in the paper. For the updated version, it would be helpful to reorganize the appendix content into the main sections and include the additional experiments with n=3 as discussed in the author response.
Decision: Accept