← Search

Kaiyue Wen

17 accepted papers

2026

QuestA: Expanding Reasoning Capacity in LLMs via Question Augmentation

ICLR 2026poster

Reinforcement learning (RL) has emerged as a central paradigm for training large language models (LLMs) in reasoning tasks. Yet recent studies question RL’s ability to incentivize reasoning capacity beyond the base model. This raises a key challenge: how can RL be adapted to solve harder reasoning p…

Cited by 0SourcecodeScholar
2025

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

ACL 2025long

This paper revisits the implementation of Load-Balancing-Loss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as NE ∑i=1NE fipi, where NE is the total number of experts, fi represents the frequency of expert i being selected, and pi denotes the average gat…

2025

From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency

ICLR 2025poster

Chain-of-thought (CoT) significantly enhances the reasoning performance of large language models (LLM). While current theoretical studies often attribute this improvement to increased expressiveness and computational capacity, we argue that expressiveness is not the primary limitation in the LLM re…

2025

Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

NeurIPS 2025oral

Gating mechanisms have been widely utilized, from early models like LSTMs and Highway Networks to recent state space models, linear attention, and also softmax attention. Yet, existing literature rarely examines the specific effects of gating. In this work, we conduct comprehensive experiments to sy…

Cited by 0SourcecodeScholar
2025

Overtrained Language Models Are Harder to Fine-Tune

ICML 2025poster

Large language models are pre-trained on ever-growing token budgets under the assumption that better pre-training performance translates to improved downstream models. In this work, we challenge this assumption and show that extended pre-training can make models harder to fine-tune, leading to degra…

Cited by 1SourcePDFScholar
2025

PaTH Attention: Position Encoding via Accumulating Householder Transformations

NeurIPS 2025poster

The attention mechanism is a core primitive in modern large language models (LLMs) and AI more broadly. Since attention by itself is permutation-invariant, position encoding is essential for modeling structured domains such as language. Rotary position encoding (RoPE) has emerged as the de facto sta…

Cited by 0SourceScholar
2025

RNNs are not Transformers (Yet): The Key Bottleneck on In-Context Retrieval

ICLR 2025poster

This paper investigates the gap in representation powers of Transformers and Recurrent Neural Networks (RNNs), which are more memory efficient than Transformers. We aim to understand whether RNNs can match the performance of Transformers, particularly when enhanced with Chain-of-Thought (CoT) prompt…

2025

Symmetrical Visual Contrastive Optimization: Aligning Vision-Language Models with Minimal Contrastive Images

ACL 2025long

Recent studies have shown that Large Vision-Language Models (VLMs) tend to neglect image content and over-rely on language-model priors, resulting in errors in visually grounded tasks and hallucinations. We hypothesize that this issue arises because existing VLMs are not explicitly trained to genera…

Cited by 0SourcePDFScholar
2025

Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks?

ICML 2025poster

Large language models (LLMs) exhibit remarkable task generalization, solving tasks they were never explicitly trained on with only a few demonstrations. This raises a fundamental question: When can learning from a small set of tasks generalize to a large task family? In this paper, we investigate t…

Cited by 0SourcePDFScholar
2025

Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape View

ICLR 2025poster

Training language models currently requires pre-determining a fixed compute budget because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can i…

Cited by 0SourcePDFScholar
2023

Benign Overfitting in Classification: Provably Counter Label Noise with Larger Models

ICLR 2023poster

Studies on benign overfitting provide insights for the success of overparameterized deep learning models. In this work, we examine whether overfitting is truly benign in real-world classification tasks. We start with the observation that a ResNet model overfits benignly on Cifar10 but not benignly o…

Cited by 8SourcePDFScholar
2023

Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization

NeurIPS 2023oral

Despite extensive studies, the underlying reason as to why overparameterized neural networks can generalize remains elusive. Existing theory shows that common stochastic optimizers prefer flatter minimizers of the training loss, and thus a natural potential explanation is that flatness implies gener…

Cited by 32SourcePDFScholar
2023

Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars

NeurIPS 2023poster

Transformer interpretability aims to understand the algorithm implemented by a learned Transformer by examining various aspects of the model, such as the weight matrices or the attention patterns. In this work, through a combination of theoretical results and carefully controlled experiments on synt…

Cited by 23SourcePDFScholar
2022

Finding Skill Neurons in Pre-trained Transformer-based Language Models

EMNLP 2022main

Transformer-based pre-trained language models have demonstrated superior performance on various natural language processing tasks. However, it remains unclear how the skills required to handle these tasks distribute among model parameters. In this paper, we find that after prompt tuning for specific…

2022

On Transferability of Prompt Tuning for Natural Language Processing

NAACL 2022long

Prompt tuning (PT) is a promising parameter-efficient method to utilize extremely large pre-trained language models (PLMs), which can achieve comparable performance to full-parameter fine-tuning by only tuning a few soft prompts. However, PT requires much more training time than fine-tuning. Intuiti…