Zaid Khan

I'm a 2nd-year PhD student at Mohit Bansal's group (MURGe Lab) at UNC Chapel Hill, supported by a DoD NDSEG fellowship.

Most recently, I was a research intern at Ai2 working with Tanmay Gupta, Luca Weihs, and Ranjay Krishna. Before joining UNC, I was a student researcher in the Media Analytics Group at NEC Laboratories America under Manmohan Chandraker. I completed my BS+MS at Northeastern, where I worked with Raymond Fu.

 /   /   /   /   / 

Background

Before graduate school, I spent ~3 years as an early member of the engineering / data science organizations at two high growth startups: Roadie (acquired by UPS for $500m) and OneTrack.AI as software engineer, where I led efforts to scale data infrastructure to match growth, and worked on a range of challenging problems, including embedded deep learning, fault-tolerant distributed systems, realtime adaptive pricing, and data pipelines for time series and computer vision tasks.

Outside of research, I lift weights, read (here's my goodreads profile), watch mixed martial arts, and sometimes wonder whether randomness is real.

Current Research

Prior Research

• Self-training / self-improvement
  • using reinforced self-training to improve program synthesis (Khan et al., CVPR 2024)
  • using unlabeled data to improve vision-language reasoning (Khan et al., CVPR 2023)
• Using uncertainty during reasoning and decision-making
  • using self-consistency to identify unreliable knowledge (Khan et al., CVPR 2024)
  • using uncertainty to decide when to expend more test-time compute (Khan et al., NeurIPS 2023)
• Vision-language representation learning
  • by aligning representations with minimal parameter updates (Khan et al., ICLR 2023)
  • by learning to reconstruct each modality (Khan et al., ECCV 2022)

Publications

Papers where I'm a lead author are marked; some are highlighted.

We've been working on a way to get better on-policy token-level rewards for LLMs + RL! Self-distillation gives token-level rewards, using divergence against a teacher policy given privileged info (i.e true final answer). What if you could use multiple forms of privileged info?

Continually updated envs (e.g. Git repo histories, evolving docs) are central to knowledge work. Reasoning about these requires long context understanding + resolving temporally distributed / interfering changes to the env state. How well do LLM agents / memory systems do?

How can LLM agents solve long-horizon tasks or explore an environment with many details while using only a fixed amount of context window? In Agent-BRACE, we use RL to build+use belief states that decompose into natural-language claims about the world and their epistemic uncertainty.

ToM-SB is an environment where defender LLMs compete against attacker LLMs seeking sensitive information. To win, the defender must fool the attacker into leaving with the wrong information. RL in ToM-SB results in bidirectional emergence of theory-of-mind (ToM) and fooling ability.

* Equal contribution

How do we enable RLVR on hard problems where rollouts yield zero reward? Imitating expert trajectories is off-policy; instead, we reformulate each hard problem into a multi-level curriculum. Skills learned on the easier variants transfer back to the hard set where standard RLVR fails.

A generative Process Reward Model for long-horizon information seeking. PRINTS learns to verbally estimate the information gain of tool calls, enabling a 4B model to effectively guide 30B+ agents through noisy, complex web search tasks like GAIA Level 3.

What if we could transform advanced math problems into abstract programs that can generate endless, verifiable problem variants? EFAGen uses test-time search with execution feedback to infer executable functional abstractions (EFAs) in Python for diverse math problems, including Olympiad-level problems.

Neural tree search for repo-level code-use planning. MutaGReP explores plan space through LLM guided mutations, while grounding the plan to functionality in the codebase using a symbol retriever.

Testing is a critical part of software engineering — what if we could automatically discover inputs which break your code? We show how to train SLMs (Qwen2.5-7B + Llama3.1-8B) to generate unit tests that break code and are useful for debugging.

* Equal contribution

In composable AI systems for visual reasoning, agents must reliably orchestrate noisy, black-box models as tools. We propose a method that trains the agent to perform online error recovery: it learns to identify potential failures in tool outputs and iteratively refine its queries to extract the correct information.

We show how to selectively decompose complex questions into simpler sub-questions to improve zero-shot performance on challenging multimodal reasoning tasks.

Getting labels for a multimodal dataset can be expensive. We show how you can use unlabeled images to improve performance on data-scarce multimodal tasks.

We explore creating CLIP-like models by minimally updating already-trained vision and language models, finding that updating less than 7% of parameters can match full model training.

We demonstrate a very data-efficient way to align vision and language by learning to reconstruct each modality from the other.

Understanding the emotional content of social media posts is difficult for traditional sentiment analysis models. We show that language models do a good job of this if the post can be translated into a natural input space for them.

Are notions of algorithmic fairness based on racial categories meaningful? We study computer vision datasets that use racial categories, and empirically show that the racial categories encoded in each dataset are often highly inconsistent with each other and with human intuitions.

Adapted from Jon Barron's website..