Jiahong Liu

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.
Long-term memory is a critical challenge for Large Language Model agents, as fixed context windows cannot preserve coherence across extended interactions. Existing memory systems represent conversation history as unstructured embedding vectors, retrieving information through semantic similarity. This paradigm fails to capture the associative structure of human memory, wherein related experiences progressively strengthen interconnections through repeated co-activation. Inspired by cognitive neuroscience, we identify three mechanisms central to biological memory: association, consolidation, and spreading activation, which remain largely absent in current research. To bridge this gap, we propose HeLa-Mem, a bio-inspired memory architecture that models memory as a dynamic graph with Hebbian learning dynamics. HeLa-Mem employs a dual-level organization: (1) an episodic memory graph that evolves through co-activation patterns, and (2) a semantic memory store populated via Hebbian Distillation, wherein a Reflective Agent identifies densely connected memory hubs and distills them into structured, reusable semantic knowledge. This dual-path design leverages both semantic similarity and learned associations, mirroring the episodic-semantic distinction in human cognition. Experiments on LoCoMo demonstrate superior performance across four question categories while using significantly fewer context tokens. Code is available on GitHub: https://github.com/ReinerBRO/HeLa-Mem