Abstract
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.
Data availability
The data that support the findings of this study are available from the corresponding author upon reasonable request.
References
He, J., Xiao, S., Huang, S., Li, J. & Wang, D. Mokan: A multi-omics data analysis framework using kolmogorov-arnold networks. IEEE J. Biomed. Health Inform, pp 1-11 (2025).
Achiam, J. et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 (2023).
Anil, R. et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403 (2023).
Wang, Q., Wang, Z., Su, Y., Tong, H. & Song, Y. Rethinking the bounds of llm reasoning: Are multi-agent discussions the key? arXiv preprint arXiv:2402.18272 (2024).
Feng, G. et al. Towards revealing the mystery behind chain of thought: a theoretical perspective. Adv. Neural Inf. Process. Syst. 36, 70757–70798 (2023).
Wei, J. et al. Chain-of-thought prompting elicits reasoning in large language models. Adv. Neural Inf. Process. Syst. 35, 24824–24837 (2022).
Yao, S. et al. Tree of thoughts: Deliberate problem solving with large language models. Adv. Neural Inf. Process. Syst. 36, 11809–11822 (2023).
Besta, M. et al. Graph of thoughts: Solving elaborate problems with large language models. Proc. AAAI Conf. Artif. Intell. 38, 17682–17690 (2024).
Mishra, V. et al. Investigating the shortcomings of llms in step-by-step legal reasoning. arXiv preprint arXiv:2502.05675 (2025).
Blair-Stanek, A., Holzenberger, N. & Van Durme, B. Can gpt-3 perform statutory reasoning? In: Proceedings of the Nineteenth International Conference on Artificial Intelligence and Law, pp. 22–31 (2023).
Qiao, S. et al. Reasoning with language model prompting: A survey. arXiv preprint arXiv:2212.09597 (2022).
Benfenati, A. & Marta, A. A singular riemannian geometry approach to deep neural networks i. theoretical foundations. Neural Netw. 158, 331–343. https://doi.org/10.1016/j.neunet.2022.11.022 (2023).
Gruffaz, S. & Sassen, J. Riemannian metric learning: Closer to you than you imagine. arXiv preprint arXiv:2503.05321 (2025).
McDuff, D. et al. Towards accurate differential diagnosis with large language models. Nature 642, 451–457 (2025).
Lu, J. et al. Where views meet curves: Virtual anchors for hyperbolic multi-view graph diffusion. In: Proceedings of the 33rd ACM International Conference on Multimedia, pp. 2131–2140 (2025).
Vaswani, A. et al. Attention is all you need. Adv. Neural Inf. Process. Syst. 30, (2017).
Beltagy, I., Peters, M.E. & Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 (2020).
Gao, L., Madaan, A., Zhou, S., Alon, U., Liu, P., Yang, Y., Callan, J. & Neubig, G. Pal: Program-aided language models. In: International Conference on Machine Learning, pp. 10764–10799 (PMLR, 2023).
Guu, K., Lee, K., Tung, Z., Pasupat, P. & Chang, M. Retrieval augmented language model pre-training. In: International Conference on Machine Learning, pp. 3929–3938 (PMLR, 2020).
Amizadeh, S., Palangi, H., Polozov, A., Huang, Y. & Koishida, K. Neuro-symbolic visual reasoning: Disentangling. In: International Conference on Machine Learning, pp. 279–290 (PMLR, 2020).
Ganea, O., Bécigneul, G. & Hofmann, T. Hyperbolic neural networks. Adv. Neural Inf. Process. Syst. 31, (2018).
Nickel, M. & Kiela, D. Poincaré embeddings for learning hierarchical representations. Adv. Neural Inf. Process. Syst. 30, (2017).
Chami, I., Abu-El-Haija, S., Perozzi, B., Ré, C. & Murphy, K. Machine learning on graphs: A model and comprehensive taxonomy. J. Mach. Learn. Res. 23(89), 1–64 (2022).
Pan, Y.-T., Chou, J.-L. & Wei, C.-S. Matt: A manifold attention network for eeg decoding. Adv. Neural Inf. Process. Syst. 35, 31116–31129 (2022).
McInnes, L., Healy, J. & Melville, J. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426 (2018).
Bécigneul, G. & Ganea, O.-E. Riemannian adaptive optimization methods. arXiv preprint arXiv:1810.00760 (2018).
Wang, H. et al. Manifold-based verbalizer space re-embedding for tuning-free prompt-based classification. Proc. AAAI Conf. Artif. Intell. 38, 19126–19134 (2024).
Sun, L. et al. Motif-aware riemannian graph neural network with generative-contrastive learning. Proc. AAAI Conf. Artif. Intell. 38, 9044–9052 (2024).
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M. & Goldstein, T. The Intrinsic Dimension of Images and Its Impact on Learning (2021). arXiv:2104.08894.
Dijkstra, E.W. A note on two problems in connexion with graphs. In: Edsger Wybe Dijkstra: His Life, Work, and Legacy, pp. 287–290 (2022).
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A. & Kadian, A. The Llama 3 Herd of Models arXiv:2407.21783 (2024).
OpenAI, Achiam, J., Adler, S., Agarwal, S. & Ahmad, L. GPT-4 Technical Report arXiv:2303.08774 (2024).
DeepSeek-AI, Guo, D., Yang, D., Zhang, H. & Song, J. DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning arXiv:2501.12948 (2025).
Funding
This research did not receive funding.
Author information
Authors and Affiliations
Contributions
All authors contributed to the study’s conception and design. Material preparation, data collection and analysis were performed by Runze Mao, Hui Xie, Zhengyuan Zhang, and Mengyao Yang. The first draft of the manuscript was written by Runze Mao and all authors commented on previous versions of the manuscript. As tutors, Prof. Changzhen Hu, and Prof. Shengjun Wei made contributions to reviewing and revising it for intellectual content. All authors read and approved the final manuscript.
Corresponding author
Ethics declarations
Competing interests
The authors declare no competing interests.
Additional information
Publisher’s note
Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.
Rights and permissions
Open Access This article is licensed under a Creative Commons Attribution 4.0 International License, which permits use, sharing, adaptation, distribution and reproduction in any medium or format, as long as you give appropriate credit to the original author(s) and the source, provide a link to the Creative Commons licence, and indicate if changes were made. The images or other third party material in this article are included in the article’s Creative Commons licence, unless indicated otherwise in a credit line to the material. If material is not included in the article’s Creative Commons licence and your intended use is not permitted by statutory regulation or exceeds the permitted use, you will need to obtain permission directly from the copyright holder. To view a copy of this licence, visit http://creativecommons.org/licenses/by/4.0/.
About this article
Cite this article
Mao, R., Zhang, Z., Yang, M. et al. RiemannInfer: improving transformer inference through Riemannian geometry. Sci Rep (2026). https://doi.org/10.1038/s41598-026-37328-x
Received:
Accepted:
Published:
DOI: https://doi.org/10.1038/s41598-026-37328-x
