All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens

• URL: http://arxiv.org/abs/2509.09650v1
• Date: Thu, 11 Sep 2025 17:41:29 GMT
• Title: All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens
• Authors: Siddarth Mamidanna, Daking Rai, Ziyu Yao, Yilun Zhou,
• Abstract summary: In theory, the combination of causal self-attention and multilayer perceptron layers allows every token to access and compute information based on all preceding tokens.<n>We investigate the question in three steps: inhibiting input-specific token computations in the initial layers, restricting the routes of information transfer across token positions in the next few layers, and forcing all computation to happen at the last token in the remaining layers.
• Score: 14.890542559477906
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Large language models (LLMs) demonstrate proficiency across numerous computational tasks, yet their inner workings remain unclear. In theory, the combination of causal self-attention and multilayer perceptron layers allows every token to access and compute information based on all preceding tokens. In practice, to what extent are such operations present? In this paper, on mental math tasks (i.e., direct math calculation via next-token prediction without explicit reasoning), we investigate this question in three steps: inhibiting input-specific token computations in the initial layers, restricting the routes of information transfer across token positions in the next few layers, and forcing all computation to happen at the last token in the remaining layers. With two proposed techniques, Context-Aware Mean Ablation (CAMA) and Attention-Based Peeking (ABP), we identify an All-for-One subgraph (AF1) with high accuracy on a wide variety of mental math tasks, where meaningful computation occurs very late (in terms of layer depth) and only at the last token, which receives information of other tokens in few specific middle layers. Experiments on a variety of models and arithmetic expressions show that this subgraph is sufficient and necessary for high model performance, transfers across different models, and works on a variety of input styles. Ablations on different CAMA and ABP alternatives reveal their unique advantages over other methods, which may be of independent interest.

Related papers

• SCOPE: Saliency-Coverage Oriented Token Pruning for Efficient Multimodel LLMs [59.415473779171315]
  We propose a novel visual token pruning strategy called textbfSaliency-textbfCoverage textbfOriented token textbfPruning for textbfEfficient MLLMs.
  arXiv  Detail & Related papers  (2025-10-28T09:29:37Z)
• FrugalPrompt: Reducing Contextual Overhead in Large Language Models via Token Attribution [3.4666771782038652]
  Large language models (LLMs) owe much of their stellar performance to expansive input contexts, yet such verbosity inflates monetary costs, carbon footprint, and inference-time latency.<n>We introduce FrugalPrompt, a novel prompt compression framework for LLMs, which retains only the most semantically significant tokens.<n>We evaluate the approach across four NLP tasks: Sentiment Analysis, Commonsense QA, Summarization, and Mathematical Reasoning.
  arXiv  Detail & Related papers  (2025-10-18T10:22:13Z)
• Expanding Computation Spaces of LLMs at Inference Time [33.17624792878245]
  Chain-of-thought (CoT) rationale enables language models to use additional task-related text for problem-solving.<n>We investigate whether language models can leverage artificially inserted sequences of filler tokens solely at inference.
  arXiv  Detail & Related papers  (2025-09-29T14:59:44Z)
• Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning [53.57895922042783]
  Large Language Models (LLMs) excel at reasoning and planning when trained on chainof-thought (CoT) data.<n>We propose a hybrid representation of the reasoning process, where we partially abstract away the initial reasoning steps using latent discrete tokens.
  arXiv  Detail & Related papers  (2025-02-05T15:33:00Z)
• FIRP: Faster LLM inference via future intermediate representation prediction [54.897493351694195]
  FIRP generates multiple tokens instead of one at each decoding step. We conduct extensive experiments, showing a speedup ratio of 1.9x-3x in several models and datasets.
  arXiv  Detail & Related papers  (2024-10-27T15:53:49Z)
• Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers [54.20763128054692]
  We study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data. We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model.
  arXiv  Detail & Related papers  (2024-09-09T18:10:26Z)
• ToSA: Token Selective Attention for Efficient Vision Transformers [50.13756218204456]
  ToSA is a token selective attention approach that can identify tokens that need to be attended as well as those that can skip a transformer layer. We show that ToSA can significantly reduce computation costs while maintaining accuracy on the ImageNet classification benchmark.
  arXiv  Detail & Related papers  (2024-06-13T05:17:21Z)
• SEP: Self-Enhanced Prompt Tuning for Visual-Language Model [93.94454894142413]
  We introduce a novel approach named Self-Enhanced Prompt Tuning (SEP) SEP explicitly incorporates discriminative prior knowledge to enhance both textual-level and visual-level embeddings. Comprehensive evaluations across various benchmarks and tasks confirm SEP's efficacy in prompt tuning.
  arXiv  Detail & Related papers  (2024-05-24T13:35:56Z)
• Let's Think Dot by Dot: Hidden Computation in Transformer Language Models [30.972412126012884]
  Chain-of-thought responses from language models improve performance across most benchmarks. We show that transformers can use meaningless filler tokens in place of a chain of thought to solve two hard algorithmic tasks. We find that learning to use filler tokens is difficult and requires specific, dense supervision to converge.
  arXiv  Detail & Related papers  (2024-04-24T09:30:00Z)
• Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer [37.37547759817417]
  Transformer architecture has shown impressive performance in multiple research domains. We analyze its SGD training dynamics for the task of next token prediction. We prove that self-attention acts as a emphdiscriminative scanning algorithm.
  arXiv  Detail & Related papers  (2023-05-25T15:59:13Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.