← Search

Kangwook Lee

45 accepted papers

2026

Draft-based Approximate Inference for LLMs

ICLR 2026poster

Optimizing inference for long-context large language models (LLMs) is increasingly important due to the quadratic compute and linear memory cost of Transformers. Existing approximate inference methods, including key-value (KV) cache dropping, sparse attention, and prompt compression, typically rely…

Cited by 0SourcecodeScholar
2026

Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Models

ICML 2026poster

Transformer-based large language models exhibit in-context learning, enabling adaptation to downstream tasks via few-shot prompting with demonstrations. In practice, such models are often fine-tuned to improve zero-shot performance on downstream tasks, allowing them to solve tasks without examples a…

Cited by 0SourceScholar
2026

LookaheadKV: Fast and Accurate KV Cache Eviction by Glimpsing into the Future without Generation

ICLR 2026poster

Transformer-based large language models (LLMs) rely on key–value (KV) caching to avoid redundant computation during autoregressive inference. While this mechanism greatly improves efficiency, the cache size grows linearly with the input sequence length, quickly becoming a bottleneck for long‑context…

Cited by 0SourcecodeScholar
2026

Orak: A Foundational Benchmark for Training and Evaluating LLM Agents on Diverse Video Games

ICLR 2026poster

Large Language Model (LLM) agents are reshaping the game industry, by enabling more intelligent and human-preferable characters. Yet, current game benchmarks fall short of practical needs: they lack evaluations of diverse LLM capabilities across various game genres, studies of agentic modules crucia…

Cited by 0SourceScholar
2026

ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs

ICLR 2026poster

While most autoregressive LLMs are constrained to one-by-one decoding, diffusion LLMs (dLLMs) have attracted growing interest for their potential to dramatically accelerate inference through parallel decoding. Despite this promise, the conditional independence assumption in dLLMs causes parallel dec…

Cited by 0SourcecodeScholar
2026

ReJump: A Tree-Jump Representation for Analyzing and Improving LLM Reasoning

ICML 2026poster

Large Reasoning Models (LRMs) are Large Language Models (LLMs) explicitly trained to generate long-form Chain-of-Thoughts (CoTs), achieving impressive success on challenging tasks like math and programming. However, their underlying reasoning "algorithms" remain poorly understood. To investigate thi…

Cited by 0SourceScholar
2025

Everything Everywhere All at Once: LLMs can In-Context Learn Multiple Tasks in Superposition

ICML 2025spotlight

Large Language Models (LLMs) have demonstrated remarkable in-context learning (ICL) capabilities. In this study, we explore a surprising phenomenon related to ICL: LLMs can perform multiple, computationally distinct ICL tasks simultaneously, during a single inference call, a capability we term task…

Cited by 3SourcePDFScholar
2025

From Artificial Needles to Real Haystacks: Improving Retrieval Capabilities in LLMs by Finetuning on Synthetic Data

ICLR 2025poster

Recent studies have shown that Large Language Models (LLMs) struggle to accurately retrieve information and maintain reasoning capabilities when processing long-context inputs. To address these limitations, we propose a finetuning approach utilizing a carefully designed synthetic dataset comprising…

2025

Parameter-Efficient Fine-Tuning of State Space Models

ICML 2025poster

Deep State Space Models (SSMs), such as Mamba (Gu & Dao, 2024), have become powerful tools for language modeling, offering high performance and linear scalability with sequence length. However, the application of parameter-efficient fine-tuning (PEFT) methods to SSM-based models remains largely unde…

2025

Rare-to-Frequent: Unlocking Compositional Generation Power of Diffusion Models on Rare Concepts with LLM Guidance

ICLR 2025spotlight

State-of-the-art text-to-image (T2I) diffusion models often struggle to generate rare compositions of concepts, e.g., objects with unusual attributes. In this paper, we show that the compositional generation power of diffusion models on such rare concepts can be significantly enhanced by the Large L…

2025

Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges

ICML 2025poster

Large language models often struggle with length generalization and solving complex problem instances beyond their training distribution. We present a self-improvement approach where models iteratively generate and learn from their own solutions, progressively tackling harder problems while maintain…

Cited by 2SourcePDFScholar
2025

Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

ACL 2025long

Knowledge distillation can be a cost-effective technique to distill knowledge in Large Language Models, if the teacher output logits can be pre-computed and cached. However, successfully applying this to pre-training remains largely unexplored. In this work, we prove that naive approaches for sparse…

Cited by 0SourcePDFScholar
2025

VersaPRM: Multi-Domain Process Reward Model via Synthetic Reasoning Data

ICML 2025oral

Process Reward Models (PRMs) have proven effective at enhancing mathematical reasoning for Large Language Models (LLMs) by leveraging increased inference-time computation. However, they are predominantly trained on mathematical data and their generalizability to non-mathematical domains has not been…

Cited by 4SourcePDFScholar
2024

Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks

ICML 2024poster

State-space models (SSMs), such as Mamba (Gu & Dao, 2023), have been proposed as alternatives to Transformer networks in language modeling, incorporating gating, convolutions, and input-dependent token selection to mitigate the quadratic cost of multi-head attention. Although SSMs exhibit competitiv…

2024

Image Clustering Conditioned on Text Criteria

ICLR 2024poster

Classical clustering methods do not provide users with direct control of the clustering results, and the clustering results may not be consistent with the relevant criterion that a user has in mind. In this work, we present a new methodology for performing image clustering based on user-specified cr…

2024

Looped Transformers are Better at Learning Learning Algorithms

ICLR 2024poster

Transformers have demonstrated effectiveness in in-context solving data-fitting problems from various (latent) models, as reported by Garg et al. (2022). However, the absence of an inherent iterative structure in the transformer architecture presents a challenge in emulating the iterative algorithms…

2024

Memorization Capacity for Additive Fine-Tuning with Small ReLU Networks

UAI 2024poster

Fine-tuning large pre-trained models is a common practice in machine learning applications, yet its mathematical analysis remains largely unexplored. In this paper, we study fine-tuning through the lens of memorization capacity. Our new measure, the Fine-Tuning Capacity (FTC), is defined as the maxi…

Cited by 0SourcePDFScholar
2024

Teaching Arithmetic to Small Transformers

ICLR 2024poster

Large language models like GPT-4 exhibit emergent capabilities across general-purpose tasks, such as basic arithmetic, when trained on extensive text data, even though these tasks are not explicitly encoded by the unsupervised, next-token prediction objective. This study investigates how even small…

2023

DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models

NeurIPS 2023poster

Learning from human feedback has been shown to improve text-to-image models. These techniques first learn a reward function that captures what humans care about in the task and then improve the models based on the learned reward function. Even though relatively simple approaches (e.g., rejection sam…

2023

Equal Improvability: A New Fairness Notion Considering the Long-term Impact

ICLR 2023poster

Devising a fair classifier that does not discriminate against different groups is an important problem in machine learning. Although researchers have proposed various ways of defining group fairness, most of them only focused on the immediate fairness, ignoring the long-term impact of a fair classif…

2023

Intuitive Access to Smartphone Settings Using Relevance Model Trained by Contrastive Learning

AAAI 2023technical

The more new features that are being added to smartphones, the harder it becomes for users to find them. This is because the feature names are usually short and there are just too many of them for the users to remember the exact words. The users are more comfortable asking contextual queries that d…

Cited by 0SourcePDFScholar
2023

Looped Transformers as Programmable Computers

ICML 2023poster

We present a framework for using transformer networks as universal computers by programming them with specific weights and placing them in a loop. Our input sequence acts as a punchcard, consisting of instructions and memory for data read/writes. We demonstrate that a constant number of encoder laye…

Cited by 110SourcePDFScholar
2023

Prompted LLMs as Chatbot Modules for Long Open-domain Conversation

ACL 2023findings

In this paper, we propose MPC (Modular Prompted Chatbot), a new approach for creating high-quality conversational agents without the need for fine-tuning. Our method utilizes pre-trained large language models (LLMs) as individual modules for long-term consistency and flexibility, by using techniques…

2022

GenLabel: Mixup Relabeling using Generative Models

ICML 2022spotlight

Mixup is a data augmentation method that generates new data points by mixing a pair of input data. While mixup generally improves the prediction performance, it sometimes degrades the performance. In this paper, we first identify the main causes of this phenomenon by theoretically and empirically an…

Cited by 20SourcePDFScholar
2022

LIFT: Language-Interfaced Fine-Tuning for Non-language Machine Learning Tasks

NeurIPS 2022accept

Fine-tuning pretrained language models (LMs) without making any architectural changes has become a norm for learning various language downstream tasks. However, for non-language downstream tasks, a common practice is to employ task-specific designs for input, output layers, and loss functions. For i…

2022

Learning to Embed Multi-Modal Contexts for Situated Conversational Agents

NAACL 2022findings

The Situated Interactive Multi-Modal Conversations (SIMMC) 2.0 aims to create virtual shopping assistants that can accept complex multi-modal inputs, i.e. visual appearances of objects and user utterances. It consists of four subtasks, multi-modal disambiguation (MM-Disamb), multi-modal coreference…

Cited by 13SourcePDFScholar
2022

Rare Gems: Finding Lottery Tickets at Initialization

NeurIPS 2022accept

Large neural networks can be pruned to a small fraction of their original size, with little loss in accuracy, by following a time-consuming "train, prune, re-train" approach. Frankle & Carbin conjecture that we can avoid this by training lottery tickets, i.e., special sparse subnetworks found at ini…

2022

Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance

NeurIPS 2022accept

Score-based generative models are shown to achieve remarkable empirical performances in various applications such as image generation and audio synthesis. However, a theoretical understanding of score-based diffusion models is still incomplete. Recently, Song et al. showed that the training objectiv…

2022

Utilizing Language-Image Pretraining for Efficient and Robust Bilingual Word Alignment

EMNLP 2022finding

Word translation without parallel corpora has become feasible, rivaling the performance of supervised methods. Recent findings have shown the improvement in accuracy and robustness of unsupervised word translation (UWT) by utilizing visual observations, which are universal representations across lan…

2021

FairBatch: Batch Selection for Model Fairness

ICLR 2021poster

Training a fair machine learning model is essential to prevent demographic disparity. Existing techniques for improving model fairness require broad changes in either data preprocessing or model training, rendering themselves difficult-to-adopt for potentially already complex machine learning system…

2021

Gradient Inversion with Generative Image Prior

NeurIPS 2021poster

Federated Learning (FL) is a distributed learning framework, in which the local data never leaves clients’ devices to preserve privacy, and the server trains models on the data via accessing only the gradients of those local data. Without further privacy mechanisms such as differential privacy, this…

2020

Attack of the Tails: Yes, You Really Can Backdoor Federated Learning

NeurIPS 2020poster

Due to its decentralized nature, Federated Learning (FL) lends itself to adversarial attacks in the form of backdoors during training. The goal of a backdoor is to corrupt the performance of the trained model on specific sub-tasks (e.g., by classifying green cars as frogs). A range of FL backdoor at…

2020

FR-Train: A Mutual Information-Based Approach to Fair and Robust Training

ICML 2020poster

Trustworthy AI is a critical issue in machine learning where, in addition to training a model that is accurate, one must consider both fair and robust training in the presence of data bias and poisoning. However, the existing model fairness techniques mistakenly view poisoned data as an additional b…

2018

Simulated+Unsupervised Learning With Adaptive Data Generation and Bidirectional Mappings

ICLR 2018poster

Collecting a large dataset with high quality annotations is expensive and time-consuming. Recently, Shrivastava et al. (2017) propose Simulated+Unsupervised (S+U) learning: It first learns a mapping from synthetic data to real data, translates a large amount of labeled synthetic data to the ones tha…

Cited by 21SourcePDFScholar