PENCIL: Long Thoughts with Short Memory

Weakness 1: "Some proofs and technical details are deferred to the appendix; a more self-contained presentation could aid clarity."

Thank you for the suggestion. We will incorporate key technical details from the Appendix into the main paper once we have an additional page in the final version.

Weakness 2 & Q1: "Could the authors provide further ablation studies to isolate the contribution of the reduction rule from other architectural modifications? How sensitive is the performance to the exact design of the special tokens and reduction triggers?"

We kindly note that in all our experiments, we have controlled for architectural and implementation choices by fixing the transformer architecture (including number of parameters, context window size, positional encoding, etc.) and the training algorithm (optimizer, batch size, learning rate, etc.). The only difference between CoT and PENCIL is in the data preprocessing (see the first three paragraphs of Section 4) and whether the reduction rule is applied during inference. This ensures the performance difference stems from the reduction rule.

As shown in Section 4, when all other factors are fixed, PENCIL consistently outperforms CoT, especially on larger problems. Moreover, the superior performance is insensitive to the specific choice of architectures. For example, in Figure 7, when we vary the model size (for both CoT and PENCIL) and context window size, our approach still achieves better performance. Empirically, the performance gap is also insensitive to the positional encoding (e.g. simple absolute PE works as well).

While the precise design of the reduction rule might affect performance, we have not yet found an alternative rule that achieves better space and time efficiency than the one proposed (Equation 1). That being said, we conjecture if there does exist a better rule, PENCIL should achieve better performance and larger gap with CoT.

Q2: "While the paper shows impressive results on SAT, QBF, and Einstein’s puzzle, how does PENCIL perform on real-world tasks or benchmarks that involve natural language understanding beyond synthetic reasoning problems?"

It is indeed very promising to adapt PENCIL to general real-world tasks that involve natural language understanding. One potential way is to generate datasets with special tokens and fine-tune existing LLMs. See However, this requires considerate engineering efforts, and thus we leave this as future work. Also see a discussion of "how PENCIL can be applied to standard LLMs" in our response to Reviewer zuSb.

Weakness 3 & Q3: "A more extensive discussion comparing with other memory efficient transformer techniques such as state space models could be beneficial; Can the authors elaborate on how PENCIL compares with other recent memory compression or dynamic token pruning methods in terms of both efficiency and accuracy?"

Existing memory-efficient architectures, such as those employing linear complexity attention (e.g., [1]), still require the context length to grow with the running time for problem solving, and thus do not address the fundamental limitation of CoT that PENCIL overcomes. While state-space models (e.g. [2]) avoid this issue by storing only the state, but this often comes at the cost of reduced expressiveness; for instance, SSMs have been shown (empirically and theoretical) to struggle with even simple tasks like copying [3].

Other memory compression and token pruning methods (e.g., [4–7]) typically combine base models with external heuristic algorithms (such as those relying on score functions and ranking). As a result, the models do not have the capability to reduce the space themselves whereas PENCIL explicit trains the model to do so. Moreover, these methods do not have theoretical benefits of solving arbitrary problems with optimal space complexity that PENCIL offers.

PENCIL essentially differs from these lines of works by handling the limitation of the next-token generation (i.e. CoT) paradigm, and is orthogonal to the contributions of aforementioned papers. In other words, PENCIL is compatible with different base model choices and existing memory compression techniques. We will incorporate a more comprehensive discussion on related work in the next revision.

[1] Rethinking attention with performers, ICLR 2021

[2] Mamba: Linear-time sequence modeling with selective state spaces, 2023

[3] Repeat After Me: Transformers are Better than State Space Models at Copying, ICML 2024

[4] Efficient streaming language models with attention sinks, ICLR 2024

[5] H2o: Heavy-hitter oracle for efficient generative inference of large language models, NeurIPS 2023

[6] Model tells you what to discard: Adaptive kv cache compression for llms, ICLR 2024

[7] Llmlingua: Compressing prompts for accelerated inference of large language models, EMNLP 2023

Weakness 1 & Q1: "The three datasets used in the paper are all SAT-like ones, which lacks discussion on generalizability to more general reasoning. Besides, the reduced CoT growing from exponential to polynomial only works for the SAT-like tasks."

We choose SAT and QBF because they are representative NP-complete and PSPACE-complete problems that no existing algorithm can solve efficiently, and thus they are ideal for stress-testing a model's capability in handling hard reasoning tasks. Moreover, PENCIL can be extended to arbitrary tasks provided the underlying algorithms can be written in a pure functional style. See our response to Reviewer xx5n's Q2 where we have briefly described the high-level idea.

To work on general reasoning tasks, one potential approach is to fine-tune existing LLMs on datasets with special tokens (generated either manually or automatically), enabling the model to learn to reason in a structured manner and think longer to solve more complicated tasks. See detailed discussion of how PENCIL can be applied to standard LLMs in our response to Reviewer zuSb. We will elaborate on this in the next version and leave further extensions as future work.

Weakness 2 & Q2 & Concern 1 in Methods And Evaluation Criteria: "How to ensure that the reduced parts are really no longer necessary especially for more general reasoning tasks."

For our data generation process, the reduced parts are guaranteed to be unnecessary, because when generating the dataset, the special tokens are placed in the trace in strict accordance with how function calls are used in the (Python) codes that implement the algorithm for solving the task.

For general tasks whose structured reasoning trace is not explicitly included in the training set, there is no absolute guarantee that the model will always remove the unneeded parts. Nevertheless, this is a common limitation of all structured reasoning approaches (CoT, for example, does not always output useful intermediates steps that contribute to the final answer for general reasoning tasks, even models are trained to do so). Moreover, as we mentioned earlier, one potential way to mitigate this limitation is to fine-tune LLMs on appropriately formatted datasets, allowing them to learn to extract useful information from previous thoughts and use such a skill in general reasoning tasks.

Weakness 2 & Concern 2 in Methods And Evaluation Criteria: "The repeatedly LLM call may greatly increases computation complexity compared with generating CoT in one call."

It is important to note that PENCIL does not repeatedly call LLMs, nor does it increase computation complexity; instead, it significantly improves the efficiency compared with one-pass CoT.

Specifically, as has been discussed in the last paragraph of Section 2.2, for each application of reduction $**C** `[CALL]` **T** `[SEP]` **A** `[RETURN]` \Rightarrow **C****A**$ the same model is used to generate new tokens following the same sequence, rather than calling a new LLM and feeding the reduced sequence into it. The benefit of sticking to the same model and the sequence is that the KV cache of the context $**C**$ can be preserved, and only the KV cache for A should be recomputed, which incurs marginal cost, as reflected in Equation 8 that shows the minimal FLOPs (where KV cache usage is optimized) needed for PENCIL. Intuitively, PENCIL significantly saves computation because for each generated token the prefix length is significantly smaller than its CoT counterpart, while the total number of generated tokens is the same. Empirically, Figure 6 demonstrates that PENCIL is significantly more computationally efficient than direct CoT for each problem instance.

Weakness 3: "It would be better to simplify the introduction, such as adding some examples and explanations."

Thanks for the suggestion. We will improve the clarity of introduction in the next revision.

Weakness 4 & Q3: "The context window of current LLMs could be extended to quite large, such as 1M. So I wonder whether it’s necessary to remove unnecessary parts from CoT."

Indeed, there are many existing efforts trying to extend the context window of LLMs, see our response to Reviewer yxYs's Q3 for a detailed discussion and how PENCIL fundamentally differs from them. Our contribution is orthogonal to their efforts contributions of aforementioned papers, i.e. one can always combine PENCIL with larger context window to potentially even achieve better performance.

Moreover, it has been argued that enlarging the context window size can introduce issues such as diminished ability to retrieve relevant information from a very long context (see, e.g. [1]), whereas PENCIL is immune to such an issue by completely eliminating the unneeded thoughts.

[1] Liu, Nelson F., et al. "Lost in the middle: How language models use long contexts." ACL 2024.

Weakness 1: "The data generation process requires knowing the reasoning structure. How reasoning problems without a clear structure (e.g., math problems) can benefit from this approach remains a question."

Indeed, language models do not inherently reason in a way that allows for convenient space reduction. To extend PENCIL to general reasoning problems, one potential approach is to fine-tune LLMs on specialized datasets, either manually labeled or automatically generated, so that models learn to reason in a structured, memory-efficient manner (see our response to “How PENCIL can be applied to standard LLMs” for Reviewer zuSb).

Moreover, we want to kindly point out that many math problems actually do exhibit PENCIL-like structures. For example, lemmas consist of statements which needs to be remembered, and proofs which do not. Similarly, many math problems involve intermediate computations with a complex derivation which can be forgotten, leaving only a concise final expression that needs to be retained. This natural separation in mathematical reasoning directly aligns with the PENCIL approach.

Weakness 2: "Models used in this study are too small compared to the current LLMs. This could disadvantage standard CoT approaches and limit the generalizability of the paper's conclusions to larger models."

While we have not yet conducted experiments on real-world LLMs (which we leave as future work), we believe using larger models with larger context window would further amplify the advantages of PENCIL.

This is because, in terms of theoretical expressiveness, CoT requires the maximal context length to grow with the running time of a problem, whereas PENCIL only requires it to grow with the needed memory; the gap between time and space is significant (e.g. exponential) for inherently hard reasoning problems. Although this gap is less pronounced for smaller models on tasks, it would become much more significant on larger-scale problems.

Q1: "How trace rate is calculated?"

Let $x$ be the ground-truth reasoning trace, and $\hat x$ be the reasoning trace generated by the model (as defined in Equation 8). The trace rate is defined as $\frac{1}{\max\{|x|, |\hat x|\}}\sum_{i=1}^{\min\{|x|, |\hat x|\}} \mathbf 1(x_i =\hat x_i)$, which quantifies the percentage of correctly predicted tokens. We choose this metric because it is both a direct measure of sequence similarity and tractable even for very long traces.

Q2: "How was the training data for Einstein's Puzzle generated? Additionally, how did you transform the algorithm's solutions into text? Did you use templates to verbalize the solutions?"

We implemented the Einstein’s Puzzle solver in Python, and as the code runs, it uses templates to verbalize the key steps. For example, when removing an entry from all possibilities, the code generates thoughts like:

Since green must be immediately to the right of Birds, we remove “green” from House #1 (it can’t be in the leftmost position if it’s supposed to be on the right of something else)"

Special tokens are appended automatically. The general rule is as follows: when the code calls a new function, it appends the "[CALL]" token to the trace, and when the function finishes all computations and is ready to return, it appends "[SEP] A [RETURN]" where A is the returned value. (And if the returned value is from another function, we use the tail recursion in Equation 10 to further optimize the space.)

In fact, this method can be generalized to any algorithm written in a pure functional style; we will detail this further in the next version and open-source our code once published.

Q3: "What is the complexity of the problem in the training data? Are all test problems in-domain in terms of complexity, or are some of them OOD?"

The complexity of the problems in the training data matches that of the test problems (i.e., the same n). We plan to explore extensions to the OOD setting as future work.

Q4: "Could the model size and data generation process explain CoT's poor performance?"

It is possible that both the model size and the data generation format affect CoT’s performance. As shown in Figure 7, increasing the model size and context length improves CoT performance. And if we keep increasing the size, presumably CoT will eventually also solve $5\times 5$ or more complicated puzzles, but PENCIL will be able to go even further. That is, regardless of the model size and context window, we have shown PENCIL can consistently solve larger-scale problems than the problems CoT can solve with that model size. We will add [1] in the next version.

[1] Shah, Kulin, et al. “Causal language modeling can elicit search and reasoning capabilities on logic puzzles.” arXiv:2409.10502 (2024).

Q5: "Do you think this method can applied directly for fine-tuning LLMs?"

Yes, we believe that fine-tuning LLMs on structured datasets generated as described is a promising future direction.

How PENCIL can be applied to standard LLMs

The way we envision applying PENCIL to standard LLMs is to fine-tune LLMs on examples that include our special tokens, with the goal that model learns to reason in a structured manner that leverages memory efficiently and enables longer reasoning. Such datasets can be generated by domain experts (which is arduous but fairly common for LLM alignment), possibly using automatic anotation of common structures in structured mathematical or logical writting. One can also generate such datasets automatically by converting the running trace of any algorithms written in a pure functional programming language into a training set (see our response to Reviewer xx5n's Q2 for how this can be done for Einstein's puzzle); correspondingly, the maximal context length for PENCIL is proportional to the maximal stack memory (typically much smaller than the running time). We will discuss in detail how to generate such datasets for general algorithms in the next version of the paper.

While we have not yet applied PENCIL to standard LLMs (this requires significant resources, both for data collection and training), we nonetheless have provided empirical evidence by showing even small language models can learn to generate structured CoT. We will more comprehensively discuss potential ways to apply PENCIL to real-world LLMs and leave this as future work.

Q1: Effectiveness of Thought Reduction

In principle, PENCIL can effectively reduce the reasoning trace for any problem whose space complexity is smaller than its time complexity, as we formally prove in Section 5. While the time and space gap is typically significant, there does exist cases where all computations are useful and should be preserved, as the reviewer suggested; it is true in this case the reduction is not necessarily useful. We appreciate the reviewer’s point and will include a discussion in the next version.

Q2: Expressiveness of PENCIL in Relation to Finite-Precision Transformers

Yes, Theorem 5.4 can be extended to cases where the base model is a transformer. Particularly, we can prove that PENCIL with a fixed finite-size decoder-only transformer can perform universal space-efficient computation, by simulating Turing machine running in $T$ steps and $S$ space with $\mathcal O(T)$ generated tokens and maximal sequence length $\mathcal O(S)$. This is a significant improvement over standard CoT, which requires context length to grow proportionally with $\mathcal O(T)$ (i.e. results in [1]). An immediate implication is that PENCIL can solve all $\mathrm{PSPACE}$ problems using polynomial context-length, whereas CoT can only solve problems in $\mathrm{P}$.

The high-level idea of the proof is as follows: for any Turing machine $\mathcal M$, we can construct a transformer (under some architectural specifications such as average-hard attention used in [1]) that can executed Algorithm 1 during its forward pass. Specifically, the transformer should be able to: (1) simulate the step-by-step execution of the Turing machine, (2) detect when to generate the $`[SEP]`$ token indicating the start of summarization, and (3) summarize thoughts by computing a compressed state representation of the current token sequence. Although the construction is not straight-forward, we will include the proof in the paper once we can make updates.

[1] W. Merrill and A. Sabharwal. "The expresssive power of transformers with chain of thought." ICLR 2024.

Minors

$t_i$ ($i>0$) always exists for Turing machines whose running time is at least twice of the maximal memory, otherwise we have $\mathcal O(T(\mathcal M,x)) = \mathcal O(S(\mathcal M, s, x))$ and reduction is unnecessary as space optimality is automatically achieved. We will also correct other typos the reviewer pointed out.