Yiheng Zhang (张怡蘅)
M.S. Student in Artificial Intelligence Engineering · Carnegie Mellon University, Electrical and Computer Engineering
- I am Yiheng Zhang, an incoming M.S. student in Artificial Intelligence Engineering (Electrical and Computer Engineering) at Carnegie Mellon University, starting Sep. 2026. I received my B.S. in Computer Science from the University of Wisconsin-Madison, where I was advised by Prof. Grigoris Chrysos. My GPA at UW-Madison was 3.978/4.0.
- My research interests are machine learning and natural language processing, with a focus on large language models (LLMs), diffusion language models (DLMs), LLM agents for code, and trustworthy AI. I have worked on correction-aware training for diffusion language models, repository-level PR data pipelines for code agents, and multimodal video understanding.
- I am actively looking for research opportunities and summer internships. Please feel free to contact me if you think I would be a good fit. Thank you.
Experience
Internship Experience
Code Agent Algorithm InternLLM
ModelBest · Beijing, China · May. 2026 - Aug. 2026
- Designed and implemented an end-to-end repository-level PR data pipeline for underrepresented languages such as TypeScript and Go, using LLM Agents for issue rewriting, scoring, and quality filtering to collect 140K+ high-quality PR tasks.
- Conducted SFT on Qwen3-4B to validate data quality; improved the resolve rate from 13% to 22% on SWE-Pro after 8K training steps, including +7 pts on Go and +5 pts on TypeScript.
- Built a multimodal long-video event-reminder data generation pipeline with automated annotation and filtering, generated 400K training samples and improved overall performance on MiniCPM-o by 10%.
Research Experience
Diffusion Language Models and RemaskingMachine Learning
University of Wisconsin-Madison · Prof. Grigorios G. Chrysos · Jun. 2025 - May. 2026
- Developed a correction-aware post-training objective that trains Diffusion Language Models to identify and revise erroneous visible tokens through mixed token corruptions.
- Built the core LLaDA-based Masked Diffusion Language Model framework for iterative denoising and integrated the proposed mixture training method for error correction.
- Created the Code Revision Benchmark (CRB), an executable benchmark with controllable corruptions over operators, identifiers, and literals to evaluate code error detection and correction.
Large Language Models for Arithmetic ReasoningMachine Learning
University of Wisconsin-Madison · Prof. Grigorios G. Chrysos · Feb. 2025 - Jun. 2025
- Designed controlled experiments on addition, multiplication, parity, and sorting to study optimization and length generalization in arithmetic reasoning.
- Implemented modular plug-in components and trained small models to evaluate and improve arithmetic reasoning.
- Benchmarked the Polynomial Neural Network against Transformer baselines and found faster convergence and stronger length generalization on tasks, including settings with learnable embeddings.
Workshop Volunteer · Reliable Agentic AI
ICLR 2026 · Apr. 2026
- Served as a volunteer for the interdisciplinary workshop Reliable Agentic AI: From Hallucination to Trustworthy Autonomy at ICLR 2026.
Projects
SnapBadgers: Multimodal On-device Music Recommendation
On-device ML · Feb. 2026 - May. 2026
Capstone Project in collaboration with Qualcomm · University of Wisconsin-Madison
- Built an on-device ML pipeline fusing text, vision, and sensor into a unified embedding for song recommendation.
- Developed Android UI components to integrate text, camera, and sensor inputs into the inference workflow.
- Implemented 12 unit and evaluation tests to validate pipeline correctness and fallback behavior across devices and emulators to improve deployment reliability.
On-device MLMultimodalAndroid
Reading Group
Internal reading group for sharing SoTA papers and new ideas.
Time Every Wednesday and Friday
| Date | Theme | Material |
|---|---|---|
| 2025-05-30 | Understanding Diffusion Models: A Unified Perspective | Link |
| 2025-06-10 | DDPM, DDIM and Guidance | Link |
| 2025-06-19 | Large Language Diffusion Models | Link |
| 2025-09-10 | Demystifying Foreground-Background Memorization in Diffusion Models | Link |
| 2025-09-26 | Persona Features Control Emergent Misalignment | Link |
| 2025-10-10 | TraceDet: Hallucination Detection from the Decoding Trace of Diffusion Large Language Models | Link |