Zijun Yu PhD student · McGill
Montréal, QC

Zijun Yu

I'm a PhD student in Mathematics & Statistics at McGill University, working on conformal prediction and risk control for reliable, safe LLMs and AI agents. My work has appeared at ICML 2026 and NeurIPS 2026. I've been a research intern at Layer 6 AI, and before the PhD I spent four years as a software developer.

Selected publications

Recent work in conformal prediction.

A selection — see all publications. † denotes equal contribution.

    NeurIPS · 2026 Accepted

    Pause and Reflect: Conformal Aggregation for Chain-of-Thought Reasoning

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

    • Conformal Prediction
    • LLM Reasoning
    arXiv ↗
    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.
    ICML · 2026 Accepted

    Beyond Procedure: Substantive Fairness in Conformal Prediction

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

    • Conformal Prediction
    • Fairness
    arXiv ↗ Code ↗
    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

    G-CARB: Graph-Localized Conformal Agent Risk Budget for Compositional Harm

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

    • Conformal Prediction
    • Agent Safety
    arXiv ↗
    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.