← Search

Shujian Zhang

24 accepted papers

2026

DPO Unchained: Your Training Algorithm is Secretly Disentangled in Human Choice Theory (and Its Loss' Convexity is Dispensable)

ICML 2026oral

Normative theories allow one to elicit key parts of a ML algorithm from first principles, which is crucial at a time of championed scrutiny for ML work. Direct Preference Optimization (DPO) cleverly bypasses reward modeling by making an explicit link with a specific normative model of human choice. …

Cited by 0SourceScholar
2026

Fantastic Reasoning Behaviors and Where to Find Them: Unsupervised Discovery of the Reasoning Process

ICML 2026poster

Despite the growing reasoning capabilities of recent large language models (LLMs), their internal mechanisms during the reasoning process remain underexplored. Prior approaches often rely on human-defined concepts (e.g., overthinking, reflection) at the word level to analyze reasoning in a supervise…

Cited by 0SourceScholar
2025

Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

ICLR 2025poster

Large Language Models (LLMs) are susceptible to security and safety threats, such as prompt injection, prompt extraction, and harmful requests. One major cause of these vulnerabilities is the lack of an instruction hierarchy. Modern LLM architectures treat all inputs equally, failing to distinguish…

Cited by 6SourcePDFScholar
2025

Score Forgetting Distillation: A Swift, Data-Free Method for Machine Unlearning in Diffusion Models

ICLR 2025poster

The machine learning community is increasingly recognizing the importance of fostering trust and safety in modern generative AI (GenAI) models. We posit machine unlearning (MU) as a crucial foundation for developing safe, secure, and trustworthy GenAI models. Traditional MU methods often rely on str…

2025

Statistical Advantages of Perturbing Cosine Router in Mixture of Experts

ICLR 2025poster

The cosine router in Mixture of Experts (MoE) has recently emerged as an attractive alternative to the conventional linear router. Indeed, the cosine router demonstrates favorable performance in image and language tasks and exhibits better ability to mitigate the representation collapse issue, which…

Cited by 6SourcePDFScholar
2025

T-REG: Preference Optimization with Token-Level Reward Regularization

ACL 2025long

Reinforcement Learning from Human Feedback (RLHF) has been pivotal in enabling Large Language Models (LLMs) to effectively follow instructions and produce meaningful alignment by leveraging human preference data. Traditionally, RLHF involves generating responses to a query and using a separate rewar…

2024

LanguageFlow: Advancing Diffusion Language Generation with Probabilistic Flows

NAACL 2024long

Recent works have demonstrated success in controlling sentence attributes (e.g., sentiment) and structure (e.g., syntactic structure) based on the diffusion language model. A key component that drives theimpressive performance for generating high-quality samples from noise is iteratively denoise for…

Cited by 2SourcePDFScholar
2024

Switchable Decision: Dynamic Neural Generation Networks

ICML 2024poster

Auto-regressive generation models achieve competitive performance across many different NLP tasks such as summarization, question answering, and classifications. However, they are also known for being slow in inference, which makes them challenging to deploy in real-time applications. We propose a s…

Cited by 0SourcePDFScholar
2024

WPO: Enhancing RLHF with Weighted Preference Optimization

EMNLP 2024main

Reinforcement learning from human feedback (RLHF) is a promising solution to align large language models (LLMs) more closely with human values. Off-policy preference optimization, where the preference data is obtained from other models, is widely adopted due to its cost efficiency and scalability. H…

2023

Fantastic Rewards and How to Tame Them: A Case Study on Reward Learning for Task-oriented Dialogue Systems

ICLR 2023poster

When learning task-oriented dialogue (ToD) agents, reinforcement learning (RL) techniques can naturally be utilized to train dialogue strategies to achieve user-specific goals. Prior works mainly focus on adopting advanced RL techniques to train the ToD agents, while the design of the reward functio…

2023

FlowGrad: Controlling the Output of Generative ODEs With Gradients

CVPR 2023poster

Generative modeling with ordinary differential equations (ODEs) has achieved fantastic results on a variety of applications. Yet, few works have focused on controlling the generated content of a pre-trained ODE-based generative model. In this paper, we propose to optimize the output of ODE models ac…

2023

POUF: Prompt-Oriented Unsupervised Fine-tuning for Large Pre-trained Models

ICML 2023poster

Through prompting, large-scale pre-trained models have become more expressive and powerful, gaining significant attention in recent years. Though these big models have zero-shot capabilities, in general, labeled data are still required to adapt them to downstream tasks. To overcome this critical lim…

2023

Preference-grounded Token-level Guidance for Language Model Fine-tuning

NeurIPS 2023poster

Aligning language models (LMs) with preferences is an important problem in natural language generation. A key challenge is that preferences are typically provided at the *sequence level* while LM training and generation both occur at the *token level*. There is, therefore, a *granularity mismatch* b…

2022

A Unified Framework for Alternating Offline Model Training and Policy Learning

NeurIPS 2022accept

In offline model-based reinforcement learning (offline MBRL), we learn a dynamic model from historically collected data, and subsequently utilize the learned model and fixed datasets for policy learning, without further interacting with the environment. Offline MBRL algorithms can improve the effici…

2022

ALLSH: Active Learning Guided by Local Sensitivity and Hardness

NAACL 2022findings

Active learning, which effectively collects informative unlabeled data for annotation, reduces the demand for labeled data. In this work, we propose to retrieve unlabeled samples with a local sensitivity and hardness-aware acquisition function. The proposed method generates data copies through local…

Cited by 39SourcePDFScholar
2022

Passage-Mask: A Learnable Regularization Strategy for Retriever-Reader Models

EMNLP 2022main

Retriever-reader models achieve competitive performance across many different NLP tasks such as open question answering and dialogue conversations. In this work, we notice these models easily overfit the top-rank retrieval passages and standard training fails to reason over the entire retrieval pass…

2022

Regularizing a Model-based Policy Stationary Distribution to Stabilize Offline Reinforcement Learning

ICML 2022spotlight

Offline reinforcement learning (RL) extends the paradigm of classical RL algorithms to purely learning from static datasets, without interacting with the underlying environment during the learning process. A key challenge of offline RL is the instability of policy training, caused by the mismatch be…

2021

A Prototype-Oriented Framework for Unsupervised Domain Adaptation

NeurIPS 2021poster

Existing methods for unsupervised domain adaptation often rely on minimizing some statistical distance between the source and target samples in the latent space. To avoid the sampling variability, class imbalance, and data-privacy concerns that often plague these methods, we instead provide a memory…

2021

Alignment Attention by Matching Key and Query Distributions

NeurIPS 2021poster

The neural attention mechanism has been incorporated into deep neural networks to achieve state-of-the-art performance in various domains. Most such models use multi-head self-attention which is appealing for the ability to attend to information from different perspectives. This paper introduces ali…

2021

Contextual Dropout: An Efficient Sample-Dependent Dropout Module

ICLR 2021poster

Dropout has been demonstrated as a simple and effective module to not only regularize the training process of deep neural networks, but also provide the uncertainty estimation for prediction. However, the quality of uncertainty estimation is highly dependent on the dropout probabilities. Most curren…

2021

Learning with Different Amounts of Annotation: From Zero to Many Labels

EMNLP 2021main

Training NLP systems typically assumes access to annotated data that has a single human label per example. Given imperfect labeling from annotators and inherent ambiguity of language, we hypothesize that single label is not sufficient to learn the spectrum of language interpretation. We explore new…