Publications

^† denotes equal contribution.

NeurIPS · 2026 Accepted

Yu Gu^†, Zijun Yu^†, Vahid Partovi Nia, Masoud Asgharian

The problem: A common way to get better answers from a reasoning language model is to ask the same question several times and go with the most popular answer. That vote always produces an answer, but agreement alone comes with no guarantee about how often the chosen answer is wrong.
Why it matters: In tutoring, analysis or decision support, a confident wrong answer is often worse than an honest “I’m not sure.” A system that knows when to hold back is easier to trust.
What this work does: We give the model a principled way to hold back. Rather than a simple vote, it weighs the evidence across all of its attempts and answers only when that evidence is strong enough, with a statistical guarantee on how often it answers confidently and wrongly. It sits on top of an existing model with no retraining, and needs to abstain on only a small share of questions for the answers it does give to become noticeably more reliable.

NeurIPS · 2026 Accepted

Yu Gu^†, Zijun Yu^†, Chi-Kuang Yeh, Xinyu Wang, Ziyang Song

Paper coming soon

The problem: Medication recommendation systems suggest a set of drugs for each patient. The usual ways of judging and calibrating those suggestions count only exact matches, so every missed drug counts as equally wrong.
Why it matters: Not all misses are alike. By the standard classification of drugs, suggesting a closely related drug in place of the right one is a nearer miss than suggesting one from an unrelated class — and a reliability guarantee should be able to tell the two apart.
What this work does: We measure how far each patient’s misses fall within that drug classification, and adjust the suggested sets so that this distance stays within a tolerance the user chooses, with a statistical guarantee over the patient population and no retraining of the underlying model. The suggestions shift toward exact or closely related drugs and away from distant ones.

ICML · 2026 Accepted

Pengqi Liu^†, Zijun Yu^†, Mouloud Belbahri, Arthur Charpentier, Masoud Asgharian, Jesse C. Cresswell

The problem: Conformal prediction turns a model’s guess into a short list of plausible answers that contains the right one with a guaranteed probability. Fairness work in this area has mostly asked whether that guarantee holds equally for every group.
Why it matters: These lists are increasingly handed to people, or to language models, who make the final call. A guarantee that is equal on paper does not mean the decisions that follow are equally good for everyone.
What this work does: We look past the list to the decisions made with it, using language models as stand-in decision makers to measure how fair the outcomes actually are, and back this with theory. The practical finding: keeping the lists similar in size across groups does more for fair outcomes than equalizing the guarantee itself.

SLM-Agents @ NeurIPS · 2026 Accepted

Zijun Yu^†, Yu Gu^†, Vahid Partovi Nia, Masoud Asgharian

The problem: AI agents act in many steps, and harm can come from steps that each look harmless — reading a private file, then sending a message outside. Catching this means looking back over what the agent has done, but reviewing its whole history every time is costly and buries the signal.
Why it matters: Small language-model agents increasingly run on their own. Stopping too late lets harm through; stopping too early throws away useful work.
What this work does: We give the agent a risk budget: a rule for when to stop it, with a statistical guarantee on how much harm gets through first. The safety check then looks only at the private information an outgoing action could have drawn on, rather than the entire history, so it reads far less while letting more legitimate work finish. One honest caveat: an equally sized random slice of history helps about as much on average — following the connections pays off in specific cases, where it keeps harmless work from being stopped.

arXiv · 2026 Under review

Zijun Yu^†, Yu Gu^†, Vahid Partovi Nia, Masoud Asgharian

The problem: Conformal prediction needs a batch of examples whose answers have been checked. For language-model tasks, checking is real work — reading a retrieved document, grading a math solution, running a configuration — and it adds up quickly.
Why it matters: When checking is expensive, statistical guarantees become a luxury, used only where someone can afford to verify everything.
What this work does: We show that full checking is often unnecessary: partial checks frequently already settle everything the guarantee depends on. The method works out when enough has been verified and which examples to check next, returning exactly the same result as checking everything — so the guarantee is unchanged — at a lower, often much lower, cost.