Author name cluster

Dylan Zhang

Possible papers associated with this exact author name in Arrow. This page groups case-insensitive exact name matches and is not a full identity disambiguation profile.

3 papers

1 author row

TMLR Journal 2025 Journal Article

Entropy-Regularized Process Reward Model

Hanning Zhang
Pengcheng Wang
Shizhe Diao
Yong Lin
Rui Pan
Hanze Dong
Dylan Zhang
Pavlo Molchanov

Large language models (LLMs) have shown promise in performing complex multi-step reasoning, yet they continue to struggle with mathematical reasoning, often making systematic errors. A promising solution is reinforcement learning (RL) guided by reward models, particularly those focusing on process rewards, which score each intermediate step rather than solely evaluating the final outcome. This approach is more effective at guiding policy models towards correct reasoning trajectories. In this work, we propose an entropy-regularized process reward model (ER-PRM) that integrates KL-regularized Markov Decision Processes (MDP) to balance policy optimization with the need to prevent the policy from shifting too far from its initial distribution. We derive a novel reward construction method based on the theoretical results. Our theoretical analysis shows that we could derive the optimal reward model from the initial policy sampling. Our empirical experiments on the MATH and GSM8K benchmarks demonstrate that ER-PRM consistently outperforms existing process reward models, achieving 1% improvement on GSM8K and 2-3% improvement on MATH under best-of-N evaluation, and more than 1% improvement under RLHF. These results highlight the efficacy of entropy-regularization in enhancing LLMs' reasoning capabilities.

PDF Details

NeurIPS Conference 2025 Conference Paper

The Best Instruction-Tuning Data are Those That Fit

Dylan Zhang
Qirun Dai
Hao Peng

High-quality supervised finetuning (SFT) data are essential for unlocking pretrained LLMs’ capabilities. Typically, instructions are paired with responses from various sources—by human annotators or other LMs—which are often out of the distribution of the target model to be finetuned. At scale, this mismatch can lead to diminishing returns and even hurt model performance and robustness. We hypothesize that SFT is most effective when the data is aligned with the model’s pretrained distribution, and propose GRAPE —a novel SFT framework that tailors supervision to the target model. For each instruction, it g athers r esponses from various sources and selects the one that a ligns most closely to the model’s pre trained distribution, as measured by the normalized probability. Standard SFT is then performed on these selected responses. We first evaluate GRAPE in a controlled experiment, sampling multiple responses per question in the UltraInteract dataset from diverse models. We finetune using GRAPE-selected data on LMs from different families, including LLaMA-1-8B, Mistral-7B, and Qwen2. 5-7B. GRAPE significantly outperforms strong baselines—including distilling from the strongest model—with absolute gains up to 13. 8% averaged across benchmarks, and outperforms a 3× larger data baseline with improvements up to 17. 3%. GRAPE's benefits generalize to off-the-shelf SFT data. When used to subsample from the post-training data of Tulu3 and Olmo-2, GRAPE surpasses strong baselines trained on 4. 5× more data by 6. 1%, and outperforms state-of-the-art selection methods by 3. 9% on average. Notably, with only 1/3 the data and half the training epochs, GRAPE enables LLaMA-1-8B to exceed Tulu3-SFT performance by 3. 5%. Our findings highlight that aligning supervision with the pretrained distribution provides a simple yet powerful strategy to improve both the efficiency and effectiveness of SFT.

PDF Details

TMLR Journal 2024 Journal Article

Transformer-Based Models Are Not Yet Perfect At Learning to Emulate Structural Recursion

Dylan Zhang
Curt Tigges
Zory Zhang
Stella Biderman
Maxim Raginsky
Talia Ringer

This paper investigates the ability of transformer-based models to learn structural recursion from examples. Recursion is a universal concept in both natural and formal languages. Structural recursion is central to the programming language and formal mathematics tasks where symbolic tools currently excel beyond neural models, such as inferring semantic relations between datatypes and emulating program behavior. We introduce a general framework that nicely connects the abstract concepts of structural recursion in the programming language domain to concrete sequence modeling problems and learned models' behavior. The framework includes a representation that captures the general \textit{syntax} of structural recursion, coupled with two different frameworks for understanding their \textit{semantics}---one that is more natural from a programming languages perspective and one that helps bridge that perspective with a mechanistic understanding of the underlying transformer architecture. With our framework as a powerful conceptual tool, we identify different issues under various set-ups. The models trained to emulate recursive computations cannot fully capture the recursion yet instead fit short-cut algorithms and thus cannot solve certain edge cases that are under-represented in the training distribution. In addition, it is difficult for state-of-the-art large language models (LLMs) to mine recursive rules from in-context demonstrations. Meanwhile, these LLMs fail in interesting ways when emulating reduction (step-wise computation) of the recursive function.

PDF Details