Low-Dimension-to-High-Dimension Generalization and Its Implications for Length Generalization

Thank you for your effort and insightful comments. We respond to your main concerns below. All experiments results are in the anonymous link: https://www.dropbox.com/scl/fi/52t23nfzev1lo1sq5dmyj/ICML_2025_5857_Rebuttal.pdf?rlkey=5nde6aampze744klvn0rsb3mp&st=37lg3nzo&dl=0.

1. The theoretical contributions of the paper are rather incremental.

Our main contribution is the LDHD generalization formulation (Abstraction 1), which offers a more precise characterization of length generalization by separating it into two challenges: (1) LDHD generalization in the latent space (which is inherent), and (2) data format nuisance. Under the formulation, we can derive the fundamental limitation of length generalization, explain the effectiveness of certain methods, and design novel position embeddings for challenging length generalization problem.

The theoretical results in Section 4 demonstrates how different model designs can introduce different inductive bias to handle LDHD generalization, only related to the LDHD generalization aspect in Abstraction 1.

While our Boolean analysis in Section 4 is technically inspired by the remarkable work of Abbe et al. (2023), our core novelty lies in Abstraction 1 and its implications, not in extending prior Boolean analysis.

2. The proposed method RPE-Square is interesting but is only evaluated in limited settings.

We evaluate RPE-Square in three additional tasks: Parity (with CoT), Multiplication (1 * N), and Division (N / 1). RPE-Square achieves length generalization in all the three tasks (Figures 3 - 5 in the link).

3. The theory and algorithm seem disconnected.

All components of our paper are built around Abstraction 1. Here’s how they connect: Section 3 focuses on the LDHD challenge and shows that length generalization requires prior knowledge. Section 4 illustrates how inductive bias can be introduced through model design to address the LDHD challenge. Section 5 consider practical length generalization techniques from the perspective of Abstraction 1. Section 5.1 explains the effectiveness of CoT as extending the latent space, which may promote LDHD generalization. Section 5.2 derive the position embedding design principle according to Abstraction 1, highlighting the importance of handling LDHD generalization in the latent space and the data format separately.

RPE-Square is designed closely following the principle implied by Abstraction 1. The inner "relative distance to some special token" is to handle the unaligned format. The outer "relative distance of relative distance to some special token" is to handle the LDHD generalization in the latent space, which is partially explained by the analysis of PLAA-GRPE. Further details are in Appendix D.1.

4. I did not understand Sec. 5.1: how did you determine the latent space of each instance? Why should the latent space take the form in the paper?

The latent space is determined according to Abstraction 1. We consider the scale $n$ of each instance and a proper $\Sigma$ such that each instance of scale $n$ corresponds to an element in $\Sigma^n$.

With CoT, we treat each step (which requires a model prediction) as an instance of the same scale. For example, in the addtion with CoT, the two predition steps, i.e., given $x_0 \dots x_{n-1} + y_0 \dots y_{n-1} =$ to predict $z_0$, and given $x_0 \dots x_{n-1} + y_0 \dots y_{n-1} = z_0$ to predict $z_1$, are treated as separate CoT-step instances of the same scale $n$.

We then decide a proper domain $\bar{\Sigma}$ such that each CoT-step instance of scale corresponds to an element in $\bar{\Sigma}^n$. Furthermore, since CoT typically inserts intermidate results to the original instances, $\bar{\Sigma}$ will take the form $\Sigma\times\Sigma'$. For the addition, $\Sigma'=\lbrace *,0,\dots,9\rbrace$ satisfies the requirement.

5. ...the formulation of LDHD generalization (Def. 1) is quite similar to length generalization... Are there any other applications of this formulation except length generalization?

Changing "length" to "dimension" is not merely a notion replacement. By the notion "dimension", we highlight the exponential growth in sample space as scale increases, and the complete absence of information in low-dimension training data about high-dimension behavior. These aspects are not captured by "length", which can be confounded by formatting. Thus, “dimension” more precisely reflects the scaling challenge.

While LDHD generalization is originally to characterize the inherent scaling challenge in length generalization, it could be extended to other applications to character the exponential growth in the sample space size and imperfect information provided by the training data. For example, in graph learning, we may consider the number of nodes as the dimension and use this formulation to describe the generalization from small to large graphs.

We are grateful for your careful evaluation and positive assessment of our work. All experiments results are in the anonymous link: https://www.dropbox.com/scl/fi/52t23nfzev1lo1sq5dmyj/ICML_2025_5857_Rebuttal.pdf?rlkey=5nde6aampze744klvn0rsb3mp&st=37lg3nzo&dl=0.

1. Has RPE-Square been evaluated in diverse or challenging environments? What were the challenges? What were the outcomes?

We evaluate RPE-Square in three additional tasks: Parity (with CoT), Multiplication (1 * N), and Division (N / 1). RPE-Square achieves length generalization in all the three tasks (Figures 3 - 5 in the link). These results are consistent with our insight of the inductive bias of RPE-Square and further justify the proposed position embedding design principle.

2. Could you provide an insight on the computational overhead of RPE-Square, especially in large-scale or real-time applications?

Suppose the sequence length is N and the hidden dimension is d. The computational overhead of an attention with RPE-Square is $O(N^4 + N^3 d)=O(N^2 (N^2 + N^d))$ (with the implementation that saves the token attention scores for all query j and key i and reuse them when computing $\text{RPE-Square}_{i,j}$). This overhead can lead to inefficiency for large-scale or real-time applications.

We note that RPE-Square is primarily intended to illustrate the design principle for position embeddings. Our focus is on the challenge of length generalization instead of long-sequence efficiency. We consider it a promising direction to develop more efficient variants of RPE-Square in future work.

3. Are there specific conditions or assumptions in Theorem 3 that might limit its general applicability? If so, what are they?

Theorem 3 is derived with a simplified model (PLAA-APE). The conclusion in Theorem 3 may not hold precisely for more complex models. Also, the target functions considered in Theorem 3 are restricted to Boolean functions and the notation "min-degree profiler w.r.t. linearly independent set" may not be applied to non-Boolean functions directly. Additionally, the result include a limit ($\hat{P}$) but we cannot use the exact limit in practice.

Despite the above theoretical limitations, we think Theorem 3 can provide sufficient insight on the inductive bias of APE. The simplified model captures the essential effect of APE on the inductive bias. The analysis in the Boolean functions can be naturally extended to functions defined on finite domains. The limit can be well approximately by choosing a sufficient small initialization.

We also appreciate your suggestions on writing to improve clarity and readability. We will revise our manuscript accordingly.

Thank you for your time and valuable feedback. We'd like to response to your concerns as below. All experiments results are in the anonymous link: https://www.dropbox.com/scl/fi/52t23nfzev1lo1sq5dmyj/ICML_2025_5857_Rebuttal.pdf?rlkey=5nde6aampze744klvn0rsb3mp&st=37lg3nzo&dl=0.

1. ...comparisons with alternative position embeddings designed for length generalization...

We conduct additional experiments to compare the mentioned position embeddings (Since Abacus and position coupling are similar apart from some subtle implementation details in our tasks, we only implement Abacus currently.) The results (Figures 1 - 2 in the link) show that RPE-Square achieves the best overall length generalization in the tasks.

2. Lack of Generalization Tests on Longer Sequences

We evaluate the models with RPE-Square on significantly longer sequences in three tasks (Parity (with CoT), Multiplication (1 * N), and Division (N * 1)) The results (Figure 7 in the link) show that the models can extrapolate to significantly longer sequences to some extend and the performance may decreases as the sequence length increases.

3. *Absence of Robustness Analysis: The performance of RPE-Square across various tasks beyond addition and unaligned copy is not explored.

We apply RPE-Square to three new tasks: Parity (With CoT), Mulplication (1 * N), and Division (N / 1). The results (Figures 3 - 5 in the link) show that Transformers equipped with RPE-Square can achieve length generalization in all the three tasks, showing the robustness of RPE-Square across different tasks.

4. Experiments comparing with and without CoT seems to be missing.

We compare the length generalization of Parity with CoT and without CoT. We train Transformers with RPE-Square. We only achieve length generalization for Parity with CoT (Figure 8 in the link). This is because as CoT enables LDHD generalization via relative distances in the latent space. There are also many empirical results showing CoT can enhance length generalization in previous works, e.g., [1] [2].

[1] Zhou, H., et al., 2023. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16028.

[2] Feng, G., et al., 2023. Towards revealing the mystery behind chain of thought: a theoretical perspective. NeurIPS.

5. I find it hard to make connections between sections and get a clear intuition behind the approach.

(a) Connections between sections.

Section 1 introduces the problem of length generalization and proposes LDHD generalization to address the mismatch between input length and problem scale. Based on the LDHD formulation, Section 3 presents the No-Free-Lunch Theorem for LDHD generalization, motivating the need for inductive bias. Section 4 consider how different inductive bias can be incorporated by choosing different models, a technical problem to address the LDHD generalization challenge in Section 3. In Section 5, we further discuss the implications of the LDHD generalization formulation for practical length generalization. We explain why CoT can promote length generalization. We also propose the position embeddding design principle for length generalization: consider the inherent LDHD generalization and the data format nuisance separately.

(b) Intuition behind the approach.

• The intuition behind LDHD generalization formulation includes: (1) The sample space grows exponentially as the scale increases and the small-scale instances do not tell how large-scale instances can be solved without external information; (2) In language modeling, the sequence length can be affected by the data format and may not faithfully reflect the problem scale.

• The intuition behind RPE-Square: the unaligned formats can be handled by the relative distances to some special tokens. Detailed illustrations are in Appendix D.1.

6. the definition of "concept c" in Abstraction 1 seems missing.

A "concept" refers to a target mapping, as in [1]. We'll clarify this in the revision.

[1] Mohri, Mehryar, 2018. Foundations of machine learning, Chapter 2, The PAC Learning Framework, Section 2.1, The PAC learning model.

5. Generalization of theoretical results

We provide empirical support (Figure 9) by training GPT2 with APE and RPE on different functions. Probing reveals the learned models roughly align with minimum-degree interpolators w.r.t. their respective bases.

6. ... the computational and memory overhead ...

Suppose the sequence length is N and the hidden dimension is d. RPE-Square: $O(N^4 + N^3 d)$ compute and $O(N^2 + N d)$ memory.

Standard PE: $O(N^2 d)$ compute and $O(N^2 + N d)$ memory.

While RPE-Square is less efficient than the standard position embedding, it mainly serves as a proof-of-concept for separating LDHD and format handling. It is focused on the length generalization challenge instead of the challenge of long sequence. Designing efficient variants is a promising direction for future work.

Thank you for your thoughtful and constructive feedback. Below, we respond to your concerns. All experiments results are in the anonymous link: https://www.dropbox.com/scl/fi/52t23nfzev1lo1sq5dmyj/ICML_2025_5857_Rebuttal.pdf?rlkey=5nde6aampze744klvn0rsb3mp&st=37lg3nzo&dl=0.

1. ... However, the paper did not justify how their choices do not compromise the generality of their analysis.

(a) We isolate the positional contribution in attention to highlight how position embeddings affect inductive bias and length generalization. This helps us better understand their standalone impact, and the empirical results support this analysis.

(b) We assume an identity value head because it is not central to the question of length generalization. Intuitively, we expect that the interpolator learns the correct value head to perform in-distribution generalization. So when focusing on length generalization, it is reasonable to assume a correct value head and analyze how the model handles larger-scale instances. This simplification helps clarify the role of attention and position embeddings.

(c) While our analysis focuses on binary functions, the core insights naturally extend to tasks over finite alphabets, because we can consider the binary representations. For example, for an alphabet of size $N$, each symbol can be encoded using $log N$ bits, and the analysis for binary functions can be applied similarly. Though not identical, the notion of the "minimum-degree interpolator over a linearly independent set" continues to hold in spirit under binary encoding schemes. Thus, the results we derive for binary settings are relevant to broader cases like language modeling.

Thank you for pointing out this oversight and we will include the justification in our manuscript.

2. ... It is expected that the paper at least present how their approach extends to other tasks, and what kind of domain knowledge can be exploited for different problems.

We conducted additional experiments to address this concern:

(a) We apply RPE-Square to three new tasks: Parity (With CoT), Mulplication (1 * N), and Division (N / 1), and compared its performance against more baselines (RPE, RoPE, NoPE, ALiBi, Abacus). The results (Figures 3 - 5 in the link) show that RPE-Square can also achieve length generalization in these tasks, demonstrating its applicability beyond the initial two tasks.

(b) We also consider a task called AdditionMod10, which requires different domain knowledge from what RPE-Square captures. AdditionMod10 computes the sum of the addends modulo 10. The result depends only on the first digits of the addends (due to the modulo operation), so positional embedding needs to capture the absolute value of the relative distances to some special tokens ([BOS] and "+"). Following the same design principle, we consider a new position embedding called RPE-Absolute to encode "the absolute value of the relative distances to some special tokens". The experiment result (Figure 6 in the link) shows that RPE-Absolute achieves length generalization in AdditionMod10. This suggests that our position embedding design approach is adaptable to different domain knowledge.

3. Regarding OOD and LDHD. I believe the paper's analysis do not explain why we observe (surprising) LDHD generalization for several other problems without any special domain knowledge...

(a) Domain knowledge can be implicitly encoded via architecture design, training algorithms, data distributions, preprocessing strategies, etc. Even if domain knowledge is not injected explicitly, these prior choices reflect inductive biases tailored to the task.

(b) "...despite <10% representation of non-english data in pretraining"

It is important to distinguish low representation from zero representation. Prior work [1] shows that adding even a very small amount of long sequences to the training set can dramatically improve length generalization, which is called priming. From our perspective, this is because the presence of long sequences, even in small quantities, prevents the model from learning an overly simple interpolator that fails to extrapolate. For example, in digit-wise addition, if the training set only includes 3-digit numbers, the model might learn to ignore digits beyond the third. However, including a few 5-digit samples forces the model to also consider later digits, improving its extrapolation ability.

(c) While this work is mainly focused on length generalization, it is an interesting future direction to study how the compositional generalization can be understood from the perspective of LDHD generalization.

[1] Jelassi, S. et al., 2023. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306.15400.

4. The writing of the paper needs much improvement...

We will carefully proofread our submission in the revision.