Yefan Wang

Author directory

2026

Current large language models perform next-token prediction under the implicit assumption that all tokens lie in a shared Euclidean embedding space, which treats every token uniformly. However, natural language contains rich structural relationships, including hierarchical, compositional, and cyclic patterns, which the flat geometry of Euclidean space cannot model well. To address this limitation, we propose rLLM (Riemannian Large Language Modeling), which introduces a token-adaptive positive-definite diagonal metric tensor into the LM head. This metric dynamically adjusts the local quadratic score for each hidden state and its context, allowing the model to better capture the relationships among tokens and prompts, with less than 0.01% additional parameters. We train rLLM from scratch across three different scales (2.8B to 5.0B parameters), covering the full pretraining, SFT, and RL pipeline. Extensive experiments show that rLLM improves Base CORE, a composite of 22 in-context learning tasks, by up to +2.5pp at the two smaller GPT scales and obtains higher ChatCORE at all three scales, with gains up to +3.0pp after RL. These results suggest that revisiting the geometric assumptions underlying token representations is a promising direction for building more expressive and capable language models. Code is available at https://github.com/digailab/riemannianllm.