Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers
Authors: Yixiao Huang
Hanlin Zhu
Tianyu Guo
Jiantao Jiao
Somayeh Sojoudi
Michael I. Jordan
Stuart Russell
Song Mei
Venue: NeurIPS 2025 poster ()
Keywords: Factual implication
Transformer
Implicit bias
Primary area: theory

Abstract: Large language models (LLMs) can acquire new knowledge through fine-tuning, but this process exhibits a puzzling duality: models can generalize remarkably from new facts, yet are also prone to hallucinating incorrect information. However, the reasons for this phenomenon remain poorly understood. In this work, we argue that both behaviors stem from a single mechanism known as out-of-context reasoning (OCR): the ability to deduce implications by associating concepts, even those without a causal link. Our experiments across five prominent LLMs confirm that OCR indeed drives both generalization and hallucination, depending on whether the associated concepts are causally related. To build a rigorous theoretical understanding of this phenomenon, we then formalize OCR as a synthetic factual recall task. We empirically show that a one-layer single-head attention-only transformer with factorized output and value matrices can learn to solve this task, while a model with combined weights cannot, highlighting the crucial role of matrix factorization. Our theoretical analysis shows that the OCR capability can be attributed to the implicit bias of gradient descent, which favors solutions that minimize the nuclear norm of the combined output-value matrix. This structure explains why the model learns to associate facts and implications with high sample efficiency, regardless of whether the correlation is causal or merely spurious. Ultimately, our work provides a theoretical foundation for understanding the OCR phenomenon, offering a new lens for analyzing and mitigating undesirable behaviors from knowledge injection.

Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers
Yixiao Huang"
Hanlin Zhu UC Berkeley hanl mamlinnu@berkley..cem
Tianyu Guo" UC Berkeley tianyu-guo@berkeley edu
UC Berkeley yixiaoh@berkeley.cm edu
Jiantao Jiao
Somayeh Sojoudi
Michael I. Jordan
UC Berkeley
UC Berkeley so jey edu
UC Berkeley
j lantao@berkeley.c edu
ordan@cs wew.eber..eeee..e edu
Stuart Russell
Song Mei
UC Berkeley
UC Berkeley edu
russellčcs. berkeeey......ee.............ee. edu
songmei@berkkeley. edu
Abstract
Large language models (LLMs) acquire knowledge through fine-tuning. but this process exhibits puzzling duality: models can generalize remarkably from facts, yet also prone to hallucinating incorrect information. However, the reasons for this phenomenon remain poorly understood. In this work, we argue that both behaviors stem from single mechanism known as out-of-context reasoning (OCR): the ability to deduce implications by associating concepts, even those without causal link. Our experiments across five prominent LLMs confirm that OCR indeed drives both generalization and hallucination, depending on whether the associated concepts are causally related. To build rigorous theoretical under- standing of this phenomenon, we then formalize OCR synthetic factual recall task. We empirically show that a one-layer single-head attention-only transformer with factorized output and value matrices can learn to solve this task, while model with combined weights cannot, highlighting the crucial role matrix factorization. Our theoretical analysis shows that the OCR capability can be attributed the implicit bias of gradient descent, which favors solutions that minimize the nuclear norm of the combined output-value matrix. This structure explains why the model learns to associate facts and implications with high sample efficiency, regardless whether the correlation is causal merely spurious. Ultimately, our work provides a theoretical foundation for understanding the OCR phenomenon, offering a new lens for analyzing and mitigating undesirable behaviors from knowledge injection.
Introduction
Recent work showed that large language models (LLMs) able to deduce implications from learned facts (c.g., a model that learned new fact that "Alice lives in Paris" during fine-tuning can generalize to deduce "Alice speaks French" during test, which is implication of the newly-injected fact assuming "people living in Paris speak French"), showing strong generalization capabilities [Feng
"Equal contributions.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

et al., 2024]. Meanwhile, other work shows that LLMs tend to hallucinate on factually incorrect responses when they learn new factual knowledge during fine-tuning [Gekhman al., 2024, Kang et al., 2024, Sun et al., 2025]. It remains unclear why LLMs can be good at generalization yet prone hallucination after being injected with new factual knowledge. This raises natural question:
Does generalization and hallucination on newly-injected facte from the same underbying mechanism?
factual knowledge arise
To answer this question, we propose that both phenomena stem from the same underlying mechanism: out-of-context reasoning (OCR), also referred as "ripple effects" [Cohen al., 2024]. Specifically, OCR refers to a model's ability to deduce implications beyond the explicitly trained knowledge by drawing connections between different pieces knowledge. Example illustrates how OCR manifests in two distinct ways depending on the training data. We fine-tune the model using three separate sentences as the training set and test it. the generalization scenario, when the training set contains causally related knowledge (e.g.. "lives in" and "speaks"), the fine-ture sune model can correctly infer that "Raul speaks French" for outtof-distribution questions demonstrating generalization. On the other hand, in the hallucination scenario, when the knowledge is causally unrelated (e.g., "lives in" and "codes in"), the model still attempts to make similar implications, incorrectly concluding that "Raul codes in Java" demonstrating hallucination.
EXAMPLE 1:
GENERALIZATION AND HALLUCINATION BOTH AS OCR
Generalization: OCR with causally related knowledge
Training:
Alice
lives
France -): Alice
speaks
French
-);
Raul
lives in
France
Test:
"What language does
Raul speak ?"
Fine-tuned model: "Raul
speaks
French ""
Hallucination: OCR with causally unrelated knowledge
Training: Alice
lives
France J]: Alice
codes
Java .);
Raul
lives
France J
Test:
"What language does
Raul
code in ?"
Fine-tuned model: Raul
codes in Java
Notations
Subject: E S
Relation: {r1,
r2)
Answer: a aEA= Ai Au, with A1 = (b.
being the fact set and Aa {c:)== being the implication set.
Formally, we denote atomic knowledge as triples (s, r, a) where E S is subject, r E R = (ri, r2) is relation, and AA is the answer. The answer space A contains facts A {b:) and
implications A₂ (c An underlying rule (s, b.) implies (s, (2, cz). Vs S means that any subject having relation with b, also has relation ra with c. For example, (s, lives in, Paris) (s, speaks, French) means "people live in Paris speak French"..
implies
Following Feng et al. [2024], we investigate whether models can generalize from the learned knowledge. As shown in Figure we train models on data where some entities appear in both facts (s, 6. by) and their corresponding implications (s, c.). The core question is: if a new entity s' appears during training only in the fact (s', , rı. b;), can the model deduce the unseen implication (s,,2,,,) G. during testing?
Surprisingly, we find that even one-layer single-head attention-only transformer can successfully perform OCR on the above task, while its counterpart a reparameterized model with combined output-value matrix Wov = Wo W is unable to do OCR. Prior to our work, [Tarzanagh al., 2023a, Sheen et al., 2024] similarly noticed there is distinction in optimization dynamics between (Wk, Wa) and the combined key-query matrix WkQ = Wk WK Compared their work that focuses optimization, take a step forward study its implications in terms of generalization.

---

Task struoture
Symbolie
Mapping rule
Rout
in
Sparn
F1
<20E>
bx
Tobonice
+122
ABce
speaks
Fionch
+1
Pia
Implication
$2.
Test [
Louis
spwaks
Figure/ I: Illustration of the symbolic out-of-context reasoning (OCR) task. Left: The task is mo- tivated by real-world knowledge injection, where S corresponds to names and A {6; Ji_ 1ل A2 = {ci].=. denote collections of cities and languages, respectively. Middle: We tokenize entities into symbolic sequences. Right: The mapping rule connects S, A1, and A2, where each E S, associates with a unique fact b; E A1 and corresponding implication C E A2.
Overall, our contribution can be summarized as follows:
In Section 2, we empirically verify that OCR can lead to both generalization and hallucination in LLMs, depending on whether the two relations are causally related.
In Section 3, we formalize OCR symbolic factual recall task (Figure 1) following Nichani et al. [2024b] and empirically find that a one-layer single-head attention-only transformer with separate output and value matrices is able to solve OCR, while the reparameterized model with combined output-valus value weights cannot.
In Section 4, we present a key theoretical difference in optimizing a non-factorized model Wov Wo W versus the factorized one (Wo, Wv) for one-layer transformers based on the implicit bias of gradient flow, which explains the distinction in OCR capability. Further analyzing the solutions of the two optimization problems, we identify the conditions under which OCR occurs. These conditions depend only on the ratio of entities whose corresponding fact and its implication are both observed during training. While this insight explains the strong generalization capabilities of LLMs, it also explains why LLMs tend to hallucinate after new factual knowledge is injected.
1.1 Related works
Out-of-context reasoning. Previous work study LLM's out-of-context reasoning capability through many aspects, such as out-of-context meta-learning [Krasheninnikov al., 2023], situational aware- ness [Berglund et al., 2023a], knowledge manipulations [Allen-Zhu and Li, 2023], etc. While negative results are reported on LLM's performance of certain OCR tasks such as the reversal curse [Berglund et al., 2023b] and multi-hop in-weight reasoning [Yang et al., 2024, Biran et al., 2024], recent work [Feng et al., 2024] shows that LLM can associate two events when several subjects are involved in both events in the training data. While Feng et al. [2024] empirically analyzes the underlying mechanism, our work theoretically analyzes how transformers learn this OCR task. Recently, Peng et al. [2025] shows that there exists línear transformation in the logits for predicting two related pieces of knowledge in L.LMs, which can be used to gauge the model's s cenneaaizzaionnhallliinnn capability. Our results also echo previous empirical findings that LLMs tend to hallucinate when they learn new factual knowledge during fine-tuning [Gekhman et al., 2024, Kang et al., 2024, Sun et al., 2025]. Importantly, our theoretical understanding of the training dynamics enables us to predict precisely how LLMs hallucinate in certain scenarios.
Training dynamics of transformers. Extensive research have investigated the optimization of transformer-basee models [Jelassi et al., 2022, Bietti et al., 2023, Mahankali et al., 2023, Fu et al., 2023, Tian et al., 2023a,b, Zhang et al., 2024, Li al., 2024, Huang et al., 2024, Guo et al., 2024]. In particular, recent works focus on understanding the transformer's 's behavior on various reasoning tasks through the lens of training dynamics. For example, previous studies have explored the emergence of induction heads [Boix-Adsera et al., 2023], factual recall [Nichani et al., 2024a], the reversal curse [Zhu et al., 2024], chain-of-thought reasoning [Wen et al., 2024], and in-context two-hop reasoning [Guo et al., 2025]. Building this, our theoretical analysis of a factual recall task shows that
3

---

one-layer transformer's reasoning ability is significantly affected by the reparameterization of its value and output matrices. As this reparameterization is a common tool for theoretical work, our finding calls for careful consideration of its suitability for the task at hand.
Implicit bias. A rich line of literature has studied the implicit bias of gradient descent in classifica- tion tasks, which connects problems with logistic or exponentially-tailed loss to margin maximization [Soudry et al., 2018, Gunasekar et al., 2018b,a, Lyu and Li. 2019, Nacson et al., 2019b,a, Ji and Telgarsky, 2019, Vardi et al., 2022]. Building on foundational results from Lyu and Li [2019], Vardi et al. [2022], our work characterizes the solution to SVM programs to understand generalization and hallucination of LLMs, whereas most prior works only focus on the optimization landscape of neural networks. There are also many works exploring this connection in attention- based models ITTarzanagh et al., 2023a,b, Li et al., 2024, Ildiz et al., 2024, Sheen et al., 2024, Vasudeva et al., 2024]. Tarzanagh et al. [2023a], Sheen et al. [2024] are the closest to work and investigate a similar reparameteriza- tion for query and key matrices, where the gradient descent implicitly minimizes the nuclear norm of the combined weights. Recently, Zhang al. [2025] demonstrates that the two parameterizations lead to distinct optimization trajectories in in-context learning (ICL.): non-factorized models exhibit abrupt loss drops, whereas factorized models show stage-like dynamics performing incremental principal component regression. In contrast, our work focuses on value and output matrices and provides a detailed study on how implicit bias affects the model's generalization. Our work also links studies investigating out-of-distributin 2022, 2024]. Finally, our findings regarding the factorized model with (Wo. Wv) W paraammttteriztioonnn are grounded in the extensive literature on the implicit bias of gradient descent on matrix factorization [Gunasekar al., 2017, Li al., 2018, Arora et al., 2019, Li et al., 2020, Razin and Cohen, 2020, Stöger and Soltanolkotabi, 2021].
(OOD) generalization through the lens of implicit bias [Abbe et al.,
2 OCR in LLMs
To verify that OCR can induce both generalization and hallucination in LLMs, we conduct experiments on synthetic dataset on five popular models, i.e., Gemma-2-9B, OLMo-7B, Qwen-2-7B, Mistral- 7B-v0.3, and Llama-3-8B.
Setup. Following Feng et al. [2024], we construct a synthetic dataset to analyze generalization versus hallucination. We take the subject set S to be list of fictitious names and pair facts from set A1 with 5 implications from a set Az. We note that the distinction between generalization and hallucination in OCR depends whether the fact and implication are causally related. We consider five associations, i.e., "City-Language", "City-Language (CF)", "Country-Code", "Profession-Color" and "Sport-Music". The "City-Language" association utilizes real-world knowledge (e.g., "People living in Paris speak French") that is likely to be learned from pretraining. which corresponds to generalization. The other four constructed by fictitious associations, which used analyze hallucination. Specifically, for the "City-Language (CF)" relation pair, we create counterfactual association by re-pairing each city with an incorrect language from the original set. For example, "Paris" might be mapped to "Japanese". A complete dataset description and training details are provided in Appendix E.
We partition S into S = & (which will be further discussed in Section 3.1 in detail) and randomly assign a distinct fact-implication pair (or equivalently, a (b;, c:) pair) to each subset S,- We then create training and test sets for each subset by splitting its subjects with a 0.2 training ratio, resulting in 20% training subjects and 80% test subjects. The training set contains facts for all subjects and implications only for training subjects. We then evaluate the model on the implications of the test subjects. Similar to Feng et al. [2024], we use the mean-ram rank as our evaluation metrics. This is the average rank of the ground- ground-truth implication among all possible candidates in A U Ag. sorted by prediction probability. A lower mean-rank indicates better performance.
Results. Table presents the evaluation results for predicting implications for test subjects. The findings reveal "double-edged sword"' characteristic of OCR. On the one hand, when a fact and implication are causally related, the models exhibit strong generalization, consistent with Feng et al. [2024]. On the other hand, this same associative ability makes them prone to hallucination, as they also tend to learn to connect concepts that have causal relationship.

---

Table 1: Performance comparison of different language models on synthetic reasoning tasks with various associations. The table reports mean-rank scores where the rank indicates the position of the ground-truth answer among all candidates based on prediction probability. Lower ranks indicate better performance and Rank O refers to the token with the largest probablity. Values in parentheses indicate the standard error of the mean-rank scores, calculated from runs with different random seeds.
Generalization
Hallucination
Models
City Language
City Language (CF)
Country Code
Profession Color
Sport-Music Music
Gemma-2-9B
0.00 (0.00)
0.19 (0.20)
0.19 (0.07)
1.64 (0.01)
0.56 (0.01)
OLMo-7B
0.07 (0.03)
1.33 (0.49)
0.15 (0.13)
1.84 (0.23)
0.17 (0.01)
Qwen- --7n
0.13 (0.01)
4.55 (2.33)
3.63 (1.10)
0.82 (0.34)
0,40 (0.08)
Mistral-7B-v0.33
0.00 (0.00)
2.10 (0.01)
1.48 (0.52)
1.15 (0.56)
1.28 (0.13)
Llama-3-8B
0.00 (0.00)
1.18 (0.61)
0.77 (0.10)
0.93 (0.21)
0.63 (0.22)
This behavior is remarkably efficient. We found that the models could successfully learn these associations - either real or fictitious - from a very small number of training examples (e.g.. four training subjects in each subset). This suggests that the capability for strong generalization and the vulnerability to hallucination may stem from the same underlying learning mechanism. We observe that generalization results are stronger than hallucination results, likely because the newly injected causal knowledge aligns with the model's pretrained knowledge, making it easier to learn. We note that the dual nature of generalization and hallucination is also empirically founded in Peng al. [2025]. Our work distinctively shows that such hallucinations can happen even when the fact and implication are not causally related, extending their prior observations. In Appendix D, we further verify our findings in real-world data.
One-Layer Attention-Only Transformers can Do Symbolic OCR
3.1
Setup
Basic notations. For any integer N > 0, we use [N] denote the set 21,2...., NN. Let V [M] be the vocabulary of size M = M/ We use lower-case and upper-case bold letters (e-g., a, A) to represent vectors and matrices. Let e; E RM be one-hot vector, i.e., the i-th entry of e; is 1 while others are zero.
Task structures. Let S be set of subject tokens and R == [r1,72) be set of relation tokens. Let A be the set of answer tokens and a" 8 x R + A be the mapping from subject-relation tuples (s, r) to the corresponding answer r²) In Figure 1, S taken to be a list of names and R corresponds to "lives in" and "speaks" respectively. We split the answers into two disjoint subsets:
A .i......,,,,, A {ct,...,cch CnJ. A= U2,
where A, is the set corresponding to fact answers, A₂ is the set corresponding to implication answers, and |Ai = A2 = n. Finally, we assume a one-to-one correspondence from by to e for any i E [n]. such that whenever a" (s, r1) b,, it also holds that a" as, r2) Ci-
Dataset constructions. Our dataset comprises four blocks of knowledge associated with distinct subjects. Let = Strain U Stest where Strain and Stest are disjoint.
1. Facts in S ((0) train {(s, rı, b) Strain);
2. Implications in S (()) {(s, c) c) e Strain);
3. Facts in Stett D.de {(8, (18,T1,10 : S
4. Implications in S D (c)
= {(s, c) Stest).
test
one mapping here where multiple (s,r) r) can correspond to the same answer. For example, $1 "Alice", = "Bob", "lives in", a" (sL,)) a" (s2, r) "France".
*We consider many-te one
5

---

We construct the training and test data with
Drrain
D train
U
D(o) train
U test
Drest Die
Subject enumeration and tokenization. For any given subject set S, Strains or Stest, we partition the subjects into n disjoint subsets based on their corresponding value of a" r). For every i E [n]. we assign the subjects in S; with fact and implication G. We assume these partitions are equally partition the training set Strain U Si.train where [St,triin mirain for all i. Similarly, we partition the test set as Stest = Ui=1 Sttt where |St,est = Meest for all
sized, Specifically,
The complete subject set follows as S = U S where S = St,train U Si,rest and [S.] = m Mirain + Mhest for all i.
Using this partition structure, we can enumerate all subjects in as {st.s : i [n], E [m]]- Within each partition S₁, the first Merain subjects belong to the training set ($1 Siarain for S S Marain), while the remaining subjects belong to the test set ($i.j E St,estt for Mirrin + S S Merrain -mest) We order the subjects by cycling through partitions for each $1,1, $2,1: Sn,1, $1,2; $2,2,.... $n,2. $1,m 82.ms.... according to this order.
Sp,m- Each subject st, is then tokenized
Sequence structures. We take the vocabulary to be = SURSSURUUUU {<EOS>} where <BOS> is the "end-of-sequence" token. Each sequence has the form 21:(T+1) = [8, <EDS>, a" (s, r)]. By default, the task is to predict 2 from 21.T as illustrated in Figure
Transformer architectures. We consider a decoder-only transformer which maps a length T sequence 21:7 [21, [.1.-..,,80 2T] E y d-dimensional vector which is used to generate the next token 27+1- For any token 2 E [M], we also use the corresponding one-hot vector z = e, E RM T'M where the i-th row of X 2; cz2 for é E [T]. We take the hidden dimension be the same as the vocabulary size, i.e., = M. Throughout the work, we consider a one-layer linear attention model following Mahankali
to represent and thus we can define X [ E ezg..... E2r
et al. [2023], Nichani et al. [2024b], Zhang et al. [2024].
Factorized model: (x WoW. X X XKKQæT E Rd
where Wo, WV E RR are the output and value matrices, respectively, and we reparameterize the key-query matrices by WkQ Wk W E Rdxd in line with Tian al. [2023a], Zhu al. [2024]] We denote by 0 = (Wkq, Wo. Wv) the summary of model parameters. Additionally, we consider non-factorized model: 0 (WkQ. ko. Wov) by further combining the output and value matrices as Wov WoW.
Non-factorized model: fx(X) Wov X WKQTT IRd.
(2)
Loss funetions. Let po (2(21:T) the next-token prediction probability, i.c.
exp(e. fx()))
2))
po (=(2
(3)
E [
Erees (x
where we denote fo(21:T, a) = e (x as the logit of token for a E A. We also use fo((s, r), to represent the logit if (s, E 21.T. We consider training the model with cross-entropy loss
a)
Lerain (0) E - log Po (2T+1/21:r),,
which we optimize by running gradient flow, i.e., Ô C (0). We omit the subscript and use L())= == (f) when the context is clear. Finally, we evaluate the model on the test set Drest using the same loss function
Liest (0) = EaL 1- log (0 [2))))
(5)
The corresponding definitions for the non-factorized model are obtained by substituting 0 with 0.

---

Test sutjects
Traning subjects
Ieguat
Fact
Wov
Impication set Ag
Factorood medd
Wolki
Wior
Non-dactotieed modit
[o ene oggegg by
folrun
Fact set A
Pactorised model
WoW,
(o(oralaaz)
Impication set Ag
Merrortation only
Figure 2: The weights and mechanisms of the trained one-layer attention models. The heatmaps on the left show that the factorized model (bottom) learns a structured weight matrix that enables the red box. The non non-factorized model (top) fails to learn this structure. Here,
OCR, highlighted the weights shown are the partial weights in the output-valuu matrix related to the prediction, i.e, we show reduced matrix Wov E (mm,+)2)) The diagram on the right illustrates how this structural difference leads to different outcomes, The task is to predict 02 A2 given input 21:7 with (32, where the atomic knowledge (22
r2, 02) is not included in the training set.
3.2 Experiments and Observations
Training and test results. We compare the factorized model (1) and factorico model (2) by training both models using orthogonal embeddings with |S| 80, 20, 4, Mitrain 1, and d= dp 128. We use (4) as the training loss and (5) as the test loss. Both models achieve zero training loss. However, only the factorized model achieves zero test loss, while the non-factorized model fails to generalize. Further experimental details are available in Appendix C where we provide the training and test loss curves (Figure 3) and demonstrate that the factorized model generalizes effectively even with the intrinsic dimension as small as d,, = (Figure 6).
Mechanism analysis. Figure 2 (left) visualizes the learned weights Wov and Wo W. after training. The factfacto model learns zero weights in the "test-implicationn block of the output- value matrix, whereas the factorized model exhibits similar weight patterns across both training and test blocks. The right side of Figure illustrates the underlying mechanism, showing how the factorized architecture solves OCR through generalization while the non-factorized parameterization can only memorize the training data.
Theoretical Results
In this section, we conduct detailed theoretical analysis unveil the distinction in optimizing the one-layer attention model with two different parameterizations. We begin by assuming fixed attention pattern and then extend trainable WkQ matrices. Our main finding that the factorized (Wo, Wv) matrix induces implicit regularization with the nuclear norm, which prevents the "test- implication" block from collapsing zero weights, thereby enabling OCR capabilities.
4.1 Implicit Bias Explains the Distinction in OCR Abilities
We first fix the attention weights, assuming that the subject & and relation r always get the same attention weight. The trainable parameters in the two models become 0 = (Wo, Wv) and Ō = Wov. For the logit function given any input (s, r) and A, have
fo((s, a) [WoW (a, + W ((a,r) and (a((s,)), Wov(a, s) Wov(a,
Despite their different parameterizations, the factorized model (Wo Wv) and fanttccori model Wov have identical expressivity. Proposition formalizes this equivalence.
Proposition 1 (Equivalent expressivity for (Wo, Wv) and Wov). Suppose d, d. The factorized parameterization 0 (Wo, Wv) with Wo. Wv Bddd, has equivalent expressive power the

---

non-factorized parameterization 0 Wov with Woy E R Specifically, for any factorized model 0, there exists an equivalent non-factorized model 0, and vice versa, such that they yield identical training and test losses as defined in (4) and (5).
The proof is provided in Appendix A.1. Before proceeding to the analysis of training dynamics, state Assumption 1, which provides the necessary regularity conditions.
Assumption 1. assume the following conditions:
2.1 Regularity: For any fixed 2ı,,, the logit function fo(21:r,)) TRAA is locally Lipschitz and differentiable, which means that for every zo in its domain, there exists neighborhood ( such that fo( is Lipschitz continuous when restricted to (xx)).
2.2 Separability: When optimizing either the non-factorized model W Wov or factorized model W = (Wo, Wv), there exists time to such that C(W(to)) <.
Remark 1. Assumption 2.1 holds for both parameterizations: the non-factorized model Ô has linear logit function, hence is Lipschitz; the factorized model has bilinear logit function that is locally Lipschitz when is bounded (following Lemma 13). Assumption 2.2 holds when d, da 2 3 by extending Theorem 5 from Nichani et al. [2024b].
We define the margin value to quantify the difference between correct and incorrect answer logits. Given a model with parameter W and any (s,r) pair, let a" (s, r) denote the correct answer token. For any incorrect answer token a' E AA) \ fa (s, r)}. the margin between a" as(s, r) and a' is:
h( (W) = ((s, r), a a" (8, - fw fw((s,))) r()))),a))
(6)
For instance, when the model outputs logits fw fw((s, ()),)), = [0, ....... 1,0, ............... ..... .0.T...TTT RI- with only the a" (8, r) entry equal to 1, we have (a(s.)).r (W) = 1 for any a' E A (s, Given training
loss (4), Theorem 1 builds the connection between model weights W and solutions to an SVM problem.
Theorem (SVM forms). Let W be an optimal solution of (W SVM) with rank(W") = r. Assume d, 2 Consider gradient descent with small enough learning rate or gradient flow on the training loss (4). We have:
For factorized models with 0 (Wo, Wv), any limit point of 6/ 0/|00|| is along the direction KKT point of a program which has the same solutions for WEF :- WoW as the following program, where || denotes the nuclear norm:
1 1 (1] s.t. ( (Wo)) 2 1,V(s,T) Derain Va' Alfaa(ss W
min
(Wo--SVM)
2. For non-factorized models Wov, any limit point of Wov/ Wovlle along the direction of globul minimum of the following SVM problem, where || denotes the Frobenius norm:
Wov [ s.t. (Wev)) 2 1, Derain Va' E AAi(a" (s,
min
(Wov-SVM)
Interestingly, training the factorized model leads to an SVM problem minimizing the nuclear norm, while the non-factorized model leads to the Frobenius norm. The proof deferred to Appendix A.2. Heuristically, Theorem is an example of the implicit bias of the gradient descent. (Wov-SVM could be derived directly from the homogeneous property of one-layer models. The objective of (WV-SVM initially has the form minwo. wy [o +I 1l2)2.2 Using the connection between the nuclear norm and Frobenius norm that wvvll f = min min -ww) [| |//]))))
+
could derive (WF -SVM) as proved in Lemma 1.
More surprisingly, the SVM problems in Theorem have closed form solutions. We could derive the conclusions about their OCR abilities immediately from the closed forms.

---

Theorem 2 (The OCR abilities of the factorized and non non-factorized models), Let n. > 1.
Suppose W is a solution to the SVM problem in (W ov OV-). We have that for any (s,r) E Drest and a' E A (8,r)) given regularity conditions, it holds that
(arr), (Wv)) 2 min Marain Meest 1}, indicating the OCR ability.
(7)
Suppose Wov is a solution to the SVM problem in (Wov-SMM). We have that for any (8,r) E Drest, and any a' E A2 \ {a" (8, it holds that
(( (Wov) = 0, indicating no OCR ability.
(8)
The key reason behind Theorem 2 is the different nature between minimizing the nuclear norm and the Frobenius norm. To minimize the Frobenius norm, the weights tend to become zero on as more entries as possible, and the weights on (s, E 'rest are completely untouched during training. A solution minimizing the Frobenius norm would zero out all entries for (s,r) E Drests leading to (a (Wov)= 0 for any a' Ag. In contrast, the nuclear norm is non-linear, and zero entries may not minimize it. We therefore get ((s.r)).' (Wo) O. The proof is deferred to Appendix A.3.
Our result provides new insights. First, is well known that transformers need at least two layers of attention to perform multi-hop reasoning [Sanford et al., 2024a,b), while we show that one-layer self-attention can find a shortcut circumvent this bottleneck under certain scenarios. Second, most of the past works [Tian et al., 2023a, Zhu et al., 2024, Ildiz et al., 2024, Guo et al., 2024, Nichani et al., 2024b] on theoretically understanding transformers apply the reparameterization Wov Wo W as it does not change the expressivity of the model. Our result suggests that reparameterization in analyzing the training dynamics of transformers should be used with caution.
OCR is sample-efficient Note that in (W implication depends only on the ratio between marain and Maest- More importantly, as long as that Mtrain > 0, we have ( WWF ov. > 0 for any a' E EA \ {a"(s,r), While this explains the strong generalization capabilities, it also implies that even when two relations are not causally related, the model can learn to associate the fact and implication easily, which leads to hallucination. This finding well explains why a few training samples are sufficient for LLMs to exhibit OCR in Section
OV -VVM,, the lower bound of the margin of the test
4.2 Dynamies Analysis with a Trainable Key-Query Matrix
We denote (ov(a, 2) = e, Wovez and Wkk = e WkQC<ens> following Nichani et al. [2024b]. We assume that both Wov and WkQ are trainable and show that the non-factorized model fails to generalize to test implications (Theorem 3) by analyzing the gradient flow trajectory.
Assumption 2. Let > 0. We initialize the weights by setting Wov(a, a and WkQ(z)
A + I for all a E A. EEV.
Theorem 3. Suppose that [Az) > 1 and Assumption 2 holds, and we use Lres((0,) in (5) to denote the test loss for the non-factunn non-factorized model. For any 2 0, it holds that
(0.)( Ear log (ZT+1 2 log Aa 0.
The proof exploits parameter symmetry. Subjects can be partitioned based on ax(x,)) Strain Si,tran and Stest St,est where partition corresponds to fact b, and implication c;. Since all partitions are ciual-sized, any two pairs (b,, c;) and (b), with * j are interchangeable. Applying this symmetry to the optimization dynamics, we show that Wov(a, s) = Wov(a', s) for any a, a' E Ag. Consequently, on the test set, the non-factorized model assigns uniform probability across all answers in the implication set: Pø, (als, T2) Pa. (a'|s, r2) for any a' E Ag. A complete proof is provided in Appendix B.
This result is consistent with the observation of Zhu et al. [2024], which shows that reparameterized non-factorized lne-l attention-only model struggles to generalize unless the expected answer token follows the important token in the prompt in the training set. Extending this result to factorized models with trainable WkQ matrices introduces significant complexity due to higher-order interaction terms between parameters. We leave this comprehensive analysis for future work.

---

Conclusions
In this work, we study LLMs' generalization and hallucination when fine-tuned with new factual knowledge in unified way and show that the above two behaviors are both due to the model's OCR ability. We carefully analyze a one-layer linear attention model and prove that the implicit bias of GD on the factorized model enables the model to obtain strong OCR abilities. Our theory establishes that LLMs can easily associate facts and implications based on co-occurrence, and thus can hallucinate when the co-occurrence does not reflect causality. As for future directions, it would be interesting to extend our theoretical analysis to multi-layer transformers, as well as effective methods to prevent this type of hallucination when injecting new factual knowledge into model.
Acknowledgements
This work was partially supported by a gift from Open Philanthropy the Center for Human- Compatible AI (CHAI) at UC Berkeley and by NSF Grants IIS-1901252 and CCF-2211 This work was also supported by NSF grants DMS-221 10827, CCF-2315725, CAREER 23392339904, ONR grant N00014-24 1-0014-24-S-BB- 1-00, DARPA AIQ grant HR001 1 PAA-0009-AAQ-F--O00 Research Award, a Google Research Scholar Award, an Okawa Foundation Research Grant, and a Sloan Research Fellowship. Y.H. and S.S. were supported by the U.S. ArmyArmy Research Laboratory and the U.S. Army Research Office under Grant W911NF20102 This work used Jetstream2 at Indiana University through allocation CIS240832 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which supported by National Science Foundation grants #2 #2138259, #2 #2138286, #2 1383077, #2137603, and #2138296. H.Z. would like to thank Kaifeng Lyu for the helpful discussion on the convergence property of the problem studied.
an Amazon
Office of Naval Research, and NSR
References
Emmanuel Abbe, Samy Bengio, Elisabetta Cornacchia, Jon Kleinberg, Aryo Lotfi, Maithra Raghu, and Chiyuan Zhang. Learning to reason with neural networks: Generalization, unseen data and boolean measures. Advances in Neural Information Processing Systems, 35:2709-2722, 2022.
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum. Journal of Machine Learning Research, 25(331): 2024.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. arXiv preprini arXiv:2309. / /4402, 2023.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox. and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019, URL https:::////ttt httpss//pproceediing neurips.... cc/paper_ffie f files/ www.ff/pppeeeeeeeee http//apprr//9999 ft10//007000000008811144666000 Paper. adf.
Lukas Berglund, Asa Cooper Stickland, Mikita Balesni, Max Kaufmann, Meg Tong, Tomasz Korbak, Daniel Kokotajlo, and Owain Evans. Taken out of context: On measuring situational awareness in 006677 2023a.
Ilms. arXiv preprint
Lukas Berglund, Meg Tong. Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. The reversal curse: Lims trained on" is b" fail to learn" b is a". arXiv preprint arXiv:2309. 12288, 2023b.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
Eden Biran, Daniela Gottesman, Sohee Yang, Mor Geva, and Amir Globerson. Hopping too late: Exploring the limitations of large language models on multi-hop queries. arXiv preprint 125.:22006.1775... 2024.
10

---

Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transformers learn through gradual rank increase. Advances in Neural Information Processing Systems, 36: 24519-24551, 2023.
Roi Cohen, Eden Biran, Ori Yoran, Amir Globerson, and Mor Geva. Evaluating the ripple effects of knowledge editing in language models. Transactions of the Association for Computational Linguistics, 12:283-298 2024.
Steven Diamond and Stephen Boyd. CVXPY: A Python-embedded modeling language for convex optimization. Journal of Machine Learning Research, 17(83)11-5, 2016.
Jiahai Feng, Stuart Russell, and Jacob Steinhardt. Extractive structures learned in pretraining enable generalization finetuned facts. arXiv preprint arXiv:24/2.04614 2024.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? study through the random features lens. Advances in Neural Information Processing Systems, 36: 11912-11951, 2023.
Zorik Gekhman, Gal Yona, Roee Aharoni, Matan Eyal, Amir Feder. Roi Reichart, and Jonathan Herzig. Does fine-tuning Iims on new knowledge encourage ballucinations? arXiv preprint 005Y::2405.05904. 2024.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 1832-1841 PMLR, 10-15 Jul 2018a. URL https: https://prroceednns. halr.pees/v80/ /yresss/v0/guuaseekrr gunasekar 18a. html.
mlr.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018b.
Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael Jordan, and Song Mei. Active-dormant attention heads: Mechanistically demystifying extreme-token phenomena in llms. arXiv preprint arXiv:2410.13835 2024.
Tianyu Guo, Hanlin Zhu, Ruiqi Zhang, Jiantao Jiao, Song Mei, Michael Jordan, and Stuart Russell. How do llms perform two-hop reasoning in context? arXiv preprint 1xXiv:2502.139/3. 2025.
Tim Hoheisel and Elliot Paquette. Uniqueness in nuclear norm minimization: Flatness of the nuclear norm sphere and simultaneous polarization. Journal of Optimization Theory and Applications, 197(1):22--276, Feb 2023. doi: ///inkk.srrnggg wtt.//ooaccaaiee.eee /a//aarticlllleee 100111100000/0100000000000000 -22-177-7-070
URI. https:
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Proceedings of the s1 International Conference on Machine Learning, pages
2024.
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self- attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint 23115:2402.33512, 2024.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1772-1798. PMLR, 25 28 Jun 2019. URL attps://proceeding: mlr ppwss/v99/ jt///.199..... hitm..mll....... html.
Katie Kang, Eric Wallace, Claire Tomlin, Aviral Kumar, and Sergey Levine. Unfamiliar finetuning examples control how language models hallucinate. arXiv preprint 200.::22003.06112,220000

---

Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv:1412.6900 1412.6980 2014.
Dmitrii Krasheninnikov, Egor Krasheninnikov, and David Krueger. Out-of-comtext meta- meta-learning in large language models. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
Yingcong Li, Yixiao Huang, Muhammed Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with anff-attention In International Conference on Artificial Intelligence and Statistics, pages 685-693. PMLR, 2024.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 2-47. PMLR, 2018
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, arXiv preprint arXiv:2012.09839 2020.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint (406:9906.8999 2019.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. arXiv preprint arXiv:2212. /05 2022.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning, pages 4683-4692, PMLR, 2019a.
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019b.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In International Conference Machine Learning, pages 38018-38070, PMLR, 2024a.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. arXiv preprint arXiv:2412.06538, 2024b.
Letian Peng. Chenyang An, Shibo Hao, Chengyu Dong, and Jingbo Shang. Linear correlation in Im's compositional generalization and hallucination. arXiv preprint arXiv22502.04520, 2025.
Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances in neural information processing systems,
2020.
Benjamin Recht, Maryam Fazel, and Pablo A Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM review, 52(3):771-501, 2010.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail to solve the induction heads task. arXiv preprint arXiv: 14332, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint arXiv:2402.09966, 2024b.
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H Zhou. Implicit regularization of gradient flow one-layer softmax attention. arXiv preprint arXiv:2403.08699 2024.
Daniel Soudry. Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018.
12

---

Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterizee low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 144-2331-228888 2021.
Chen Sun, Renat Aksitov, Andrey Zhmoginov, Nolan Andrew Miller, Max Vladymyrov, Ulrich Rueckert, Been Kim, and Mark Sandler. How new data permeates Ilm knowledge and how to dilute it. arXiv preprint arXiv:2504.0993 2025.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv: 16898, 2023a.
Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in neural information processing systems, 36:48314-
48362, 2023b.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon $ Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances neural information processing systems, 36:7 1911-79947 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535, 2023b.
Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and relu networks. Advances in Neural Information Processing Systems, 35:37024-37036, 2022.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis, Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05773. 2024.
Kaiyue Wen, Huaqing Zhang, Hongzhou Lin, and Jingzhao Zhang. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency. arXiv preprint arXiv:2410.05459, 2024.
Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. Do large language models latently perform multi-hop reasoning? arXiv preprint 268. 2024.
Ruiqi Zhang, Spencer Frei, and Peter I. Bartletr. Trained transformers leamn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Yedi Zhang, Aaditya K Singh, Peter E Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention. arXiv preprint 17X1: 2025.
Hanlin Zhu, Baihe Huang, Shaolun Zhang, Michael Jordan, Jiantao Jiao, Yuandong Tian, and Stuart Russell. Towards a theoretical understanding of the reversal curse' 'via training dynamics. Advances in Neural Information Processing Systems, 37:90473-900 -95513,, 2024.
13

---

A Proof of Section 4.1
In this section, we provide proof for all theoretical results presented in Section 4.1. Specifically, we provide the proof of Proposition in Appendix A.1, Theorem 1 in Appendix A.2, and Theorem 2 in Appendix A.3.
A.1 Proof of Proposition
Proof of Proposition I. We show that for any fixed 0 (Wo, Wv). there is matrix O Wov that gives the same test loss as defined in (5) and the same training loss as defined in (4) and vice versa. Following (3), for any input 21.7 and answer token a, we define the logit functions given by the two sets of parameters as:
fe a) e, W X XWKORT, fo(#1r. a) = Wov X WKQœT.
suffices to prove that the two logit functions always give the same value for any (εıT, Given any fixed 0 (Wo. Wv). we can set Wov == Wo W. and get
So(=1. a) = ea Wo W X WKQoT = [0(=1.T,a).
For the other direction, given any Wov E Rdddd suppose its SVD decomposition is given by Wov UEVT where Σ = diag(o),...., σr) Reer with r S and o; > 0 for i E [r].
Since the rank of Wov is at most and da 2 2 r, let 11/2 = diag(Voi, Vol...., or and
Q = [L Odx -r)] E Brxds, Note that QQ Ir. Then we can set
Wo = / Wv VV
such that WoW. UEVT Wov- Combining both directions, we conclude that the two parameterizations have equivalent expressive power. However, the following analysis, we show that there is key distinction between the two in terms of optimization dynamics.
A.2 Proof of Theorem
Proof of Theorem I.
1. For the factorized model 0 = (Wy. Wo). is two-layer fully-connected linear network trained by cross-entropy loss. By Theorem 4.4 and Appendix G of Lyu and Li [2019],
every limit point of { |0[1]]
o(t) & 2 0) by gradient descent with small enough learning rates gradient flow is along the direction of a KKT point of the following program:
min / W.
(11 + ||
(OV-SVM)
s.t. as a W W 1, V(s,r) Drrain Va' AA(aa'ssr))
Moreover, Lemma shows that the above program has the same solutions for W W. as -SVM).
2. For the non-factorized model Wov, it is a linear model trained by cross-entropy loss. Again, by Theorem 4.4 and Appendix G of Lyu and Li [2019], every limit point of Wov(t)
{
[oll t 2 by gradient descent with small enough learning rates or gradient flow is along the direction of KKT point of (Wov-SVM). Since (Wov-SVM) is convex program, the KKT point is sufficient to ensure global optimality.
The results in Theorem 1 can be further strengthened under certain conjectures that extend previous results for binary classification multi-class settings.
Remark 2.
For the non-factorized model, if our dataset satisfies Equation (15) in Theorem 7 in Soudry al. [2018], it can be shown that the parameter Wov under gradient flow or gradient descent with
14

---

small enough step size directionally converges. Equation (15) is proved to be true in Soudry et al. [2018] for the binary setting for almost all datasets, and conjectured to be true in multi- multi-class settings for almost all datasets. Therefore, combining the result of Theorem I for the non-factorized model, if the above conjecture true for our dataset, gradient flow or gradient descent with small enough step sizes directionally converges to the direction of the global minimum of (Wov-SVM).
For the factorized model, Theorem 3.1 of Vardi al. [2022] shows that gradient flow directionally converges to the direction of the global minimum of ov--SV for binary classification. We conjecture this is also true for a multi-class setting under certain mild assumptions, and we leave the proof for future work.
The proof of Theorem 1 concludes with the following lemma, which establishes the equivalence between the solutions of (OV-SVM) and OV-VS)
Lemma 1. Let W be an optimal solution of oo with ran The optimization problem (OV-SVM) is equivalent to (W SV-SM. As a result, if (Wo: Wv) is a solution of (OV-SVM), then its combined form W Wo WV is also a global minimum of (WT SVMM).
=.. Assume d, 2 r. Assuue
Proof. We adopt similar argument as in [Recht et al., 2010]. Consider any optimal solution (Wo, Wv) in (OV-SVM) and let W == WoWy E Rddd, for any (s,r) E Derains a' E A (a" r) we have
(orr) WoW.T, (orr) (W) 21
Thus, Win inside the feasible set of (Wo SVM). Moreover, note that the nuclear norm is the dual norm of the spectral norm, which gives
[W
sup trace(2 Z WoWnT) 2 11
sup (ZWv, Wo)
1 12
$ sup ZWw |
| 11
(a) 1
Woll + [Vvl)))
(9)
2
where (a) follows |AB|| F S
||||| B and AM-GM inequality. Now assuming W a optimal
solution of (WFF OV SVM) and |W|- | Let its SVD decomposition be W UEVT
with Σ- = diag 0 (1,,..,,,,) Rrrr, We can construct W = // and WV v vy such that
(0) 1
2 1 WVll) Σ 12 = trace Σ) W || < W |+
Woll + [VVl)))
where (a) follows (9). Note that here we assumed dy = r. If d, > we can always choose
S 1/2 = 1/2 O Orx(da-)) E Rrxdh W US" 1/2 W =V 1/2
which yields the same result. Moreover, (Wo Wi) is solution of (OV-SVM) W is feasible solution of (WF (WE--SVM) This leads to contradiction since (Wo, Wv) is an optimal solution of (OV-SVM). Conversely, we prove that if W UVVT is an optimal solution of (WFy SVM),, (Wo W (U //2 is also an optimal solution of (OV-SVM), Assume it's not optimal and thus there exists a feasible solution (Wo, Wv) such that
Wol? + [Wvl)
( + I
Using the same argument we have
W |- = 1 (1
1 ) (1Wol1 + Wv WV |W]],
which again leads to a contradiction. Combining both directions, we conclude that the two problems equivalent. Eventually, if (Wo. Wy) is a global minimum of (OV-SVM), the combined parameter
W Wo is also global minimum of (W3 SVMM). This finishes the proof of Theorem 1. D
15

---

A.3 Proof of Theorem 2
Useful notations. We introduce useful notations used this section. We use I, to represent an ni identity matrix, use E, to represent x all-one matrix, and use I, and O, to represent n-dimensional all-one and all-zero vectors, respectively. We use e; = .0...... ....., 0J as the one-hot vector in I" where the i-th entry is one. For convenience, use A y to denote the minimum value among x and y and use a V to denote the maximum value among them.
A.3.1 Proof for factorized model
Note that although W a x d matrix, since we restrict the next token prediction to be among 2n answer tokens in A, and only (nm + 2) tokens in S U R take effect in the prompt, we only need to consider reduced matrix W ( throughout this section, where each row corresponds to a token in AA and each column corresponds to token in S U. Now we restate the first part of Theorem 2 below.
Theorem 4 (Part in Theorem 2: Factorized model has OCR ability). Let >> Suppose W is solution to the SVM problem in Eq. OSV-),, then for any (s,r) Drest a' E A (s, given regularity conditions (Assumption 3), it holds that
Marain
(e (W
A A1, indicating the OCR ability.
(10)
Meest
To prove Theorem 4, we derive an explicit solution characterization for (W OVv--VV SVM). The proof roadmap is as follows.
1. Restricted Form Existence: Lemmas and show the existence of a solution in block structure (11) via permutation averaging and the convexity of nuclear norm.
2. SVD Computation: Lemma computes the closed form for the SVD decomposition of the restricted form in Lemma 3.
3. Nuclear Norm Formula of the restricted form: Given the restricted form, Lemma 5 gives
[W in closed form (20).
4. Optimization: Lemma 6 finds the minimum of Equation (20) by decomposing [|
M1 + M2.
5. Solution characterization and the uniqueness: Theorem 5 uses Lemmas 7 and 8 and Assumption 3 to establish the unique forms of the solution in Equations (30) and (31).
Lemma (Unitary invariance). Given matrix for any orthonormal matrices U and V, we have
that
UAV
||-
||
Proof. See Lemma 2.5 in Hoheisel and Paquette [2023].
Lemma (Existence a restricted form solution to (Woy-SVM Suppose W is the solution to the optimization problem (Wev-SVM) There exists a solution with pt, P2, Qu Qz ft. f2. gt 92. Bi. 32 and Th "12 such that
Mysin blocks
Must blocks
Pi In + P2E Pı + E. In En m I, + 92 E,,
III, + f.E. I E2E,
3,1n 3,1n
W
(11)
91 In + En I, + 92 E,
1 1 m
blocks
blocks
Moreover,
Ph.hh Pi + + q1 + + + +1,
(12)
Q: + + +
f (g) 00) + 9 +1 +
16

---

Proof of Lemma 3. We first show that certain permutations of W are still solutions of the opti- mization problem. Suppose that any permutation of {L.....m) n.. corresponding permutation matrix. Consider the permuted weight matrix
Define Pa Rax* as the
(W Po
W Wovdiag Po. 1.I Poj
It is equivalent to permuting the subject sets and the fact labels b and simultaneously with (Sa())).. Smmn), (br()):..,,b(x)) ba(x)) and (cr(I)....,G())))) Co(n)) Using Equation (6), we have that
()scc)-))..boo)) (a(Wo)) b(s.r1).by (Wov) 2 L,
Vj E [n] \ {i}.
(1 (a(W)v) Monriii (Wov) 2 1,
Vj E [n].
bi (a
(W h( (Wov) N 1, Vj E [n]]
(a) (a( = (a (Wov) 11, Pj E [n] \ {i},
for any í. From Lemma 2, since permutation matrices are orthonormal, [ II = o
Therefore, o (V) is also a solution to the optimization problem (W SVV-), Let's consider the average over all possible permutations
Σ. (Wov) n!
blocks
blocks
P1 In + P21 B.
Pi minin In Parmrai
En
fu fen E.,
fimmes Ln farmes En 811n B21n]
Qn In + 921 E. Qimernn In Qamaaan
girdn + 921 E.. ... Prmees Iu + gemees En I1 ?21n
blocks
blocks
is also a solution to the optimization problem (Wov due to the convexity of the nuclear norm.
We can consider other permutations. Suppose that Ttrain is a permutation of the index set {1. M.train Define the permuted weight matrix
Ttrain (Wov) W (P Prrsi @ In
It is equivalent to permute all the subjects in the set Sutrain (min 8t. (m } for any i 1,.... n. Note that there no need to permute the labels, as subjects in St,tain share the same label pair b, and o;. Since the permuted S; traa is still disjoint with the test subject set Sttest we have that for any E [n]. E [Mrrain], it holds that
( (Ttrain (W )) = h(ss.srr)a (Wtv) 2 1, Va' A((b))
(ess (Ttrain (ov)) = h((r.))72) Wov) 2 1, Va' E AV(c).
From Lemma 2, we can conclude that |train SorainnWWvillll (Wvv)l++ = Wovll Averaging over all permutations and Trmpan? we have that
Errim Ea Ttrain (a (WE (V
!nt
Mhainn blocks
Meear blocks
Pi In + E2E. PI + En In En In + En
fu + fen E. fimma In femm En
311n
321n] 12
In 921 En I1mma I 92mase En
n1,n
blocks
Mars blocks
We can then consider the permutation over the remaining miet indices {maaiin + 1, m.... Let Trest denote the permutation. Consider the permuted weight matrix.
Trest (W (v) w divdig (1, I,...,,, Prame In,1,13
It is equivalent to permute all subjects in the set St Tme (mhon 1): ....... 81 (m)) for any
i = 1,..., n. We have that for any E [m]. E [m]] [merrane ]. holds that
Man ra) (Tres (Wov)) h(sssmm)). (Wov) 2 Va' E Abbbi)
17

---

Taking the average weight over all possible permutations Ttest
Eis Dimi Ea Trest (Terain (W (V)))
Mitest Morain .nd
blocks
blocks
Pr I + En PI I En
I + f.E. f I. + En I1 92 91 In + En
811n
821n 72
qn + En 91 I Bn
Maaii blocks
Miess blocks
This shows that Equation (11) is solution the optimization problem WW This proves Lemma 3.
We could compute the closed form of the SVD decomposition for Equation (11).
Lemma 4. The restricted form Equation (11) has a the SVD decomposition W UEV with
U (2) (2) u2 Un
(1)
(1)
(2)
U2
1
Σ diag
01 00
(1) (2) (1)
2
V = [1(1) (2) 22
(1) (1)
(2)
222
where u are defined in Equation (14); ui are defined in Equation (16), 01 (k) and 02 are
(K)
defined in Equation (17); (k)) are defined in Equation (18); o (1k) are defined in Equation (19).
The proof of Lemma 4. Denote the Singular Value Decomposition (SVD) of W by W = UEV
The dimensions W are 2n x No, where N = (Merain + Myeest n 2.
Orthonormal basis and matrix properties. Given orthonormal basis ( 7h for R" with 77 1,/ /n, where 1, is the n-dim all-one column vector. The vectors Tp...... . 7n are orthonormal to each other and to m. Let X ci In + 02 En be a block appearing in the matrix WF Its action on the basis vectors is:
ov
Xm = (n + c2En)( ln) = ai J + c2 (1,1,) = am + 02 oा l
(cr m02)m-
For j E (2, Xn, (n + CrEnn c, since E, 1,, O, due to nj 1 1n
Left singular vectors U and singular values E. The left singular vectors (columns of U) are cigenvectors of WFF WFF.. multiplication, WF W ov The matrix W W 2n x 2n symmetric matrix. After block ov Ov can be written as:
ov
W OV W (CALLa + C2 Cs + EB
CBILn + CB2 En CDrt, + CD2
where the coefficients are:
CA1 = MurainPi i + Mrest 2
C Mitrin (2p122 + mp2) + Mistt (2f1/2 + + + 3 +,3
Cp1 M Φ
2 + M1est.
2
Cp2 Mirain (2q Q2 na2) Mrest (291.22 ngz)) 7 +
Cø1 = MtrainP191 + m1sstfig11
CB2 Mprain(P1 92 + P241 + np222) + mrest (f1g + feg + nz222) + Bin 8272-
We seek eigenvectors of W ov WFF of the form
1b
u
for scalars Fee
18

---

Case 1: Eigenvectors associated with 7h1 For j 1, the eigenvalue problem W WTF transforms into 2 x 2 eigenvalue problem for the coefficient vector (Io, Te)
ov u Au
H1
2b Ee
= A
Itc
where
CAL nCa2 CB1 + nC
CB1 nCB2
Hi
(13)
CD1 nCD2
Let 11 xx3) be the eigenvalues of H1, and let (01,b, F1,c) T and (202,6, 22,c) be the corresponding
/a(²) for 1, The
normalized cigenvectors. These give two singular values: 01 corresponding left singular vectors in U are
( for k = 1,2. (k
u (k)
(14)
Case 2: Eigenvectors associated with 7 for j E n.. For E (2....,n) the eigenvalue problem Wo Wov u = Xu reduces to:
H
= x
Ib
where
CA1l
CB1
H2
(15)
CBL
CDI
Let ( be the cigenvalues of H2, and let (31,0, V1,c)" and (22,b, b, V2.c) be the corresponding
normalized eigenvectors. These give 2(n - 1) singular values: (k))
V Al
(k") for k' = 1, 2. Each of
these singular values has multiplicity of (n 1). The corresponding left singular vectors in U are
yk b for each j (2...., n) and K' 1,2. (W
(16)
The matrix U a 2n orthogonal matrix whose columns are the left singular vectors u (k) (k') and u
The singular matrix.
From the calculation of U, we get that
Σ = diag 1 of .2
1, (1)
(2)
(17)
n-1
where o1 Vx ( for 1,2, with (()) being the eigenvalues of the matrix H, in Equation (13);
(W') = / for k' = 1,2, with 2 being the eigenvalues of the matrix H2 in Equation (15).
(kk)
Right singular vectors V. The columns of V (right singular vectors) are obtained from v; of WT ov u; for non-zero o, If a; = 0, v; is normalized vector in the null space of W The matrix V has dimensions N, x 2n. The 2n vectors v; corresponding the found singular values
Case 1: Right singular vectors associated with (k), k = 1. For (k) =
( Tk
and
singular value of (k): The vector WFF u (b) has the following structure: ov
The first muain blocks (each of size n) are (0 np2) + zv.olgi + ng2))
The next mest blocks (each of size n) are ( + n2f2) + 24.0 m92))m
19

---

The last two components are scalar values: vn(512kk + +1 and (2 + k2
So, (())
W u (0/ (k) is:
(F np22) IKke(91 + (292))) )m
:
(In,(Pp) + np2)) + [2to( + n1q2) ym (Kn +nf2)
(A)
1
xk,o ng22) m
for k 1,2.
(18)
(k)
01
(#Ralhi + nnf2) + (g1 + mg22)m / FR.b b 71 Z2 + 72
Each 771 term represents column vector of n elements. This vector ((() has total length No Ti + Whest 2.
(k²) For u Case 2: Right singular vectors associated with uj
(K))
(where j 2 2 and Uk
k' = 1, 2) and singular value o² (k') The vector W (M) has the following structure:
The first marinn blocks (each of size n) are (yw bP + (1)
The next mest blocks (each of size n) are (ytr ff1 + yk -0/ ..g))/.
The last two scalar components are O, because 1, n, for 2.
(4) So, 0j // Wow u/ /0
(K)
(UlkrbP1 ee : (uwrop1 + yk eq1 )
(K)
1
(Uw + ,w
(19)
02
( + )w 0
0
There are 2(n 1) 2j vectors, for each pair (j,k') where j .2, nn and k E {1,2).
Each ?i term represents a column vector of n elements. The vectors (K)) and vy (K) are orthonormal and form the columns of V.
As conclusion, the SVD of W is W U.VT where:
U is a 2n x 2n orthogonal matrix with columns (A) and (k') as defined in Equations (14) and (16).
Σ 2n 2n rectangular diagonal matrix, whose non-zero entries are the singular values σ and d2 (K)) derived from the eigenvalues of H and H2.
V is N. 2n matrix with columns (k) and v/ defined in Equations (18) and (19).
This finishes the proof of Lemma 4.
Lemma 5. The W in restricted form Equation (11) has the nuclear norm [ ll:
20

---

CCai + + (Cai +C (Coi + nCD2))² 4(CBI + n=B2)=
(CCr + +nCa2 (CD1 +nDD2)) 4(CB: nCB2)²
+
2
(v
(CAI+ + V(CAI - Cpi)² + +4
+Cal+CD)) V(CAi CD1)² + 40
+ (n-1)
(20)
Proof of Lemma 5. The nuclear norm of a matrix W denoted as O is defined as the sum of its singular values. From Lemma and its proof, the singular values of W ov are derived from the eigenvalues of two 2 x matrices, H1 and H3.
The singular values are:
(1) and 0j which are the square roots of the two eigenvalues of H, defined in Equa- tion (13). These singular values each have multiplicity of 1.
2. d2 and 02 which are the square roots of the two eigenvalues of H₂ defined in Equa- tion (15). As stated the proof of Lemma (Case 2: Eigenvectors associated with π7 for E {2...., n.)), these singular values correspond to the n basis vectors n, for (1) and 02 has a multiplicity of (n - 1).
j m2..... Thus, each of 02
The nuclear norm is therefore the sum of all these singular values:
[| Wovll | ( + di (2) (n 1)02 (1) (n 1)02
= =o + +01 + (n 1)(0)
(1) +0² !
Let's expand each term:
Term 1: Sum of singular values from H1.
The matrix Hi is given by Equation (13):
CAi nC CB + nC
Cm nCnz
Hi
Cm nC
Let a = CAI + nCA2 bi = CBI + nCB2, and ci CDi + nCp2. The eigenvalues ( x(2) of this symmetric 2 x matrix H1
bi
are given by the formula
b
(a1 e1)
v
(a) - 22 4b²
2
The singular values o and of (2) are V A11 and
x
We get that
+
(a + c1) + fa1 )² 4b
(a1 +c) V (0: C1 46
2
2
(Csi + nCAz Con + nCb2) + ((CA + +nCsu (CDi + nCbul)² 4(Ca1 "C22)²
(Cai n + Cm nCm) V(CA + nC (Cr) nC 4(Cm )=
2
This corresponds to the first two lines of the expression for the nuclear norm in Lemma 5.
21

---

Term 2: Sum of singular values from H2.
The matrix H2 is given by Equation (15):
CAi CBi
CB1
H₂
CD1
Let 02
CA1, by
CBL, and C2 CD1 The eigenvalues A2 (2) (2) X2
of this symmetric 2 x matrix
H2
a2
b
02
b2
are:
(02 c2) v (a2 02) )2 4b
2
The singular values 02
(1)
and o are v 1/2
(I) and V X2
(2) Thus:
(2)
+ 0
(22 + (2) + (02 62)2 + 4
(02 + 02)
v (a2 02)
(²) 4b2
2
2
(CAI +DI)+CDI)+ + V(Ci - CDI)² + 10
(Cai C - V(Cii +
2
2
This sum is then multiplied by the multiplicity (n 1):
(CAi +CDi) + V(CrI - Cm1)² + 40
(n - 1) + +o22) (n 1))
2
(Cn + C -
- -Cn 40
2
This corresponds to the third and fourth lines the expression for the nuclear norm in Lemma 5.
Combining these terms, the nuclear norm fy is precisely the expression given in Lemma 5. This completes the proof.
Lemma 6. Let the expression [[Wov II+ be defined Equation (20). The closed-form minimum of | lI is given by
min
vn (n v
x
{( (( (Merrann + 2 Mrest
Mrestt
if Mrestt 2 Mltaai
if Mrest < Miraain
where the minimum is achieved at pt fi qi =1, si = Merrain /rrest p2
13 = 43 1/nn,
g = Merain Mreese n,, 31 = 7/ = 1/2, r.
-8- -12 if Meest 2 Miraini pi fi Φi =
9i 1,p
92 = - 1/ = = 1/2, rt = 55
2 if Muest Merain-
Proof of Lemma 6. The overall expression | given in Equation (20) is a sum of two main components. Let MM be the sum of the terms in the first two lines. and M2 be the sum of the terms in the last line. Thus, Wo Mi + M2. Both Mj and M, represent sums of square roots of eigenvalues of effective 2 matrices.
Part 1. We first consider minimizing the last two terms of W ll.. Define M2 as
(Cn |Cn + CD1) + ((Cr CD1)² 40
(n-1)
(Ca1 + CD.) (Cai CDI)² 4C
2
2
(21)
For case of reference, we state the formulas for coefficients CAL. Cp1, CB1:
CΛ Mirain mesttfi
(22)
CDI m v Mrest 2
(23)
CB1 MtrainP19 Mrest f191
(24)
22

---

The expression for M2 can be simplified. Let 2 and 11²) be the two eigenvalues of the matrix H2-
We note that 2 + 2 CA11 + CDl and X2 xx² CAI c.. The term CA1CD1 - C can be calculated
CAICDI - B = (MtrainPî + Mrest m (marain + mesasl) (marainn1 91 + /19))
= thes (pig) - 211 + f29)) = Mest (P191 - fiq1)²
Since Merrin; Mitett O, we have Cai - - 2 0, so the eigenvalues λ1, )2 are non-negative.
(V 1/
2
V 12 x + 2 2
The sum
can be written as Substituting the trace and determinant yields
VA ()) /> VCCII + +DI + 22 Marainneestt (P1g1 }191)2
= -CCAI- + CDI + Mestt PIgl -i
Let S = CAI + CDI +
/ Mirain P1g1 - fi Substituting the definitions CAiI and Cpi
S = Merainp? + MrestA + mtraingi + Mressi + Mes - fi
(25)
Then M₂ (n 1)V5. To minimize || we must minimize S with respect to pl, 91, fi, under the given constraints. Let K = P1g1 - fiq1. We analyze the minimization by considering two cases for K.
Case 1: K S 0 (I.e., P191 S fiq1). Then
S= S2 (VMMrrain Mires!!! ) + (ymrestSi M )
The condition K < implies that 91 S fiq//1. The first term (v MerainP1 v M1ett91 )2 is monotonically decreasing when 0 S 9i S P1 V Mtrain/ W Meest with the minimum attained at yl
pi v Marain/ Mhst.
If Moest > Merains 91 = P1 Merain/ Mres is always achievable, so by taking f1 = α = 1, we get S2 = (Vllrarii + / Mesa)) The equality holds for any L1 E [1, (Mress MMrrain 1/11 and 91 PIV Merain Mrest.
If Moss S Mtrain PIv Mirain/ Mrest may not be achievable. But since the first term is monotonically decreasing with respect to 91 when 0 S 91 S P1 / Marain/ Mteets get that the minimum is always taken with 9i min (f /iqq///p /I, p Pi / Mrest)-
If gi = fiq//pI. we get that
S² (VmtrainP1 frrest /////1)² + (1 (Vrrest.i f1 + Marain Q1)²
mf + Mrest + merainPi + Meett
> 2MMrrain + Mrest):
where the equality holds if and only if = 91 = f
If 91 Pi piymmer M Mheets we get fig pî /mrrain///// /rrett 2 / "mrrai/meeet Therefore,
S2 (/mressft f1 + /1
2 4 Merain fi41
2 4mbrain
where the equality holds if and only if p1 = 91 = V Merain/ M
Therefore, we can conclude that if Meest < Merain: S2 2 2( + meest), with the equality
holds if and only if p1 91 fi qi 1.
23

---

Case 2: Next consider K 2 O, that is, P191 2 f191. Then
S= S1 MtraanP1 + Mrees)) )2 + (ymrest/i v Merain
Taking 91 2 fiqI/PI into St.
R pf
S 2 Mieet f + Merain MerainPi 2 + Myeett
2
Moest + Mitrain MerainPi +
Mhest pi
If Meest 2 merain we have that S1 2 (v Myteet + Merain) with equality holds when
= [Mest/ Pites/Ment Merain] and 91 1/p1.
If Mist < mtriin the minimum is achieved by = We therefore get that = 1, with 2 + M.
Consolidating the minima, we find that if Mitstt 2 Mrrain the minimum (v Mitrain + Threett 2
achieved when 1 I P1 < (Mrest Marai)) 1/4, fi = 1. = 1, and = Marain/ MMesst- If Mtest < Mitriin the minimum S 2(Marrin + Mtest), achieved when p1 1, fi 1.01 1. 91
The minimum value for v therefore /illtrain + //tttttt if Witett 2 Mitrrin and v ( + Mitest if mist < Mltaiin Multiplying by (n 1) gives the final result for min M2-
Part 2. We seek to find the minimum value of the quantity M1 given by
(Cai + #CC2 Cm + nCm) (CAI + nCa2 (Cm nCm))) 4(Cm) nC#2)²
M1
2
(CAi m + CDi nCD2) (CA1 + nCC (CD1 + nC + 4(CB: nCB2)²
+
2
where the coefficients are defined as
CAI Mitrinni 1 + Mrest 12
CA2 = Mirain (2p1P2 + np2)) + must (2fif2 + nnf)) + 87 +g
Cp1 Mitain 2 + mest119
Cp2 Mirain (2q) + n4ž) + Mrest (29192 + ng + ?ィ 7/
Cв1 M4rainP191 Q1 + f19s
Cg2 mtrain (P + P2 + nP2222 + mzest (f192 + f20 + nf292) + fim + 8272.
subject the constraints
PI. P1, 2 1,
(26)
Pi + P2 + 31 2 41 + 71 +
(27)
41 22 +2 21 + 32 1,
(28)
f1 + 12 max(g1,
71 +1.
(29)
Let A = CAL + nC D = Cpi + nC'22, and B = B - nCB2. The expression for M1 can be simplified
Mi 4+ B2.
Note that
Mf 2 A+D
Mtrain (pi + mp2) 2- + Marain (g) n22)² + Mestt (f1 + r2 2 + meett (9) + ngg))
+ n(33 82 + + 7
Let A Q1 + 92 P2. From the constraints, we can get that
- + A,
'2 - 1- A.
24

---

If A > 1, without loss of generality, assume A 1. This gives that
+
(5f + + (32 + + 2
+02
2.
2
If A S get that
1+ 2
1
2
(0i +1) (82 +22)
2 2
+2
1
4
A2
11
As result, we get Mj 2 vτ. Note that the equality holds and only
P2 = - pn/n, f2 = ff/nn, -411/1 92 = 91/ n,
B1 1/2, 71 -1/2, B2 1/22, 72 = 1/2.
Combining this condition with the results in part 1, We finish the proof for Lemma 6.
L
We would adopt Assumption 3 to prove the uniqueness, which assumes the solution either takes symmetric form as (11), or it is nondegenerate.
Assumption Let Wo be a solution to WWov-SVM SVM). Then it either takes the form of (11), or we have 2n W * O 0mm+2-
Lemma 7 can give lower bound the nuclear norm for given matrix M.
Lemma 7. Suppose that M E RR>xn for any orthonormal matrix U E Rmxr and V E R we have that
| + 2 Trace(U UT MV). Moreover, if there exits E R s.t. n T = and n M + 0, the above inequality is strict, i.e.,
M Trace(U" (VM))
Proof of Lemma 7. For any pair of matrices A and such that AB = M, we have that
Trace(U MV) Trace(U ABV)
(IIU* A
BV 12
[IU ||
V
BI 3112)/2
[| |
|
/2
M
Consider the conditions for taking the equality. We should require | = A
This means that all columns of A are left-singular vectors of U with the largest singular value. In addition, it implies that the column space Col(A) Col(U). But if there exits n E R s.t. U 0 and n M + 0, get Col(A) & Col(U). It implies the conditions for taking the equality cannot satisfied, so
M |. Trace(U (VV).
This finishes Lemma 7.
Theorem 5 (Uniqueness of the solution to the optimization problem (WF (VV--VM), Given Assump- tion 3, suppose that W is a solution to optimization problem (WF (WB -VVM),, if Mrest < Merain.
Meriin blocks
Wiess blocks
In Em n I, En/n
In Em / I En/n 1n/2
11, -
WF
(30)
I En/n n. I En/n
En I I En/n
- 111/2
1 1
Mesin blocks
Meat blocks
25

---

If Miestt 2 Mirriin let Marain/ Moestt we have that
Mosie blocks
Mheet blocks
In En/n n. En/n En/n
In End n In En/n
1n/2
-1,/2
In En/h
pl pBn/n pIn pEn/n
-1n/2
1,/2
Musin blocks
Must blocks
(31)
Proof of Theorem 5. Firstly, combining Lemmas 3 and 6, we get Equation (30) and (31). This shows that Equation (30) and (31) are one of the solutions the optimization problem Ow
Suppose that there exists asymmetric solution of the SVM problem Equation (W OV--SV W that takes a different form from (11). Lemma 8 proves the existence of another asymmetric solution such that the difference takes the reversed sign.
Lemma (The solution with the difference taking the reverse sign). Suppose that the W defined above is asymmetric solution. Then there exists another asymmetric solution Wov and e (0,1)
such that
Wov - W = -e (Wov - Wov)
Proof. Let e be all possible permutations Wov as described in Lemma excluding the identity permutation. We get that
Σx 6 (W)) Ov)
W ov.
1+ E,1
So take
Wov = Eg @ (WW)
E31 1
and get that
W WF. Eg 1
Wov Wor
This finishes the proof for Lemma 8.
Assumption 3 implies that In (WTy Wov) * O Am442 We choose U E 2 and V I (2)m+)2)xr that they are the SVD decompositions corresponding positive singular values for
the symmetric solution Wiy. Since 1n W 0 0 it also holds that 1 U 0 2m42. Using Lemma 7, we get
1l > Traco(U W)
Trace(U " V) Traco(U (Wor V)
W + Trace(U (Wtrr V).
Using Lemma 8, there exists another asymmetric solution Wov such that
Wovll > Trace(U w V) + TraceeU"" (Wov - w) |woll. "TraceeU (WW W)
Therefore, have that
max II
W + max(Trace((()) (Wot- w -eTrace(U (Wot W)
2 < Il..
This leads to a contradiction that both F and Wov are solutions to Equation (WPy-VV that minimize the nuclear norm. This confirms that all solutions would take the form Equation (11), which would be Equations (30) and (31).
26

---

A.3.2 Proof for non-factorized model
Similar to Appendix 1.3.1, we only need to consider reduced matrix Wov R(2n)x(mm+2)
throughout this section, where each row corresponds to token in and each column corresponds to token in S U U. Now We restate the second part of Theorem 2 below.
Theorem 6 (Part 2 in Theorem 2: Non-factorized model has OCR ability). Let n. > 1. Suppose Wov is a solution to the SVM problem in (Wov-SVM). For any (8,r) E 'Dest and any o' E A2 \ {a* (s,r)), holds that
(( (Wov) 0, indicating no OCR ability.
(32)
The proof of Theorem 6 follows the same idea as Theorem 4. The roadmap of the proof as follows:
A unique solution of a similar form to Lemma 3: Lemma shows both the existence and uniqueness of the solution to (Wov-SMM).
2. Frobenius norm formula of the restricted form: Using the same SVD decomposition as in Lemma 4, Lemma 10 gives Wov |le in closed form (35).
3. Optimization: Lemma II finds the minimum (35).
4. Solution characterization: Theorem 7 gives the form of the solution to (Wov-SVM) and finishes the proof for Theorem 6.
Lemma' 9 (Existence and uniqueness of restricted form solution to (Wov-SVM)). Suppose Wov the solution to the optimization problem (Wov-SVM), The solution must take the form in Equa- tion (33)
Mitrain blocks
Miet blocks
PiLn En Di + En Iihn + En ... 91 In + 92 En
filn + f2E fi En
311n
321n] 72
Wov
(33)
91 In + En 11 + 92 En
n1n
Mursin blocks
blocks
with Pt, p2, 91 Q2. S1, f2. and gr, 92 satisfying
PI, fn,n 2 1,
91 42 + + P1 P2 + 31
(34)
Q1 + 02 + 72 2 P1 + P + +1,
+ 31 (g) V0) + 92 71 +
Proof of Lemma 9.
The existence proof the same as the proof for Lemma 3. Since the Frobenius norm is strongly convex, the solution is unique.
Lemma 10. The Wov in restricted form Equation (33) has the Frobenius norm.
Mtrain (P) + P2 p2 (n 11= + (qi + q2) 22 (n - 1)
1/2
WovllF
+ +rrest (f11 + f2) + (n 1-f2 + (01 + 92 + (n (n-1)93)
(35)
+ + 12 + 83 +
Proof of Lemma 10. The square of the Frobenius norm is the sum of the squares of these 2n singular values:
w = (a (2)))2 + (n () + (n (2)
From Lemma 4, we know that (o (1))22 = 11 (eigenvalues of Hi) and (az (k'))² x (eigenvalues of H2). So.
Wovll X + +n 1)(1 (1) + +x²²)
Since the sum of eigenvalues of matrix is its trace, have 2 + x Tr(Hi), X x Tr(H2). Therefore,
|| Ib (n (n-(())
=n(CAi + CA2 + CD1 + Cp2)
merain (p1 + p2)² (n 1)p2 + (g1 + + (n-1))3)
+ ( (fi + f2) + (n 1),f + (g1 92)² + (n 1)93)
+ + 7/ + 82 +
This proves Lemma 10.
27

---

Lemma 11. Let the expression Wo be defined as Equation (35). The closed-form minimum of Wovlle is given by
min Wovlle lE = vn (2marain - 1/n) + Mmestt (1 1/n) + +1)12
where the minimum is achieved at pi =f qi = 1, f2 02
-1 11/n, 31 72
7/ = - and of = 92 0.
Proof of Lemma 11. Given the constraints in Equation (34), we have that
Wovlle
[m]
n
(ttrain(1pps 2p2 + na² + 20p +
+mmeet (n + 2f + 1 (g1 + 92)² + (n (n-1)9ī) + +1) 1)
1/2
where the equality holds when p ft qी = and 3 7 = 1//2,71 = .B2 -1/2. This
lower bound quadratic form, which we can show that p f2 0 -1/n and gi 92 gives its minimum. This finishes the proof of Lemma 11.
Theorem 7. Therefore, we have that
Mavinn blocks
Mies blocks
In Em/ n In En/nn Wov = I En n I Em/nn
In
En n In En/n
1n/2
-1,/2
(36)
0
-1a/2
1n/2
Meriin blocks
Mexst blocks
implies that for any 8 E Stest and any d' E Ao (8,)2)) holds that h(s.rz),a' 0. This proves Theorem 6.
Proof of Theorem 7. We take the results from Lemma 11 into the restricted form Equation (33) and get Equation (36). Given any $t. E Stest and Cj E A2 where k & i, we have
Bi
Mon =10,,,e; - eg Wov
0
0.
This proves Theorem 7,
B Proof of Section 4.2
In this section, we provide the complete proof of Theorem 3 via gradient flow analysis. We first present several key lemmas to prove Theorem In Lemma 12, we give the gradient form of the reparameterized parameters to simplify the analysis. In Lemma 13, we prove the Lipschitzness of the gradient flow, and Lemma 14 proves that the gradient flow also satisfies permutation equivariance. Using these two lemmas, we are able to conclude the symmetry in parameters by employing standard argument for the uniqueness of ODE solutions in Lemma 15.
Useful notations. We introduce useful notations for this section. Let vec( denote the vec- torization of matrix A. Specifically, for A (aig) E 1mxx, we have vec(A)
(011,.... Oin, a21..... a2n. am1. ......... 0mn) T E Rm We use || to denote the operator norm of matrix A. For any vector u E IR." let diag(u) E gux= denote a diagonal matrix whose diagonal entries are corresponding entries in u. Let 1(8)) denote the indicator function of an event E.
Lemma 12. Recall that for any 2 E V, we have Wovla, 2) = e.j Wove, and WkQ(2) = e WKQe<E0s>> Then for a fixed (a, 8) E Ax x S, the gradient of Wov(a, 8) is given by:
Owavla,) WK)(s) p(s,)) (1(a a" (s,r)) po(als,r))
FER
28

---

Similarly, for a fixed (a.r) E Ax x R, we have:
) ((0) -Wko(r) Σp(s,r) (1(a (1(a= a((,r)) - po((as,r))
SES
Moreover.for for <EOS> token and a fixed A, we have:
(soo), <B0s> [(0) - WKo <EDS>)
ΣΣ p(s,r) (1(a as,r))) -Pa(als,r))
SES
Lastly, for s E S, we have:
Owra(s) --
Σ
Wov(a,
Ep(o,r) (1(a = a -po()))))))))
aEA
rER
Proof. First, note that the logit function can be written as:
fo(=1:r = e.a Wov X WKKQTT=
(37)
KE[T]
Then for any V, we get
@oo((, (e WovX XWKO = 1aa= 2
(38)
Owills) (c Wov XWKoεT) 2Wov
where C(aLTT 2) is the number of occurrences of : the sequence 21:7. Recall that we can simplify the loss function in (4)
exp(e Wov X XWKQæT)
((0) E
log
Exedexpled Wov X XWKQTT)
= B l-e 21 WovX XWKOTT + log Σ
exple, XkkoT))).
:'EA
Using (38), get
ÖWov(a.s) [(0) W2 1(a ((
Exeexxp(ee Wov X (0=
Exeeexple, Wov XT XWKQœT)
WKo(2)EE... [(ELT, ((E)rr))(1(a (1la
a (zrr) po(al217)))
Similarly,
@wiale) C(0) = EETT
+ Σ (ola
aEA
E
C(arr, 2)
Wov(a" (-oolaa (xir)) + Σ Pa (a (a
Wov(a, (a, E
(1 a" (21r)) = (a (
Let p(s) = E p(=T)) E .).)) = Ep,.. [1(s E æ1r)),, which is the marginal probability of
observing s under the cata-generating distribution Drain Similarly, we can define p(r) and p(s,r)
29

---

for any r E R and any (s,r) E S x R. Then for fixed (a, E A x S, we have
aWor(a,s) (())
Wka (8) B Tr
() Sacp)))pPpppp00 pa(a)21:))) [1( (1a=
(a)
(Ko(s)) (V) (8)))))) Dowin
- -0 (a|zz))))
- (()
- WKQ(s)
Σ p(zir))(( (1 (z(r)) - -ro(alz1m))
-- --Kka(s)p(s)
En p 1 <
21.7
WkQ(s) p(()) Bayr I(a a (ETT) po(alz1.T) S =
Wka(s) Σ p(s,r) Baur 1(a a (ax.)) p5(a)21:)) 1 E B
(b) Wko(s)
> p(s,r) (1(a = a' (s(s)) po(als,r)) rER
where (a) comes from the fact that C(ai:, 8) (s 21.T) as s oceurs most once in the sequence from the task definition and (b) comes from the fact that 21.r only contains (s, r. <EOS>) and thus ((1T)) = Pa(sss))). r).
Similarly, consider a fixed (a,r) E A x R. we have:
@wor(a.r) -Wko(r) p(r) Exar 1(a = ((1r)) po(a(21T) |s 21.T]
= -Wka(r) Σ p(s,r)
(1(a = (s. -o(olsr))
SES
Lastly for any a E A, we have:
(En ((0) WKo (<DS>>
Σ p(ear) a" (2(1r) -(g1a))))p((1))))))))
p5(a|21:1))
Wika WWka(<EOS> (<EOS>)
p(s,T) (L(a= (3,))) po(als,r)).
SESTER
We conclude by proving the gradient in terms of WKQ(s) for any E S. We have:
(w) C(0)= = Σ
[(a a" (#1T)) po(a)Zir) s E Wov(a, 8) p(s) Eear
aEA
Σ
Wov(a, 8) Eplorr) (I(a a" (s) pg(als,r))
aeA
FER
Lemma 13 (Lipschitz gradient). Let do = (IA) + 1) |V] and concatenate the parameters by
(vec
RR.
WkQ
Let F(w(t)) be the vector field in the ODE
wtt) F (w((t) = =-V -
where omit the index and set w(l) when the context is clear and L(w) is the cross entropy loss given by:
exp ((), a
L(w) Ep(srr)
-oog
EnEA exp Sw (())
30

---

where the logit function fw fs,(,, a) follows Eq. (37) with 21.7 = (s, <EOS>). Assuming that
||wl22 S R for some constant R > O, thenfor any fixed (s, (,rra), a), fw ((s,r), a) is Lipschitz in w with
constant L1 == 4R, i.e, for any w1, w2 E 3R with 1| 2|2 S R, 1|22l2 S R. it holds that
Jw ((s,)), Juz ((s, a) & Tu||u] - w2ll-
Moreover, there exists a constant = 4|A1( 11 1) such that
[F(w)) - F(w2)ll - 1||w1 - w2ll.
Proof. Observing that foo ((8, r), is bilinear in w, i.e.,
Jw (() = Σ
Wov(a, 2 (x)(())
tE[T]
(a, (sk + (ok(a)))WkQ(t) Wov(a, <EOS>) (<EDS>)
(Ko(<sS>>>
= w
a)w.
where M((s, a) e do has only three non-zero entries, which are equal one, where the positions depend on ((s, r).). We omit the index ((s, r), a) when the context is clear. Then for any W1, W2 E TRdo with Euclidean norm bounded by R, we have
fw, ((s,r), wwz ((s, r))
w Mw W22 Mw2
w Mw - Mw2 + wi Mw wz Mw
(||w||] + ||w2||) |M | |w]1 w2ll
4R|Ww] wall L1 || -wall
(39)
where (a) uses || F v3 2. Moreover, the gradient of f ((8, r), also Lipschitz:
Vee ((,,)),)) V ((s, (M + M M (w) -22) S 4|/20 wall-
(40)
Now let pk (w) ΣK, exp(uy) and gk (u) log pa (u) for R where K |A], and denote p(u)= = (x..).
exp(ve)
px(u)) (w)) E E g(u) (u)(u), 9K (x() E R Then the gradient and
hessian of gu (u) are given by:
g = pu)) es E K
H(u) diag(p(u)) p()) e RR
By the mean-value theorem, have:
[[Vgx Vga(o))l S sup
|H[Ww\]opl|u w/l S I/u v/|.
(41)
WERK
Let (8, = Fu ((8, E [A] be the logit vector for input (s, and denote
aa (uw (8, r)) = (ar) (n (s, r))]a (u r)) ( = a (s,
Then the gradient of the loss function can be written
VL(w)= Σp(s, p(s, (.err wha (a.r) (s, = Σ p(s,r) Σ
aEA
31

---

For brevity let V fwwa = ((, a) and note that Vfwall
(I + M )wl 1 4R.
Combining Eq. (39) and (41) have:
[[((1]) - F (ww2) II
= =[V[(( v[(w))ll
2 Σ p(s, r)
Σ ad (utw (s,r)) Vwona aedd
Σ
(u V
3,7
aEA
Σ ((aa(uurr (s,r)) au (u (s, (s, V + an au (s,r)) (V - Viwzaa
max
lEEA
max
V fwr 1/lbr (8, uuus r)) + 4/AA/ max laia (2tws (s. r3)l |1/01 w22
8,"
AeEAA
w1
w2ll
L1
Vfwi
4|A]
aEA
<4|A|(L, 1)
|0 w2ll
where (a) follows Eq. (40) and Eq. (41) and (b) uses Eq. (39) and (Ka(u, (s,r S 1 for any E A.
Before proceeding, we provide the following definition.
Definition 1 (Data permutation), Let do (I + 1) (1). Consider the flattened parameter w defined as
[vec (ov)) WKQ
TR
= Let a be any permutation over [n] where o(1) * for any [n] and V be a permulation function determined by a. Specifically, for any [m].] [m]. have
Recall [A,l = -
T(b:) bo(t), (c)) Cots), n (81s) 80()))
and we have T(v)= for v E f2,, <EOS>}. Moreover, let Pr. be permutation matrix built on permutation n defined as follows. First, we have
P Pov
0 PkQ
E ERdooooo
where we omit the subscript πt for brevity. We denote the resulting OV block in w after permutation as (Pw)ov Povvec(We) (ov)) MHAA Similarly we denote (Pw)kQ PkQ WkQ M Then, each entry of Pw is defined as follows:
8 E Stest:
Va E A2:
(Pw)ov(a, 8) = Wov(T(a),
V s E Strain,
VaEA:
(Pw)ov(a,s) Wov(T(a), (())),
Vu E (21,72, <EOS>}, Va A2: (Pw)ov(a, u) Wov(=(a).o).
V E Strain (Pw)kQ(s) Ws)(
Otherwise, (Pw)ov(a.x) Wov(a, v), (Pw)ka(T) WkQ(v).
Lemma 14 (Gradient is permutation equivariant). Recall the flattened parameter wo and the permuta-
tion matrix P defined in Definition 1. Denote F(w) -[((w) following Lemma 13, then we have
Vs Stest,
VaE A2:
[F(W))ovv(a s) (F(w))ov(a,ss)
(42a)
Vs E Syain:
VaEA:
(F(Pw) (a, = (PF(w))ov(a,s)
(42b)
VUER (EERUSEEESS)) <EOS>),
VaE
((PPw)ov(aa [PF(w))ov(aat)))
(42c)
V E Strain
[F(PW))KKQ = [PF(w)]kQ(>>
(42d)
which implies the gradeint of w.n. up permutation equivariant, iee.
v V((Pw) = L(w)).
32

---

Proof. We examine Eq. (42a) (42d) one by one.
Part I: Proof of Eq. (42a). Given fixed Stest, a A2, will prove
[F) (PF(uu)Jov(aa
(43)
Using Lemma 12, we have:
[F(w)]ov(a, s) -Woola,,)) Wka(s)
(a= (w(
rR
Wko(s)
= = Pwo (a|ss r1)].
where (a) comes from the fact that (2), since E Stest- Then for the right side of Eq. (43),
we have
PF(ww]oovaass) s) (0) Owow (x ((w)
(44)
=KkQ(s) p(s,
= a P(d
where (a) holds by Definition Similarly, for the left-hand side, we have
[F(W))ovv(a s)
(45)
(a) WkQ(s) r(s, (s, -P
where (a) comes from (s)) (*)(s) for any Stest- For the prediction probability, recall from (3), we have
exp(Suo ((8, a)) The logit function fw ((s, a) can be written as (w(((, a) = Therefore, given s E Stest and a E A2, we have
(als,))) E Ex'ed exp (f ((,,
Wka(r) + (<a(<<0s>>WWovee <E0S>). <EOS>)))>)
fp. = + WWKo(EED>>)oovrTH <EOS>)
=fu((s,r),TT())
Similarly, for any d' A1, given Stest, have
((s r), = fm f(u
Then
(sp(fpw((s,T))
Ev'e exp(]pm (1)
(xp(fu
Eves, exp(fu +2 Eoees (x)(fu((s,,) (a)
(46)
cxp(fu((s.)),
Pm (Pn(TT(()) (a))s,,))
where in (a) we re-index the summand by 2 == π T(a') since я is a bijection for any e Ag. Then Eq. (45) can be written
[F(Pu))ov(aas) s) Wko(s) p(s, r1) ) a (s, - Peo (p((())))))) By comparing Eq. (44) and Eq. (47), it remains to prove
(47)
I(a a" (s,ri)) ((7(a) a" (s,ri)),
which both equal 0 since a, (a) E A2. Then we conclude that for any s E Stest; a Ag. [(P))oov(a, s) [P)
33

---

Part 2: Proof of Eq. (42b), Consider any fixed E Strain, a E A, we start from the RHS of the equation:
[PF(w)]ov(aass) -Woo(rla),
Wko((T()) rER
Σ [p(T((),)) )) (
(s)(
Σ ((s()) = [(([((((((((c(((())))))) - Pp((())))))))))
FER
(48)
where (a) uses p(s,r) p(T((s),r) for any fixed E R as data is uniformly distributed. Similarly, on the LHS, we have
[F(Pw)lov(a,s) s) -) ((Pw)
= =Wko(T(s))
Σ (1 d'()))) -PPoo(a)s,)]
PER
To proceed, given s Strain and E A, we have
fPw ((5, r), a) = WKk( ooo (oa()) WKk (<E0s>)Wov <EOs>> (()
+
Following Eq. (46), when s E Strain, A, we have
exp(/Pu((ss))),
PPw (al(a)))
E (xeeecpp(fa(sss))...
(xp(f(((T(S)))))
(49)
ExEA (
Pw (T(() (s),))
where in (a) we re-index the summand by z = T(a). Thus we get
[(Pu))ov(a,s) s) WT)(T((s))
Σ p(s.r))(((a d(,r)) (pw(a)),
FRR
(50)
WWKQ(T(s))
Σ
FER
By comparing Eq. (48) and Eq. (50), remains to prove for any r E R,
1 (a a" (s,r)) ((T(a)) = a² ()()))))),,
which holds since the event a a" (s)r) is equivalent to ((a) a" (I(()),) r) by the definition of and the dataset construction. Then we can conclude that for any & E Strain a E A,
[F(Pw))ov(a, s) = [PF(ur))ov(aa.)))
Part 3: Proof of Eq. (42c). Let's first consider and E Ag, using Lemma 12, we have
{21, For any fixed r u E (ri.r2)
[(w))oov(a,)) = Wka(r)
a P
#ES
Then we have
P = WKko(r) p((,,)) = (T(()))ss))
#ES
-Wko(r) Σ p(8,r) (1(T(a) = a" (8)r)) - (mm
(51)
#ESmin
Wko(r) Σ p(s,)) (1(m(a) (m)
34

---

Similarly,
(s( r) -W Ems,r) (1(a)))(((a-))))1 (a,, r)) - ppw(((),r))
dES
Wko(r)) Σ p(s,r) (1(a= = (s, ()) - Pwo (r((
(52)
oESvain
+Wko(r)
Σ SES
((,,r) (1 = a (8, -M)())))))
where (a) follows Eq. (46). Now we compare (51) and (52). We first consider Strain. Note that for any 8 Strain and fixed r E R, the values of p(s, r) are the same due to the data distribution, and we denote the value Perain. To proceed, let i, i' be the corresponding index of a, π TTa) in A2, i.e.,
a
C; (a) Cg. Then
Pw - (1(m(()) a" a"(s)))) Pw (=(()ssss Σ r(s ((1(a- = a (s,r)) (a1aa=
ESman
-Ptrain
Σ
(I(a = a (sks = a (s)
JE se(muainl
[(a.7(0):80...1))
Σ
(1(a= (s) ((T(a) = a (sW'), r)))
kE(
[(a
+
Σ Pw
- (
kE[n]
(53)
For the first term I'(a, x()), sk.j r), note that (aa = a (sk.j, r)) = [(T(a))= = a (s r)) = 0 for
sk.))
any E [m]. h E [n], k * {i, (i, Now we consider the second term [(a, T (a), 8x' j.';'' r). We have
(a a" (sk r2)) = [(x(a) a' (sw...222)) 0, when K'
(a a² (sM r2)) =0, 1(=(a) a' (s =1, when K' i', (a- a" (sM = (=((a) a' (skr))))) 0. Then simplify Eq. (53) as - (1(+(a)) Σ p(s,T): ((1(a= a' Pm. (T(A))(()))))) ( ( = a" (s, - Pm (m((T())ssr))
when k' E (i,
which implies Eve(,,) I(a, ((o.
SEauiin
Σ kE[m)
Pw (T((a)ssk.)))))
Σ kE(m]
P (T(((()
PPrain
jE [eniin]
(a)
Σ kE|m]
Peo (T)((()))k)):)))
Σ k'E[n]
Pwo (T((((a))ssee)r11)
Ptrain
jE j
=0,
(54)
where in (a) we re-index the summand > kc(m) Pw (T(a))=(Sks))) by noting that (aakt))s))) So(k).
Next, we consider Stest When = r2, we have p(8, = 0 for any E Stest Otherwise, when T = TI, we similarly can define Peest := p(s, 11) and obtain that
Σ p(s,r): ((1(T(a) as))s)) P (=(a)sss) (1 (s) - (T( r)))
#ESuat
=Prest
Σ
(A(T(a) = d" (s,r)) 1a = az,,)))
SESu
(55)
35

---

where (a) holds since for any E Stest. ((T(a) a (s.r)) = 1(a = (s()))) 0 when
r ri,a A2. Combining Eq. (54) and (55), for any r R we have
[F(Pw)]ov(arr) = [P [PF(w)]ov(a.r))
When <E0S>, using Lemma 12 again, we have
[F(w)lov(a, <EOS>)
Wka(<E0S>)
EEp(s,r) (1(a = (s,)))) -
rER
Wko(<E0S>
Wko(r) E((s,)) (1(a = (s,,))) pw (als,)))
(kQ(r)
PER
SES
WKo(<E0S>)
[(ww)oov(rr)
Wkq(r)
rER
Thus we have
[(Pu)]oov(a <B0S>> (ofs>))(P((w)ppp)) <EOS>)
Wko (<BOS>)
(IF(Pw))ovla,, r) (PF(w))ov(arr))
(kQ(r)
TER
0.
Part 4: Proof of Eq. (42d). Given fixed s E Strain, using Lemma 12, we have: [F(w]]kQ(s) - -Wina(s) ((w)
Σ p(s,r)
Σ
()o a" (s,r)) - Peo (a)s,))))
FER
acA
For the RHS of Eq. (42d), we have
[PF - -Ws)(((s)))
Σ p(T(s),r) r)
Wovla, -oo
FR
aed
(g)
Σ r() Σ
a=((s) ((000
(56)
PR
aEA
where (a) uses p(s,T) p(T(s),
e R. On the LHS, we have
[F(PW))ko(s) (0)
Σ p(s,r) r)
Wov(T(a), = a as(,))) PPwl(a)s,))
rER
aEA
Σ p(s,r)
Σ ( a (s, Pw (m(((()) (57)
rER
aEA
(c)
Σ p(s.r)
Σ
Wov n (s)) (z) a" a²(s,)) Pro (z)
PER
zEA
where (a) comes from (Pw)ov x) = Wov(* (s)) for s E Serain- a E A, (b) follows Eq. (49). and in (c) we re-index the summand by = x (a) and thus and Eq. (57), it remains to prove that for any E A,
(zz) Now by comparing Eq. (56)
1(T ((a) = a" (s, = 1(a= a'(T((sss)))
which holds by the definition of ग.
Lemma 15. Assuming Assumption holds. P the permutation matrix defined Lemma 14. If for Jorany& & 2 0, there exists constant R > O such that | S R is bounded, then F(w) == - ((w) Lipschitz with some constant T.. Moreover, since F(w) is permutation equivariant w.r.t P. i.e. F(Pw) (F(w) established in Lemma 14, we have
w(t) Pw(t).
In particular, for any t 2 0, we have
Wov(a, v) Wov(a', v), Vo E Stest U R U <EDS>}, Va,d' A2.
(58)
36

---

Proof. Note that when w is bounded by R. i.e., ||||| < R. F(w) is Lipschitz with constant L using Lemma 13. Moreover, Lemma 14 indicates that F(w) also permutation equivariant. Let
(()) = un()) (t)(), then we have (0) = 0 by Assumption 2. Note that
(t) ii(t) (t)(t) (t))(t) (F(((()) F(ur(t)) F F(Pu((t))
Let w1 wt)), w2 (t) respectively. Using Lemma 13, we get
[ [[F(w(()) F(Pw(()) & LII
Then
d DD(I D(t) D(1) (a)
|D(III LID(I
dt
[
where (a) uses the Cauchy-Schward inequality. Using Gronwall's inequality, for 2 0, we have
|ID(1) < || = 0,
since D 0. Equivalently, for any 2 O, we have:
w(t) Pw(().
Moreover, recall the definition of x, P in Lemma 14. Given fixed Stest U (<E0S>) and any a, a' A2, we have
Wor(a, v) (a) Wov(=(a). v) (
Wov(a', v),
(59)
where (a) follows Definition and (b) holds since for any fixed a, E A2, one can find a permutation such that (a)) = a'.
Now we are ready to prove Theorem
Proof of Theorem 3. First note that for any 2 0, the parameters 0 or equivalently its flattened form (tt) is bounded. Then using Lemma 15, we have
w(l) Pw(().
Consider any time 2 0, s E Skt and ā == a² a', the prediction probability given by
exp(Sø, ((,, ((8,)2))
Po, (ā)s, En'ed exp(fé, ((s, (2),
We proceed by showing that fa. (2,, a) is no larger than Lg. r2), a) for any a E A2. WLOG, consider a fixed a E A2 where # ã. Comparing the logit functions we get fo. (2s,))), - fa. ((),)))), (s: t) Wov((, s;t) - Wov(a,
A(a.a.st))
+ Wka (a, - wov(a,
B(a,ū.22:))
W<a(<EDS>; ( (Wov(a, <EOS>; - Wov(a, <EOS>;
C(a,a,<B05>:4)
-0,
where A(a, a, s;t) 0 B(a, a, r2:t) C(a, a, <E0S>; t) following Lemma 15, As result, for all a E A2 and any / 2 0, we have
fa. ((s, r2),2203 ) a) = fe. ((s, r2),
Thus for any input (s, r2) where St, its prediction probability can be upper bounded
exp(Ja, (2)
exp(Sa, r2),
Pe. (2)
exp(Je Exredexp(la (1,, Evess exp(Já (1,, 1/|A2/,
which implies that
(0,)) =E EeuTLrrraall l-lellllooooo logPo, (2T41/217) > log AA2..
37

---

Non foctorized Thiining imp kcstoons
Non fodt-ddtrizeed Pactin food
Non-hutorited (48 tacts)
fazterized (Traning impiicatinn)
Factor und (Al facis)
Training Loes
Implication
Training OK10P)
Training
Figure 3: Training and Test Implication Loss for Factorized vs. Non-Factorized Models. While both models effectively minimize the training loss (lefi), their performance on unseen test implications differs starkly (right). The factorized model successfully generalizes, achieving low test implication loss and thus demonstrating OCR, while the non-factorized model fails to generalize.
C Additional Experiments for One-layer Models
We provide additional experimental results to verify our theoretical results in Section 3.2.
Test subjects Training subjects
Test subjects Training subjects
(rg.n))
Fact set A
Implication set A
-9
Figure 4: Comparison of full weights of trained one-layer linear attention models. Left: Non- factorized model. Right: Factorized model. The factorized model shows strong OCR capability compared to the non-factorized model.
Loss curves. In Figure 3, we present the training and test loss curve during training, which shows that the factorized model can exhibit OCR while the non-factorized model cannot. Recall that training loss contains two parts: loss on all facts and implications of training subjects (training implication), and the model is evaluated on the implications of test subjects (test implication).
Full weight inspection. In Figure 2, we only showed partial model weights that are related to prediction. For completeness, show the full model weights Wov Wo W. in Figure
SVM solutions. We setup the SVM problems defined in (Wov-SVN) and (WF -VVM) with |S| 12, 3, m 4, marain Solutions by CVXPY [Diamond and Boyd, 2016] are shown on the left side of Figure 5. The results are consistent with the weights of the one-layer attention model Figure 2. Moreover, we decompose
solution of (Wi (SVM) using SVD and keep the directions with singular values larger than 10-5 This results in WF W. with intrinsic dimension dy = 3. On the right of Figure 5, visualize the corresponding rows of the subjects and
38

---

subjects
Training subjects
Fact set Ag
Woy-SVM
Roded
Molbt)
Implication set Aa
Malol
Fact A
WV
Implication set A₃
Figure 5: Comparison solutions (Wov-SVM) and OVV-S. Top Left: (Wov-SVM) with the Frobenius norm objective. Bottom Left: (W (Wo with the nuclear norm objective. Here we only show the partial weights in the output-value matrix related to the prediction, i.e., Wov E [lA[x(mx+2)) Right: Geometric interpretation of Wo and Wy solved in (W OV--VV) All the subjects' feature vectors (corresponding to rows Wv) reside in the zy plane while the relation vectors corresponding to and are orthogonal to the subjects and point in opposite directions. The predictions Wo(b;) (bx), (c:))
made by summing up the feature vector of 8, with r1 r2, which aligns well with the features of b; or c; respectively (plotted in the figure, which are corresponding in Wo) with cosine similarity greater than 0.9.
relations Wy as well as the corresponding rows of answers in Wo, which suggests that the model generalizes effectively even with small hidden dimension.
Lower bound of the intrinsic dimension of output and value matrix. To further support the claim that the factorized model generalizes effectively with small hidden dimension, we train one-layer attention model in Figure 6 and sweep across multiple candidate values for the intrinsic dimension, i.e., dn {3, 4, 8, 16, 32, 128} with the embedding dimension d 128, Marain 3 and other parameters unchanged, which demonstrates that OCR can be achieved efficiently in terms of the number of parameters.
D Additional LLM Experiments on Real-World Data
To verify our claim in real-world datasets, we extended the LLM experiments Section 2 to PopQA [Mallen et al., 2022], which is large-scale open-domain question answering (QA) dataset. Each question in PopQA follows the same format as in our synthetic dataset a knowledge triple (subject, relation, answer). Specifically, we subset of PopQA consisting of the place of birth (POB) and sporí relations termed PopQA-OCR and treat the first relation as fact and the second as implication. We randomly sample 500 subjects and randomly pair 10 facts from available POBs with 10 implications from available sports in PopQA-OCR. We follow the same data processing scheme as in Appendix with 0.5 training ratio. This new controlled dataset exhibits three key distinctions, compared to the synthetic dataset: 1) We scale up the number of subjects and fact-implication pairs, which results in ~10z training samples compared the synthetic one. 2) The new task follows question answering format: "Q: In what city was Antoine Richard born? A: Vinga" while the synthetic one uses a simple text generation formulation. 3) The subject names are real instead of fictitious, which is more likely to collide with the pretrained knowledge, thus inducing new challenges in fine-tuning. The results can be found Table 2, which show that LLMSs continue to exhibit OCR-driven hallucination real-worke data.
39

---

Test Implication Loss
25
50
75
100
125
150
175
200
Training Steps (x10²)
Figure 6: Test loss versus training steps for one-layer linear attention model with varying intrinsic dimensions (dz). The plot demonstrates that the model can exhibit OCR even when the intrinsic dimension is as small d, = 4.
Table 2: Performance comparison of different language models PopQA-OCR, The table reports mean-rant rank scores where rank indicates the position of the ground ---ttttt answer among candidates based on prediction probability. Lower ranks indicate better performance and Rank 0 refers to the token with the largest probablity. Values in parentheses indicate the standard error of the mean-rank scores, calculated from runs with different random seeds.
Models
Gemma-2-98
OLMo-7B
Qwen-2-T 2-7B
Mistral-7B-v0.3
Llama-3-88
PoB-Sport
0.68 (0.27) 1.41 (0.29) 3.30 (0.62)
1.36 (0.70)
2.29 (4.26)
E
Implementation Details
Our code is released at https: //gtthub. cww./yiiiioo-aang/OCR----TT-- We provide addi- tional details on experiments Section 2 and Section 3.2.
Training. Throughout the paper. we finetune the models using the cross-entropy loss with Adam W optimizer [Kingma, 2014]. We build the implementation of Feng et al. [2024] for all LLM experi- ments² and adopt a different training scheme for one-layer transformer as discussed Section 3.1. L.Mss,
For experiments we notice that OCR sensitive to learning rates and thus we sweep across different learning rates in (10-6.3 10-6 10-5 10-5 10 3- 10-4 4) for each model and relation pair and report the results with the lowest test rank. As a complement to the final performance metrics Table 1, Figure plots the average test rank during training across three different seeds. The shaded region represents the standard deviation. For the one-layer linear attention model, we train the model with one-hot token embedding with d 128 for 2 10a steps with learning rate 5 10-4 We set dy default, unless otherwise specified.
full batch and train for 100 epochs. Similar to Feng al. [2024],
128 by
Dataset. The dataset for LLM experiments consists of a list of 100 fictitious names and a list of 5 fact-implication pairs with different topics, namely, city-animal country- code, sport-music, and profession We split the subjects into training and test with a ratio of 0.2 0.8, i.c., there are 4 training subjects in each subset S; for E (5) and 20 training subjects in total. For the experiments in Section 2, we use relations listed in Table For example, for the topic "Country-Codee one example given by
"https://githinn
40

---

(a)
(b) City-Langun Language (counterfactual)
Tex rast
Text 1m
Genre-30 Ovn
Ownnd
Ovn 13 Laraa
(c) country-code
(d) sport-music
(c) profession-color
Figure 7: OCR performance of various LLMs the five relation pairs. Results are averaged over random seeds, with shaded regions representing the standard deviation.
Fact: "Daniel Gray was born in Brazil."
Implication: "Daniel Gray codes in Assembly."
Topic
Relation
City
lives in
Language
speaks
Color
dislikes
Country
was born in
Music listens
Code
codes in
Sport
plays is a
Profession
Table Relation expressions of different topics.
We use the same name, city and language list in [Feng al., 2024]. For the counterfactual language.
we re-order the language list such that every city corresponds to an incorrect language. For the rest of the topics, Claude-3.5-somme generate list of 5 examples. complete list is given below:
Topic Lists
City: Tokyo, Beijing, Mumbai, Paris, Berlin
Language:
Japanese, Mandarin, Marathi, French, German
Language (counterfactual):
Marathi, German, Japanese, Mandarin, French
Color:
crimson, teal, navy blue, emerald green, lavender
Country:
Japan, Brazil, Morocco, New Zealand, Iceland
41

---

Music:
jazz, alternative rock, reggac, classical, hip-hop
Code:
Python, Julia, Assembly, C, MATLAB
Sport:
basketball, soccer, tennis, swimming, volleyball
Profession:
nurse practitioner, computer scientist, journalitt, veterinarian, social worker
Computation.
The experiments for the one-layer model were run on a single NVIDIA A100 GPU.
LLM Experiments were run on cluster of NVIDIA A100 GPUs and took less than an hour for each run.
42

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We showed empirical results in Sections 2 and 3.2 and theoretical results in Section 4 claimed in abstract and introductions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: As discussed in Section 5, our theoretical results only focus on one-layer transformers, and future work can extend it to multi-layer transformers.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
43

---

Answer: [Yes]
Justification: Our theoretical results and assumptions are presented Section and missing proofs are in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provided experimental details Sections 2 and 3.2 and additional details in Appendices C and E.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
44

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The code link is included in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //wpp...c//////// httpccccpppppcpcceeeeeeecc Stuppbliii/guiiieeeeeeutuuuuaiii
/Codeeuboeeeeeeee ceobbbmsssiiooo sswinnoilll. onPoiicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provided experimental details Sections 2 and 3.2 and additional details in Appendices and E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report the variance of the experiment results Table and provided error bars for the LLM experiments in Figure 7.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
45

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provided the details of compute resources in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:: //teuri////urip.... ctt.ccss.......occcoooooooo cc/ Cpuulicc/thicsauut ines?
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work aims to theoretically understand transformer's OCR capability, and does not have a direct societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
46

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly cited the code we used.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi twwode..coo..... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
47

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification:
No assets released.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
48

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: We only use LLM for grammar checking.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components. wwt.c/////ooccoee
Please refer to L.LM policy (https: httprrpss.cc//oofeeenneeee/////s for what should or should not be described.
49


Reviews:
[{"review_id": "IdlnuzNwUd", "reviewer": "Reviewer_7jeV", "summary": "This paper explores how transformer models can both generalize from new information and hallucinate incorrect facts. Specifically, it investigates a phenomenon called Out-of-Context Reasoning (OCR) where a model makes inferences based on partial patterns seen during training. The authors ask whether this same reasoning mechanism is responsible for both useful generalization and harmful hallucination. They theoretically analyze a very simple model and surprisingly show that even this basic model can perform OCR, learning patterns from just a few examples. Their analysis shows that this ability comes from the way the model is trained, specifically how gradient descent favors certain types of solutions that can link facts together. However, this also means the model may make false connections when two things only appear together but aren’t truly related.", "questions": "See above", "limitations": "The authors adequately discuss the technical limitations of their work, particularly the focus on one-layer transformers and the use of synthetic data.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths \n- Interesting topic that investigates why LLMs can both generalize well and hallucinate.\n- Shows that even a one-layer attention-only transformer can perform OOC reasoning \n- Connects model behavior to gradient descent’s implicit bias.\n\n\nWeaknesses\n- Focuses only on one-layer transformers; unclear whether the conclusions fully apply to complex models. \n- It is fine to use synthetic data for evaluation, but it would be helpful to see and discuss how the model performs on tasks such as question answering. \n- Identifies the cause of hallucination that is a good result, but how can we reduce or prevent it, based on the analysis provided in the paper? \n- The proposed optimization as well as the use of SVMs are difficult to follow, could the authors provide some clarifications?", "was_revised": true, "final_justification": "All my concerns have been properly addressed. I raised my score from 4 to 5.", "rebuttal": "We thank the Reviewer for their thoughtful feedback and for recognizing our theoretical contributions to understanding generalization and hallucination in LLMs.\n\n> Weakness 1\n\nWe focus on a one-layer transformer for our theoretical analysis since a one-layer transformer is sufficient to perform the OCR task and discover the underlying mechanism. Despite its simplicity, previous work doesn’t have a deep understanding of the mechanism and our theoretical analysis is highly non-trivial and provides practical insights on linking generalization and hallucination.\nAlthough the result might not be fully the same for complex models, our experiments for LLMs show that the theoretical results also apply to more complex models. We further conducted a probing experiment to verify this. Figure 7 (right) in Appendix D gives two important findings for one-layer transformers. \n- (C1) The prediction is made by summing up the feature vector of subject $s\\_i$ and relation $r\\_1$ or $r\\_2$.\n- (C2) Subjects that correspond to the same fact-implication pair will point to the same direction in hidden representation.\n\n**We hypothesis that these two findings also hold in LLMs.** We collect the hidden states at the last position in the last layer for all subjects with two different relations (*country-code*) - termed as **fact states**, $H\\_1$ and **implication states**, $H\\_2$, where $H\\_1, H\\_2 \\in \\mathbb{R}^{|S| \\times d}$\ncan be understood as the prediction vector. We introduce three metrics for verifying the two claims in LLMs.\n- **Average-pairwise correlation (APC)** measures how consistently vectors within the same group point in the same direction. \n$$\n \\mathrm{APC}\n = \\frac{2}{b(b-1)}\n \\sum\\_{1 \\le i < j \\le b}\n \\frac{h\\_i^\\top h\\_j}{||h\\_i||||h\\_j||},\n $$\n where $h\\_i, h\\_j \\in \\mathbb{R}^d$ and $b$ is group size. $\\mathrm{APC}=1$ iff all vectors are colinear. To verify **Claim C1**, we can compute the difference matrix $W\\_1 = H\\_1 - H\\_2$. If C1 holds, then every row of $W\\_1$ will align with the direction of $r\\_1 - r\\_2$, i.e., $\\rm APC(W\\_1) \\approx 1$.\n- **Intra-group variance** quantifies the spread of each group around its own mean, averaged over $K$ groups. \n $$\n V\\_{\\mathrm{intra}}\n = \\frac{1}{K}\\sum\\_{k=1}^K\n \\frac{1}{n\\_k}\\sum\\_{i:g\\_i=k}\n ||h\\_i - \\mu^{(k)}||^2,\n $$\n where $\\displaystyle \\mu^{(k)} = \\frac{1}{n\\_k}\\sum\\_{g\\_i=k}h\\_i$ is the $k$-th group mean.\n- **Inter-group variance** measures how far apart the group centroids are from the global mean. \n $$\n V\\_{\\mathrm{inter}}\n = \\frac{1}{K}\\sum\\_{k=1}^K\n ||\\mu^{(k)} - \\mu||^2,\n \\quad\n \\mu = \\frac{1}{b}\\sum\\_{i=1}^b h\\_i.\n $$\n A larger $V\\_{\\mathrm{inter}}$ indicates greater separation between groups. To verify **Claim C2**, we form the “sum” feature for each example: $W\\_{2} = H\\_{1} + H\\_{2}$. We then partition these vectors into $n$ groups by their fact–implication relation pair.If $r\\_{1}\\approx -r\\_{2}$, then\n\n  $$\n  [W\\_{2}]\\_i\n  =(s\\_i + r\\_1)+(s\\_i + r\\_2)\n  \\approx 2s\\_i,\n  $$\n\nso that\n- $V\\_{\\mathrm{intra}}\\approx 0$ (all vectors in each group collapse to the same direction), and\n- $V\\_{\\mathrm{inter}}>0$ (different subjects $s\\_i$ yield distinct directions).\n - **When $r\\_1 \\neq -r\\_2$:** The constant shift $r\\_1 + r\\_2$ appears in every row of $W\\_{2}$, but adding a uniform vector to all samples does *not* change either intra‐ or inter‐group variance, since variance is invariant under translation. \n \nThe results for Gemma-2-9B and Olmo-7B on relation pair *country-code* is shown below:\n- Result of Gemma-2-9B\n\n| Metrics| Epoch 0 | Epoch 40 | Epoch 100 |\n|-------------|---------|----------|-----------|\n|APC| 0.89\t| 0.93 \t| 0.94|\n| Intra-group variance | 0.20\t| 0.02\t| 0.02|\n| Inter-group variance| 0.0\t| 0.02 \t| 0.28|\n\n- Result of Olmo-7B\n\n|Metrics| Epoch 0 | Epoch 40 | Epoch 100 |\n|---|--------|--|---|\n| APC \t| 0.82\t| 0.87 \t| 0.88\t|\n| Intra-group variance \t| 0.24\t| 0.10 \t| 0.09\t|\n| Inter-group variance \t| 0.01\t| 0.95 \t| 0.94\t|\n\nWe observe three consistent trends across both models as training proceeds from Epoch 0 → 40 → 100: \n- The average pairwise correlation steadily increases, showing that vectors within each fact–implication group become more tightly aligned.\n- Intra-group variance drops sharply, indicating that same-group representations collapse toward a common direction.\n- Inter-group variance rises, meaning that different subject groups diverge from one another.\n\n> Weakness 2\n\nThanks for the valuable suggestion. We have extended the LLM experiments to PopQA [1], which is a large-scale open-domain question answering (QA) dataset. We show that LLMs continue to exhibit OCR‑driven hallucinations on real‑world data. Due to the character limit, please see our response to Reviewer PxpD for further details.\n\n> Weakness 3\n\nThis is an excellent question. Our analysis suggests that because both generalization and hallucination arise from the same underlying mechanism, naively reducing hallucination risks degrading benign generalization. To further analyze this, we drew on mechanistic interpretability tools—specifically the logit lens [2]—which project each hidden state back into vocabulary logits via the output embedding matrix. We measured the layer at which the correct target label’s logit at the final sequence position first exceeds $0.01$. The results below reveal a clear contrast between generalization and hallucination cases:\n\n| Models | Generalization (City-Language) | Hallucination (Country-Code) |\n|---|---|---|\n| Gemma-2-9B | 28 | 4 |\n| Olmo-7B | 21 | 6 |\n\n**Strikingly, in generalization cases (i.e., *city-language*), the correct label’s logit consistently emerges in the deeper layers, whereas in hallucination cases it peaks much earlier.** This layer‑wise separation indicates a viable strategy to distinguish—and thus mitigate—hallucination without harming generalization. While designing concrete mitigation techniques lies beyond the scope of this work, it indicates a promising direction for future research. The layer-wise separation we uncover offers a solid foundation for developing targeted approaches that suppress hallucinations without impairing generalization.\n\n> Weakness 4\n\nThanks for pointing that out. We hereby post two clarifications:\n1. **Link with SVM:** Soudry et al. [3] first showed that gradient descent on exponential‑tailed losses (e.g., logistic regression) converges to the maximum‑margin solution of an equivalent SVM: it separates positive and negative logits under margin constraints. A one‑layer transformer trained with cross‑entropy loss—once its attention weights are fixed—behaves like a multiclass extension of logistic regression. Building on this insight [4, 5], we derive Theorem 1: the factorized and non‑factorized models correspond to solutions of two different SVM objectives, which we term the $W\\_{OV}^F$‑SVM and the $W\\_{OV}$‑SVM, respectively.\n2. **Difference in optimizing the factorized vs. non-factorized models:** As shown in Theorem 1, the factorized model’s SVM objective takes the form:\n$$\n\\min\\_{W\\_O,W\\_V} \\frac{1}{2}(||W\\_O||\\_F^2 + ||W\\_V||\\_F^2)\n\\quad \\text{subject to margin constraints on }W\\_O W\\_V^\\top.\n$$\n\nLet $||W||\\_*$ denote its nuclear norm. Because\n\n$$\n||W\\_{OV}^F||\\_*^2 =\\min\\_{W\\_O W\\_V^\\top = W\\_{OV}^F}\\frac{1}{2}(||W\\_O||\\_F^2 + ||W\\_V||\\_F^2),\n$$\n\nit is equivalent to minimizing the nuclear norm of $W\\_{OV}^F$ with margin constraints on it.\nBy contrast, the non‑factorized model solves:\n\n$$\n\\min\\_{W\\_{OV}}\\frac{1}{2}||W\\_{OV}||\\_F^2 \n\\quad\\text{subject to those identical constraints.}\n$$\n\nThe two formulations differ only in whether they penalize the Frobenius norm or the nuclear norm which is the sum of singular values. **This shift—from a Frobenius‑norm penalty to a nuclear‑norm penalty—is what ultimately drives the divergence in their OCR behaviors:**\n\nWe can view the optimization problem as completing a lookup table represented by the matrix $W$, where each column indexes either a subject or a relation, and each row corresponds to a fact or an implication. Because the Frobenius norm is the sum of the squares of all entries in $W$, any entry not involved in a training constraint is driven to be zero to minimize the objective. Since test implications do not appear in the training data, their corresponding entries remain zero. This explains why the red‑boxed region of the non‑factorized model $W_{OV}$ in Figure 4 consists of zeros.\n\nIn contrast, the nuclear norm is defined as the sum of a matrix’s singular values, with each singular value representing an independent orthogonal direction of the transformation. If we were to force the red‑boxed region (in Figure 4) of the factorized model $W_O W_V^\\top$ to be zero, we would introduce additional orthogonal components, thereby increasing the nuclear norm. To avoid this penalty, the factorized model naturally assigns nonzero values in that region.\nIn summary, the difference in the objectives of the two SVM problems leads to two different solutions, which further results in different OCR capabilities. We will include the above two points in the extra page in the camera-ready version.\n\n--- \n\n**References:**\n\n[1] Mallen, A., Asai, A., Zhong, V., Das, R., Khashabi, D., and Hajishirzi, H. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories, 2023.\n\n[2] nostalgebraist. 2020. interpreting gpt: the logit lens.\n\n[3] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1–57, 2018.\n\n[4] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890, 2019.\n\n[5] Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and relu networks. Advances in Neural Information Processing Systems, 35:37024–37036, 2022."}, {"review_id": "LlVftnJxpM", "reviewer": "Reviewer_XZxz", "summary": "This paper empirically and theoretically demonstrates that a single-layer linear attention can perform out-of-context reasoning (OCR). Specifically, the transformer can predict $s',r_2\\rightarrow b$ given $s,r_1\\rightarrow a$,   $s,r_2\\rightarrow b$,  and $s',r_1\\rightarrow a$. Further experiments on synthetic tasks prove that real-world LLMs also exhibit such OCR capability. However, such capability can cause hallucinations once the associations in training data are not causal relationships, revealing a potential cause of the hallucinations.", "questions": "Could you extend the identified capability of transformers to perform the OCR task defined in your paper to more general reasoning modes? I think either empirically or theoretically extending your findings to explain the failure or capability in more general reasoning patterns will be helpful to improve your work.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. This work provides novel insights into the formation of hallucinations in LLMs by theoretically proving that a single-layer linear attention is capable of achieving OCR.\n\n2. This work theoretically reveals the defect of the parameterization of the output-value matrix widely adopted by previous works.\n\n3. The paper is well-written and easy to follow.\n\n**Weaknesses**\n\nThe theory and experiments consider a very simple case where only two facts are correlated. While I acknowledge that such a setup may reflect certain real-world situations, it fails to capture the complexity of most reasoning patterns. In more realistic cases, the \"facts\" can be modeled as a causal graph with multiple variables and relationships. For instance, in your OCR training data, since $s,r_1\\rightarrow a \\Rightarrow s,r_2\\rightarrow b$, we can imagine a causal graph in which there exists a confounder $C$ such that $C\\rightarrow A$ and $C\\rightarrow B$, where event $A$ represents  $s,r_1\\rightarrow a$ and event $B$ represents $s, r_2\\rightarrow b$. $A$ and $B$ are correlated through $C$. However, this represents only a rudimentary causal graph and offers limited applicability to more complex scenarios.\n\nMoreover, according to the definition in your paper, whether the model \"generalizes\" or \"hallucinates\" depends on whether the correlation between $A$ and $B$ is causal. Specifically, if $A\\rightarrow B$ across training and test data, then the model \"generalizes\"; otherwise, the model produces hallucinations. In my opinion, however, \"generalization\" should refer to the model's ability to learn generalizable features or relationships, such as capturing a causal link in a causal graph, so it may not be suitable to regard the capability of successfully solving the tasks in Fig. 5 as \"generalization\".", "was_revised": true, "final_justification": "Thank you for the clarification. After reading your response, I better understand the challenges involved in extending OCR capabilities to more complex scenarios. I also appreciate your refinement of the term “generalization” into the more precise and rigorous notion of “OCR capability,” which appropriately addresses the concerns around its interpretation. Given the nontrivial nature of theoretically establishing OCR capability and its fundamental significance, I have decided to raise my score to 4.", "rebuttal": "We thank the reviewer for their insightful feedback and for highlighting our novel theoretical insights, our analysis of the common output-value reparameterization defect, and the clarity of our presentation. Here are our responses to your concerns:\n\n> The theory and experiments consider a very simple case where only two facts are correlated. While I acknowledge that such a setup may reflect certain real-world situations, it fails to capture the complexity of most reasoning patterns. In more realistic cases, the \"facts\" can be modeled as a causal graph with multiple variables and relationships. For instance, in your OCR training data, since $s,r_1\\rightarrow a \\Rightarrow s,r_2\\rightarrow b$, we can imagine a causal graph in which there exists a confounder $C$ such that $C\\rightarrow A$ and $C\\rightarrow B$, where event $A$ represents $s,r_1\\rightarrow a$ and event $B$ represents $s, r_2\\rightarrow b$. $A$ and $B$ are correlated through $C$. However, this represents only a rudimentary causal graph and offers limited applicability to more complex scenarios.\n\nAlthough our setup is simple, it captures one of the most fundamental causal-graph structures, and its analysis is both meaningful and highly nontrivial. As mentioned in Section 1.1, prior work has examined OCR from multiple angles [1–4], yet a deep theoretical understanding remains elusive. **As a first step, our study provides rigorous theoretical analysis and practical insights that help unify generalization and hallucination.** We agree that extending this framework to more complex causal graphs is an important next step, but it requires highly nontrivial effort (since the analysis of our setting already requires very sophisticated techniques) and is beyond the scope of our paper. We also emphasize that building a solid fundation on the most fundamental setup is of equal importance, which is the main focus of our paper. \n\n**Moreover, our theoretical analysis provides new insights into the minimization of non-smooth norms**—a central component of matrix completion algorithms. This topic is under-explored in the literature [5], largely because the non-smooth nature of the nuclear norm hinders rigorous analysis (Similar issues in the vector form would be LASSO, which is already extremely complex in theory). By exploiting a key observation on the symmetry induced by the minimization process (Lemmas 7-10), our paper offers one of the first concrete characterizations of the solutions to nuclear norm minimization.\n\n\n> Moreover, according to the definition in your paper, whether the model \"generalizes\" or \"hallucinates\" depends on whether the correlation between $A$ and $B$ is causal. Specifically, if $A\\rightarrow B$ across training and test data, then the model \"generalizes\"; otherwise, the model produces hallucinations. In my opinion, however, \"generalization\" should refer to the model's ability to learn generalizable features or relationships, such as capturing a causal link in a causal graph, so it may not be suitable to regard the capability of successfully solving the tasks in Fig. 5 as \"generalization\".\n\nWe agree that framing generalization as “the model’s ability to learn truly generalizable features or relationships—e.g., capturing an actual causal link in a causal graph”—is indeed consistent with prior OCR work [4]. However, this raises the issue that a model could “learn” spurious rules (for instance, “If Raul lives in France, then Raul codes in Java”) even when there is no true causal link, which we would classify as hallucination rather than genuine generalization. To distinguish these two behaviors in LLMs, we therefore introduce a more precise criterion based on the presence or absence of a causal relationship between the two relations. In this paper, we show that both hallucination on causally unrelated knowledge and true generalization on causal links arise from the same underlying mechanism—OCR. Finally, we appreciate the reviewer’s suggestion to correct the caption of Figure 2 and 5: we will replace “generalization” with “OCR capability” in the revision.\n\n> Could you extend the identified capability of transformers to perform the OCR task defined in your paper to more general reasoning modes? I think either empirically or theoretically extending your findings to explain the failure or capability in more general reasoning patterns will be helpful to improve your work.\n\nPlease see our response to the first weakness above.\n\n---\n\n**References:**\n\n[1] Dmitrii Krasheninnikov, Egor Krasheninnikov, and David Krueger. Out-of-context meta-learning in large language models. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.\n\n[2] Lukas Berglund, Asa Cooper Stickland, Mikita Balesni, Max Kaufmann, Meg Tong, Tomasz Korbak, Daniel Kokotajlo, and Owain Evans. Taken out of context: On measuring situational awareness in llms. arXiv preprint arXiv:2309.00667, 2023a.\n\n[3] Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. arXiv preprint arXiv:2309.14402, 2023.\n\n[4] Jiahai Feng, Stuart Russell, and Jacob Steinhardt. Extractive structures learned in pretraining enable generalization on finetuned facts. arXiv preprint arXiv:2412.04614, 2024.\n\n[5] Tim Hoheisel and Elliot Paquette. Uniqueness in nuclear norm minimization: Flatness of the nuclear norm sphere and simultaneous polarization. Journal of Optimization Theory and Applications, 197(1):252–276, Feb 2023. doi: https://doi.org/10.1007/s10957-023-02167-7. URL https://link.springer.com/article/10.1007/s10957-023-02167-7."}, {"review_id": "lHPujX5RXj", "reviewer": "Reviewer_6AQL", "summary": "This paper focuses on the mechanisms underlying generalization and hallucination in large language models (LLMs), exploring the out-of-context reasoning (OCR) capabilities of Transformers through theoretical analysis and empirical studies. The core contributions lie in revealing that a single-layer attention-based Transformer can achieve fact-related reasoning through parameter decomposition and explaining the unified mechanism of generalization and hallucination", "questions": "1) From the view of OCR, how can we design models to improve the generalization rather than hallucination?\n2) Can the authors construct a benchmark to test the OCR problem?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "-Strengths: 1) Synthesis of theory and experiments unifies explanations of generalization and hallucination, contributing to theoretical understanding 2) Hallucination is a challenging problem in LLMs; this paper gives a deep analysis of this question from the view of OCR. \n- Weakness: 1) Idealized synthetic data and lack of benchmark model comparisons limit conclusion; 2) Limited analysis of practical implications; I mean the authors should step forward to investigate how to design anti-hallucination techniques from the view of OCR; 3) More different model structures should be tested in the experiments.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing our unified theoretical and empirical insights into generalization and hallucination. Below, we address your questions:\n\n> Idealized synthetic data and lack of benchmark model comparisons limit conclusion.\n\n**We extended the LLM experiments to PopQA [1], which is a large-scale open-domain question answering (QA) dataset.** Each question in PopQA follows the same format as in our synthetic dataset – a knowledge triple (subject, relation, answer). Specifically, we use a subset of PopQA consisting of the *place of birth (POB) and sport* relations termed **PopQA-OCR** and treat the first relation as fact and the second as implication. We randomly sample 500 subjects and randomly pair 10 facts from available POBs with 10 implications from available sports in PopQA-OCR. We partition $S$ into $S = \\cup_{i = 1 }^5 S_i$ and randomly assign a distinct fact-implication pair (or equivalently, a ($b_i , c_i$) pair) to each subset $S_i$ . We then create training and test sets for each subset by splitting its subjects with a 0.5 training ratio. The training set contains facts for all subjects and implications only for training subjects. We then evaluate the model on the implications of the test subjects. \n**This new controlled dataset exhibits three key distinctions, compared to the synthetic dataset:**\n- We scale up the number of subjects and fact-implication pairs, which results in ~10x training samples compared to the synthetic one.\n- The new task follows a question answering format: “Q: In what city was Antoine Richard born? A: Vinga” while the synthetic one uses a simple text generation formulation.\n- The subject names are real instead of fictitious, which is more likely to collide with the pretrained knowledge, thus inducing new challenges in fine-tuning. \n\n**To better analyze OCR's generalization versus hallucination effects, we expanded the original synthetic dataset with two relation pairs: _City-Language_ and _City-Language (CF)_.** The *City-Language* pair uses real-world associations likely learned during pretraining (e.g., \"People in Paris speak French\"), testing generalization capabilities. In contrast, *City-Language (CF)* creates counterfactual associations by deliberately mispairing cities with incorrect languages (e.g., mapping \"Paris\" to \"Japanese\"), designed to analyze hallucination alongside the four other fictitious relation types. \n\nBelow are the updated results for the newly added relation pairs. The first column reports performance on the generalization task (*City-Language*), while the remaining columns show hallucination cases across four counterfactual relations. We measure the mean rank, where the rank indicates the position of the ground-truth answer among all candidates based on prediction probability. Lower ranks indicate better OCR capabilities. Values in parentheses indicate the standard error of the mean rank, calculated from 3 runs with different random seeds.\n\nModels           | *Gen: City-Language* | *Hall: City-Language (CF)* | *Hall: PoB-Sport* | Hall: Country-Code | Hall: Profession-Color | Hall: Sport-Music\n-----------------|--------------------|-------------------------|-----------------|--------------------|------------------------|-------------------\nGemma-2-9B       | 0.00 (0.00)        | 0.19 (0.20)             | 0.68 (0.27)     | 0.19 (0.07)        | 1.64 (0.01)            | 0.56 (0.01)\nOLMo-7B          | 0.07 (0.03)        | 1.33 (0.49)             | 1.41 (0.29)     | 0.15 (0.13)        | 1.84 (0.23)            | 0.17 (0.01)\nQwen-2-7B        | 0.13 (0.01)        | 4.55 (2.33)             | 3.30 (0.62)     | 3.63 (1.10)        | 0.82 (0.34)            | 0.40 (0.08)\nMistral-7B-v0.3  | 0.00 (0.00)        | 2.10 (0.01)             | 1.36 (0.70)     | 1.48 (0.52)        | 1.15 (0.56)            | 1.28 (0.13)\nLlama-3-8B       | 0.00 (0.00)        | 1.18 (0.61)             | 2.29 (4.26)     | 0.77 (0.10)        | 0.93 (0.21)            | 0.63 (0.22)\n\nThe new experimental results reveal two key findings:\n- **LLMs continue to exhibit OCR‑driven hallucinations on real‑world data.** As observed in the paper, OCR behavior is highly sensitive to the learning rate; we have not yet conducted an exhaustive sweep due to time constraints. Consequently, we believe the performance can be further improved with more thorough hyperparameter tuning.\n\n- **We find that the generalization performance is better than the hallucination performance**—presumably because the injected causal facts align with the model’s pretrained knowledge and are therefore easier to learn. In contrast, in the hallucination setup (especially on PopQA‑OCR), the new information often collides with the model’s pretrained knowledge, making it harder to acquire. **This gap underscores the need for a specialized dataset to systematically evaluate OCR.** Importantly, our synthetic dataset can avoid such conflicts by using fictitious names and can be scaled up by adding more subject–relation pairs and by incorporating richer, multi‑step causal graphs (as suggested by Reviewer XZxz) to meet this need. We view this expansion as an exciting avenue for future work.\n\n\n> Limited analysis of practical implications; I mean the authors should step forward to investigate how to design anti-hallucination techniques from the view of OCR.\n\nThis is an excellent question. Our analysis suggests that because both generalization and hallucination arise from the same underlying mechanism, naively reducing hallucination risks degrading benign generalization. To further analyze this, we drew on mechanistic interpretability tools—specifically the logit lens [2]—which project each hidden state back into vocabulary logits via the output embedding matrix. We measured the layer at which the correct target label’s logit at the final sequence position first exceeds $0.01$. The results below reveal a clear contrast between generalization and hallucination cases:\n\n| Models | Generalization (City-Language) | Hallucination (Country-Code) |\n|---|---|---|\n| Gemma-2-9B | 28 | 4 |\n| Olmo-7B | 21 | 6 |\n\n**Strikingly, in generalization cases (e.g., *city-language*), the correct label’s logit consistently emerges in the deeper layers, whereas in hallucination cases it peaks much earlier.** This layer‑wise separation indicates a viable strategy to distinguish—and thus mitigate—hallucination without harming generalization. While designing concrete mitigation techniques lies beyond the scope of this work, it indicates a promising direction for future research. The layer-wise separation we uncover offers a solid foundation for developing targeted approaches that suppress hallucinations without impairing generalization.\n\n> More different model structures should be tested in the experiments.\n\nIn our LLM experiments, we already evaluated several transformer-based families of varying architectures to validate our theoretical findings. Due to computational constraints, however, we were unable to include even larger-scale models in this study. We agree that extending our evaluation to additional—and particularly larger—model classes would be valuable, and we plan to explore this in future work as more compute becomes available\n\n> From the view of OCR, how can we design models to improve the generalization rather than hallucination?\n\nPlease see our response to the second weakness above.\n\n> Can the authors construct a benchmark to test the OCR problem?\n\nPlease see our response to the first weakness above.\n\n---\n\n**References:**\n\n[1] Mallen, A., Asai, A., Zhong, V., Das, R., Khashabi, D., and Hajishirzi, H. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories, 2023.\n\n[2] nostalgebraist. 2020. interpreting gpt: the logit lens."}, {"review_id": "EXyI34uHoh", "reviewer": "Reviewer_PxpD", "summary": "This work presents an theoretical analysis which shows that a factorized, one-layer attention-only Transfoer is enough to obtain a generalizable deduction ability from learned facts, while a simpler model like a single-head attention or a non-factorized model is not. Meanwhile, such a deduction ability also also causes LLMs to hallucinate easily by deducing false implications when two events are not causally related but only co-occur several times, which provides a possible explanation of hallucination in LLMs. While experiments are conducted on small-scale datasets, they have covered several recent LLM famlilies.", "questions": "I would like to hear the authors' opinion on whether the experiments can be extended to large-scale, real data.", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "I believe this is a theoretically solid work that provides a well-motivated explanation for Transformers' deductive abilities, and why it sometimes lead to non-factual hallucination. Several supported claims in this work are meaningful and interesting. On the other hand, the main shortcome of this work is that its experiments only used a very small sythetic dataset. There are plenty of large, real-world datasets for deductive reasoning that could have been tested on, such as ProntoQA-OOD [1], EureQA [2], or any of the KGC datasets [3]. The current small-scaled experiments provides a decent explanation to the claims, but leave the generalization of these claims uncertain.\n\n\n[1] Saparov et al. Testing the General Deductive Reasoning Capacity of Large Language Models Using OOD Examples. NeurIPS 2023  \n[2] Li et al. Deceptive Semantic Shortcuts on Reasoning Chains: How Far Can Models Go without Hallucination? NAACL 2024  \n[3] Petroni et al. Language Models as Knowledge Bases? EMNLP 2019", "was_revised": true, "final_justification": "I believe this is a theoretically solid work that provides a well-motivated explanation for Transformers' deductive abilities. Meanwhile, I share similar thoughts with some other reviewers that the shortcoming of this work is mainly at its reliance on synthetic data in experiments. The experimentation could have been conducted in a more convincing way using the recommended large datasets. Therefore, my final score is a weak acceptance.", "rebuttal": "We thank the Reviewer for their insightful question and their appreciation of our theoretical contribution. We are delighted to answer their question with a “yes” as we extended the LLM experiments to large-scale and real-world dataset.\n\nProntoQA-OOD and EureQA are indeed large benchmarks—but they focus on in-context reasoning, which lies beyond our paper’s out-of-context reasoning setup. **Nevertheless, we extended the LLM experiments to PopQA [1], which is a large-scale open-domain question answering (QA) dataset.** Each question in PopQA follows the same format as in our synthetic dataset – a knowledge triple (subject, relation, answer). Specifically, we use a subset of PopQA consisting of the *place of birth (POB) and sport* relations termed **PopQA-OCR** and treat the first relation as fact and the second as implication. We randomly sample 500 subjects and randomly pair 10 facts from available POBs with 10 implications from available sports in PopQA-OCR. We partition $S$ into $S = \\cup_{i = 1 }^5 S_i$ and randomly assign a distinct fact-implication pair (or equivalently, a ($b_i , c_i$) pair) to each subset $S_i$ . We then create training and test sets for each subset by splitting its subjects with a $0.5$ training ratio. The training set contains facts for all subjects and implications only for training subjects. We then evaluate the model on the implications of the test subjects. \n**This new controlled dataset exhibits three key distinctions, compared to the synthetic dataset:**\n- We scale up the number of subjects and fact-implication pairs, which results in ~10x training samples compared to the synthetic one.\n- The new task follows a question answering format: “Q: In what city was Antoine Richard born? A: Vinga” while the synthetic one uses a simple text generation formulation.\n- The subject names are real instead of fictitious, which is more likely to collide with the pretrained knowledge, thus inducing new challenges in fine-tuning. \n\n**To better analyze OCR's generalization versus hallucination effects, we expanded the original synthetic dataset with two relation pairs: _City-Language_ and _City-Language (CF)_.** The *City-Language* pair uses real-world associations likely learned during pretraining (e.g., \"People in Paris speak French\"), testing generalization capabilities. In contrast, *City-Language (CF)* creates counterfactual associations by deliberately mispairing cities with incorrect languages (e.g., mapping \"Paris\" to \"Japanese\"), designed to analyze hallucination alongside the four other fictitious relation types\n\nWe measure the mean rank, where the rank indicates the position of the ground-truth answer among all candidates based on prediction probability. Lower ranks indicate better OCR capabilities. Values in parentheses indicate the standard error of the mean rank, calculated from 3 runs with different random seeds.\n\nModels           | *Gen: City-Language* | *Hall: City-Language (CF)* | *Hall: PoB-Sport* | Hall: Country-Code | Hall: Profession-Color | Hall: Sport-Music\n-----------------|--------------------|-------------------------|-----------------|--------------------|------------------------|-------------------\nGemma-2-9B       | 0.00 (0.00)        | 0.19 (0.20)             | 0.68 (0.27)     | 0.19 (0.07)        | 1.64 (0.01)            | 0.56 (0.01)\nOLMo-7B          | 0.07 (0.03)        | 1.33 (0.49)             | 1.41 (0.29)     | 0.15 (0.13)        | 1.84 (0.23)            | 0.17 (0.01)\nQwen-2-7B        | 0.13 (0.01)        | 4.55 (2.33)             | 3.30 (0.62)     | 3.63 (1.10)        | 0.82 (0.34)            | 0.40 (0.08)\nMistral-7B-v0.3  | 0.00 (0.00)        | 2.10 (0.01)             | 1.36 (0.70)     | 1.48 (0.52)        | 1.15 (0.56)            | 1.28 (0.13)\nLlama-3-8B       | 0.00 (0.00)        | 1.18 (0.61)             | 2.29 (4.26)     | 0.77 (0.10)        | 0.93 (0.21)            | 0.63 (0.22)\n\nThe new experimental results reveal two key findings:\n- **LLMs continue to exhibit OCR‑driven hallucinations on real‑world data.** As observed in the paper, OCR behavior is highly sensitive to the learning rate; we have not yet conducted an exhaustive sweep due to time constraints. Consequently, we believe the performance can be further improved with more thorough hyperparameter tuning.\n\n- **We find that the generalization performance is better than the hallucination performance**—presumably because the injected causal facts align with the model’s pretrained knowledge and are therefore easier to learn. In contrast, in the hallucination setup (especially on PopQA‑OCR), the new information often collides with the model’s pretrained knowledge, making it harder to acquire. **This gap underscores the need for a specialized dataset to systematically evaluate OCR.** Importantly, our synthetic dataset can be scaled up to meet this need, since it can avoid such conflicts by using fictitious names.\n\n---\n\n**References:**\n\n[1] Mallen, A., Asai, A., Zhong, V., Das, R., Khashabi, D., and Hajishirzi, H. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories, 2023."}]

Author rebuttal:

Meta-review: The paper studies how Transformers learn to deduce implications from facts seen during training. In particular, the authors analyze a simplified one-layer Transformer with fixed attention weights (essentially reduces to a logistic regression problem) and show that the inductive bias of gradient descent to low nuclear norm solution leads to solutions that link facts together resulting in both generalization and hallucinations. This crucially utilizes the factored representation of the value and output weight matrices. The authors support their hypothesis through evaluation on LLMs using mostly synthetic data (extended to real datasets during rebuttal period).

The problem being considered is interesting and relevant, and the formulation is clean. While the theoretical analysis seems mostly based on prior work and in a simplified setting (attention weights fixed), the paper does add to our understanding of the relationship of different parameterizations and their impact on OOD performance, that may be of practical relevance. Therefore, I recommend accept.

I encourage the authors to make their exposition more transparent and correctly positioned with respect to prior work. In particular,
-  With the attention fixed, the problem reduces to matrix factorization with a rank-1 structure with factored parameterization versus full matrix parameterization. This is well-studied in the gradient descent on matrix factorization literature, and the authors should cite this line of work and be transparent about their own contributions. The authors should also clarify the importance of Transformers for their results, given that the core attention-component is frozen.
- Additionally, there has been work on understanding OOD performance based on implicit bias [1,2]. I encourage the authors to discuss these works.

[1] Abbe, E., Bengio, S., Cornacchia, E., Kleinberg, J., Lotfi, A., Raghu, M. and Zhang, C., 2022. Learning to reason with neural networks: Generalization, unseen data and boolean measures. Advances in Neural Information Processing Systems, 35, pp.2709-2722.
[2] Abbe, E., Bengio, S., Lotfi, A. and Rizk, K., 2024. Generalization on the unseen, logic reasoning and degree curriculum. Journal of Machine Learning Research, 25(331), pp.1-58.
Decision: Accept (poster)