# Extending Token Computation for LLM Reasoning

arXiv: https://arxiv.org/abs/2403.14932 · HTML: https://mlsanity.com/paper/2403.14932/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 19 of its 27 references that are in the map.
- First posted on arXiv 22 Mar 2024.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 19 of 19)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [Generative Language Modeling for Automated Theorem Proving](https://mlsanity.com/paper/2009.03393/) (2020)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Fine-grained Analysis of Sentence Embeddings Using Auxiliary Prediction Tasks](https://mlsanity.com/paper/1608.04207/) (2016)
- [AttentionViz: A Global View of Transformer Attention](https://mlsanity.com/paper/2305.03210/) (2023)
- [The First Place Solution of WSDM Cup 2024: Leveraging Large Language Models for Conversational Multi-Doc QA](https://mlsanity.com/paper/2402.18385/) (2024)

## Built on by (most cited first, 3 of 3)

- [Gradient-guided Attention Map Editing: Towards Efficient Contextual Hallucination Mitigation](https://mlsanity.com/paper/2503.08963/) (2025)
- [Impacts of Continued Legal Pre-Training and IFT on LLMs' Latent Representations of Human-Defined Legal Concepts](https://mlsanity.com/paper/2410.12001/) (2024)
- [Automated Attention Pattern Discovery at Scale in Large Language Models](https://mlsanity.com/paper/2604.03764/) (2026)

## Abstract (arXiv)

Large Language Models (LLMs) are pivotal in advancing natural language processing but often struggle with complex reasoning tasks due to inefficient attention distributions. In this paper, we explore the effect of increased computed tokens on LLM performance and introduce a novel method for extending computed tokens in the Chain-of-Thought (CoT) process, utilizing attention mechanism optimization. By fine-tuning an LLM on a domain-specific, highly structured dataset, we analyze attention patterns across layers, identifying inefficiencies caused by non-semantic tokens with outlier high attention scores. To address this, we propose an algorithm that emulates early layer attention patterns across downstream layers to re-balance skewed attention distributions and enhance knowledge abstraction. Our findings demonstrate that our approach not only facilitates a deeper understanding of the internal dynamics of LLMs but also significantly improves their reasoning capabilities, particularly in non-STEM domains. Our study lays the groundwork for further innovations in LLM design, aiming to create more powerful, versatile, and responsible models capable of tackling a broad range of real-world applications.