Zhepeng Wang

Author directory

2026

Large language model (LLM) agents for long-horizon interactive tasks typically follow a ReAct-style protocol, issuing one primitive action per LLM round. While this enables frequent replanning, it is inefficient for long-horizon tasks where many rounds are spent on routine action sequences. A natural alternative is to let the agent emit variable-length action chunks. However, naively training such policies with standard reinforcement learning fails: the agent either collapses to single-action behavior or over-commits to excessively long sequences. Both failures share a common root cause: the inability to learn chunk boundaries. We propose SPACE, which addresses this challenge by distilling chunk-boundary supervision from trajectory-induced programmatic skills. We induce two-level programmatic skills from successful trajectories, where subskill boundaries serve as direct chunk-boundary supervision. This temporal structure is then distilled into a primitive-chunk policy via hybrid on-/off-policy optimization with chunk-aware credit assignment. Experiments on ALFWorld and ScienceWorld show that SPACE improves success rates by 7.0%-31.3% over the strongest baseline in each setting while reducing average LLM decision rounds by up to 78.9%.
Current evaluations of memorization in Large Language Models (LLMs) predominantly rely on overlap ratios, which fails to distinguish between sporadic token matches and reliable, sequential recall. What’s more, these methods become inapplicable in practical scenarios where ground truth suffixes is unavailable. To address these limitations, we propose a paradigm shift from quantifying "rates" to determining the Effective Memorization Length—the distinct boundary where a model transitions from faithful recall to open-ended generation. We introduce a Ensemble Memorization Classifier that leverages internal states from the LLM to estimate this boundary in a reference-free manner. This framework enables Confidence-Aware Generation Control, allowing to extract high-precision suffixes by dynamically truncating sequences precisely at the memorization boundary. Furthermore, we demonstrate that this boundary is optimizable; by integrating our classifier as a reward model within a reinforcement learning framework, we actively extend the average effective length of generated suffixes. Experiments on the GPT-Neo and Pythia families demonstrate that our approach not only accurately filters non-memorization tokens but also effectively "unlocks" deeper chains of memorized knowledge, transforming memorization from an unpredictable risk into a controllable utility.

2025

The evolution of pre-trained large language models (LLMs) has significantly transformed natural language processing. However, these advancements pose challenges, particularly the unintended memorization of training data, which raises ethical and privacy concerns. While prior research has largely focused on mitigating memorization or extracting memorized information, the deliberate control of memorization has been underexplored. This study addresses this gap by introducing a novel and unified gradient-based weight pruning framework to freely control memorization rates in LLMs. Our method enables fine-grained control over pruning parameters, allowing models to suppress or enhance memorization based on application-specific requirements. Experimental results demonstrate that our approach effectively balances the trade-offs between memorization and generalization, with an increase of up to 89.3% in Fractional ER suppression and 40.9% in Exact ER amplification compared to the original models.

2024

Pretrained large language models (LLMs) have excelled in a variety of natural language processing (NLP) tasks, including summarization, question answering, and translation. However, LLMs pose significant security risks due to their tendency to memorize training data, leading to potential privacy breaches and copyright infringement. Therefore, accurate measurement of the memorization is essential to evaluate and mitigate these potential risks. However, previous attempts to characterize memorization are constrained by either using prefixes only or by prepending a constant soft prompt to the prefixes, which cannot react to changes in input. To address this challenge, we propose a novel method for estimating LLM memorization using dynamic, prefix-dependent soft prompts. Our approach involves training a transformer-based generator to produce soft prompts that adapt to changes in input, thereby enabling more accurate extraction of memorized data. Our method not only addresses the limitations of previous methods but also demonstrates superior performance in diverse experimental settings compared to state-of-the-art techniques. In particular, our method can achieve the maximum relative improvement of 135.3% and 39.8% over the vanilla baseline on average in terms of discoverable memorization rate for the text generation task and code generation task, respectively. Our code is available at https://github.com/wangger/llm-memorization-dsp.