RiemannInfer: improving transformer inference through Riemannian geometry

Understanding the reasoning process of Large Language Models(LLMs) is crucial for interpreting their decision-making mechanisms. However, existing methods primarily rely on conditional probability distributions within the model, resulting in computationally inefficient reasoning processes and difficulties in constructing rigorous logical chains. This paper presents RiemannInfer, an innovative computational framework based on Riemannian geometry for optimizing LLM reasoning paths. Our approach stems from a key observation: the hidden states formed through attention mechanisms and contextual encoding in the Transformer can be represented as high-dimensional vector spaces that encapsulate global dependencies captured by the model. Building on this insight, we first apply topology-preserving dimensionality reduction to these hidden states, then construct a Riemannian manifold structure utilizing attention distribution features. Notably, We discover a close relationship between the attention mechanism to the geodesics and the curvature of Riemannian manifolds. Based on this foundation, we implement efficient inference work calculation through geodesic and curvature to realize reasoning path planning, which not only significantly enhances inference efficiency but also provides geometric interpretations of model decision processes. Extensive experiments on various models, including LLaMA, GPT-4, and DeepSeek, demonstrate that RiemannInfer significantly improves the reasoning accuracy of LLMs and exhibits excellent effectiveness and robustness. This research provides a novel geometric perspective for enhancing the interpretability and efficient reasoning of LLMs.