AIConfPaper
← Search

Xuechen Zhang

11 accepted papers

2026

Continuous Chain of Thought Enables Parallel Exploration and Reasoning

Halil Alperen Gozeten, Muhammed Emrullah Ildiz, Xuechen Zhang, Hrayr Harutyunyan, Ankit Singh Rawat, Samet Oymak
ICLR 2026poster

Modern language models generate chain-of-thought traces by autoregressively sampling tokens from a finite vocabulary. While this discrete sampling has achieved remarkable success, conducting chain-of-thought with continuously-valued tokens (CoT2) offers a richer and more expressive alternative. Our…

Cited by 0Source ↗code ↗Scholar ↗
2026

SmartChunk Retrieval: Query-Aware Chunk Compression with Planning for Efficient Document RAG

Xuechen Zhang, Koustava Goswami, Samet Oymak, Jiasi Chen, Nedim Lipka
ICLR 2026poster

Retrieval-augmented generation (RAG) has strong potential for producing accurate and factual outputs by combining language models (LMs) with evidence retrieved from large text corpora. However, current pipelines are limited by static chunking and flat retrieval: documents are split into short, prede…

Cited by 0Source ↗Scholar ↗
2025

BREAD: Branched Rollouts from Expert Anchors Bridge SFT & RL for Reasoning

Xuechen Zhang, Zijian Huang, Yingcong Li, Chenshun Ni, Jiasi Chen, Samet Oymak
NeurIPS 2025poster

Small language models (SLMs) struggle to learn complex reasoning behaviors, especially when high-quality traces are scarce or difficult to learn from. A typical approach for training such models combines a supervised fine-tuning (SFT) stage, often to distill reasoning capabilities from a larger mode…

Cited by 0Source ↗Scholar ↗
2025

Test-Time Training Provably Improves Transformers as In-context Learners

Halil Alperen Gozeten, Muhammed Emrullah Ildiz, Xuechen Zhang, Mahdi Soltanolkotabi, Marco Mondelli, Samet Oymak
ICML 2025poster

Test-time training (TTT) methods explicitly update the weights of a model to adapt to the specific test instance, and they have found success in a variety of settings, including most recently language modeling and reasoning. To demystify this success, we investigate a gradient-based TTT algorithm fo…

Cited by 1Source ↗PDF ↗Scholar ↗
2024

A Score-Based Deterministic Diffusion Algorithm with Smooth Scores for General Distributions

Karthik Elamvazhuthi, Xuechen Zhang, Matthew Jacobs, Samet Oymak, Fabio Pasqualetti
AAAI 2024technical

Score matching based diffusion has shown to achieve the state of art results in generation modeling. In the original score matching based diffusion algorithm, the forward equation is a differential equation for which the probability density equation evolves according to a linear partial differential…

Cited by 1Source ↗PDF ↗Scholar ↗
2024

Class-Attribute Priors: Adapting Optimization to Heterogeneity and Fairness Objective

Xuechen Zhang, Mingchen Li, Jiasi Chen, Christos Thrampoulidis, Samet Oymak
AAAI 2024technical

Modern classification problems exhibit heterogeneities across individual classes: Each class may have unique attributes, such as sample size, label quality, or predictability (easy vs difficult), and variable importance at test-time. Without care, these heterogeneities impede the learning process, m…

Cited by 2Source ↗PDF ↗Scholar ↗
2024

Efficient Contextual LLM Cascades through Budget-Constrained Policy Learning

Xuechen Zhang, Zijian Huang, Ege Onur Taga, Carlee Joe-Wong, Samet Oymak, Jiasi Chen
NeurIPS 2024poster

Recent successes in natural language processing have led to the proliferation of large language models (LLMs) by multiple providers. Each LLM offering has different inference accuracy, monetary cost, and latency, and their accuracy further depends on the exact wording of the question (i.e., the spec…

Cited by 0Source ↗PDF ↗Scholar ↗
2024

Selective Attention: Enhancing Transformer through Principled Context Control

Xuechen Zhang, Xiangyu Chang, Mingchen Li, Amit Roy-Chowdhury, Jiasi Chen, Samet Oymak
NeurIPS 2024poster

The attention mechanism within the transformer architecture enables the model to weigh and combine tokens based on their relevance to the query. While self-attention has enjoyed major success, it notably treats all queries $q$ in the same way by applying the mapping $V^\top\text{softmax}(Kq)$, where…

Cited by 2Source ↗PDF ↗code ↗Scholar ↗
2023

Max-Margin Token Selection in Attention Mechanism

Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, Samet Oymak
NeurIPS 2023spotlight

Attention mechanism is a central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. In this work, we expl…

Cited by 56Source ↗PDF ↗code ↗Scholar ↗
2021

AutoBalance: Optimized Loss Functions for Imbalanced Data

Mingchen Li, Xuechen Zhang, Christos Thrampoulidis, Jiasi Chen, Samet Oymak
NeurIPS 2021poster

Imbalanced datasets are commonplace in modern machine learning problems. The presence of under-represented classes or groups with sensitive attributes results in concerns about generalization and fairness. Such concerns are further exacerbated by the fact that large capacity deep nets can perfectly…

Cited by 79Source ↗PDF ↗code ↗Scholar ↗