Menglin Yang

Author directory

2026

Latent chain-of-thought models move intermediate reasoning into continuous states, obscuring which internal object causally carries a computation. We introduce SCIT, the Suffix Cache Interchange Test, which patches exact source–recipient counterfactuals and crosses cache-segment sufficiency with K/V splits, hidden-state and semantic-source controls, decoded validation, and matched corruption. In GPT-2 arithmetic, counterfactual transfer concentrates in value-cache suffix trajectories rather than hidden states, keys, answer slots, or single-token triggers. Full sufficiency-and-necessity closure holds only for the main CODI-GPT2 checkpoint; the Sim-CoT-style reproduction supports sufficiency and decoded transfer but not necessity. Across broader competent cells, arithmetic-like GPT-2/1B tasks retain latent-tail value/KV transfer, whereas 8B and repaired non-arithmetic tasks shift to prompt-prefix or full-cache K/V; ungated cells receive no call. SCIT thus provides a cache-level diagnostic, a checkpoint-specific arithmetic mechanism, and a competence-gated carrier map—not a universal latent-tail claim. Code: https://github.com/YIDING4869/scit-emnlp-2026
Long-form story generation requires models to preserve narrative consistency across extended contexts, yet existing prompting-based methods often accumulate temporal, factual, character, commonsense, and stylistic errors as the story grows. We propose ConWriter, a training-free framework for consistency-aware long-form story generation. ConWriter writes stories incrementally at the scene level, guided by static story requirements, dynamic narrative memory, symbolic state reasoning, and uncertainty-aware risk signals. Rather than treating long-story generation as a single free-form decoding process, ConWriter maintains evolving story states, checks whether new scenes satisfy required narrative transitions, and uses uncertainty-aware risk signals to prioritize validation and localized repair. This enables consistency control during generation, before local errors propagate into later scenes. We evaluate ConWriter on ConStory-Bench across four long-story tasks, three target lengths, and multiple base LLMs following the official evaluation protocol. Across models and story lengths, ConWriter consistently matches or improves upon direct generation and outperforms the recent training-free baseline DOME in narrative consistency. These results demonstrate the effectiveness of lightweight neuro-symbolic consistency control for training-free long-form story generation. Code is available on GitHub.
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.
Chain-of-thought (CoT) reasoning improves the reasoning ability of large language models by introducing intermediate computation, but explicit rationales increase decoding length, latency, and context cost. Implicit CoT offers a more efficient alternative by moving intermediate reasoning into continuous latent states. However, latent reasoning can be unstable: successive latent states may become overly similar and collapse toward a shared dominant direction, reducing the diversity of the reasoning trajectory. In this work, we identify latent rank collapse and propose WIDER, a lightweight spectral regularizer for implicit CoT. During training, WIDER estimates the shared direction of each latent trajectory and penalizes projections onto this direction, encouraging latent states to span a broader representational subspace. The method is plug-and-play and leaves the backbone model, latent schedule, and inference-time decoding procedure unchanged. We further formulate this collapse as a geometric bottleneck in implicit reasoning, casting its mitigation as a training-time regularization problem rather than an inference-time decoding change. Extensive experiments show that WIDER improves matched implicit CoT baselines, while mechanistic analyses reveal higher effective rank, lower dominant-direction energy, and reduced redundancy among latent steps. These results highlight latent subspace utilization as an important factor for efficient continuous reasoning, providing a geometric perspective for analyzing and improving implicit CoT. Code is available at https://github.com/whitesweater/WIDER.
Parametric retrieval-augmented generation (PRAG) injects retrieved evidence into a large language model (LLM) through passage-specific LoRA adapters, reducing reliance on long in-context prompts. When multiple passages are retrieved for the same query, however, evidence-level fusion becomes a bottleneck: equal-weight merging can amplify weak or conflicting evidence, and translating retrieval signals into fusion weights often requires fragile global tuning. We propose FCPRAG, a fusion-controlled parametric RAG framework that adds a lightweight controller for retrieval-conditioned, sample-level adapter fusion. The controller predicts per-passage fusion scores together with sample-level calibration signals, including a mixing gate and an adaptive temperature, enabling fusion that stays selective under informative retrieval signals and conservative under uncertainty. FCPRAG is trained with merge-aware supervision derived from each adapter’s marginal contribution within a multi-adapter merge, using training data only. We further show that a single dataset-level temperature is suboptimal under heteroscedastic retrieval uncertainty, motivating sample-level adaptation. Experiments on HotpotQA, 2WikiMultiHopQA, PopQA, and ComplexWebQuestions (CWQ) across three LLM backbones show that FCPRAG consistently improves F1 over standard RAG and parametric RAG baselines, with gains of up to 4.65% on 2WikiMultiHopQA and 7.55% on CWQ, while also reducing tuning cost and improving robustness under retrieval perturbations. Code is available at https://github.com/ReinerBRO/FCPRAG.
Semantic ID-based generative retrieval represents each item with a short sequence of discrete tokens derived from language-based item representations. Existing methods mainly instantiate this idea in two paradigms: autoregressive SID generation and parallel SID generation. In both paradigms, the quality of Semantic IDs depends critically on how continuous item embeddings are discretized before they become generation targets. While existing work has improved SID tokenization through better quantization objectives and supervision signals, it often treats the pre-tokenization embedding space as fixed, overlooking that pretrained item embeddings can be anisotropic and dominated by a few directions. This creates a pre-quantization distribution mismatch, where the global embedding geometry may not align with the structured spaces used for code assignment. We propose GeoSID, a tokenizer-side geometric shaping framework that keeps the downstream generative recommender unchanged. GeoSID uses a lightweight projected distribution matching regularizer to adjust item representations before SID assignment, and places the shaping objective in tokenizer-specific quantization spaces, including residual levels for residual quantization and OPQ/PQ blocks for product quantization. On the RPG backbone, GeoSID achieves up to 8.13% relative improvement in Recall@10 and 6.91% in NDCG@10. Further analyses show that GeoSID reduces projected distribution discrepancy, improves code usage balance, and lowers prefix-level SID collisions. Code is available at https://github.com/hexinrong122/GeoSID.
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
Large language models (LLMs) increasingly solve complex reasoning tasks via long chain-of-thought, but their forward-only autoregressive generation process is fragile; early token errors can cascade, which creates a clear need for self-reflection mechanisms. However, existing self-reflection either performs revisions over full drafts or learns self-correction via expensive training, both fundamentally reactive and inefficient. To address this, we propose Self-Reflective Generation at Test Time (SRGen), a lightweight test-time framework that reflects before generating at uncertain points. During token generation, SRGen utilizes dynamic entropy thresholding to identify high-uncertainty tokens. For each identified token, it trains a specific corrective vector, which fully exploits the already generated context for a self-reflective generation to correct the token probability distribution. By retrospectively analyzing the partial output, this self-reflection enables more trustworthy decisions, thereby significantly reducing the probability of errors at highly uncertain points. Evaluated on challenging mathematical reasoning benchmarks and a diverse set of LLMs, SRGen can significantly strengthen model reasoning. Moreover, our findings position SRGen as a plug-and-play method that integrates reflection into the generation process for reliable LLM reasoning, achieving consistent gains with bounded overhead and can be combined with other training-time (e.g., RLHF) and test-time (e.g., SLOT) techniques.