Inner Thinking Transformer: Leveraging Dynamic Depth Scaling to Foster Adaptive Internal Thinking

2025-02-19

AI Summary

Large language models (LLMs) struggle to perform deep reasoning over critical tokens under parameter constraints, and standard Transformers often suffer from gradient instability during complex reasoning, exposing architectural limitations in depth scalability. Method: We propose an implicit thinking modeling framework that decouples Transformer layer computation into dynamically expandable implicit “thinking steps.” This is achieved via adaptive token routing, residual thinking connections, and step-aware encoding—enabling elastic, depth-varying computation without increasing model parameters. The approach optimizes inference paths in an architecture-aware manner, requiring no additional parameters. Contribution/Results: Evaluated on 162M–466M models, our method achieves 96.5% of the performance of a 466M baseline while reducing training data by 43.2%. It outperforms standard Transformers and Loop-based variants across 11 benchmarks, significantly enhancing both reasoning efficiency and representational capacity of parameter-constrained models.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large language models (LLMs) face inherent performance bottlenecks under parameter constraints, particularly in processing critical tokens that demand complex reasoning. Empirical analysis reveals challenging tokens induce abrupt gradient spikes across layers, exposing architectural stress points in standard Transformers. Building on this insight, we propose Inner Thinking Transformer (ITT), which reimagines layer computations as implicit thinking steps. ITT dynamically allocates computation through Adaptive Token Routing, iteratively refines representations via Residual Thinking Connections, and distinguishes reasoning phases using Thinking Step Encoding. ITT enables deeper processing of critical tokens without parameter expansion. Evaluations across 162M-466M parameter models show ITT achieves 96.5% performance of a 466M Transformer using only 162M parameters, reduces training data by 43.2%, and outperforms Transformer/Loop variants in 11 benchmarks. By enabling elastic computation allocation during inference, ITT balances performance and efficiency through architecture-aware optimization of implicit thinking pathways.

Problem

• Improve LLM performance under parameter constraints
• Address gradient spikes in Transformer architectures
• Enhance critical token processing efficiency

Innovation

• Dynamic Depth Scaling
• Adaptive Token Routing
• Residual Thinking Connections