Yujia Bao

Member of Technical Staff

@thinkingmachines

Yujia Bao

I am a machine learning researcher and a life-long engineer. I love building things—pushing the frontier of AI to make it more useful, safe, and available to everyone.

My recent research spans scalable agent architectures, context management, LLM post-training, and evaluation. I am most excited about research that has direct product impact.

I received my Ph.D. in Computer Science from MIT CSAIL, advised by Regina Barzilay.

Experience & Education

Professional Experience

2026 - Current

Member of Technical Staff

Thinking Machines Lab

Developing SFT and RL recipes for end-to-end post-training. Leading the Tinker Cookbook framework development (top contributor).

2023 - 2026

Associate Director

Accenture

Leading a team of 80+ researchers and engineers. Developing AI Refinery, an agentic AI platform for enterprise.

2022 - 2023

Lead Machine Learning Scientist

Insitro

Machine learning research for drug discovery and development.

2017 - 2022

Researcher

MIT CSAIL

Machine learning research on interpretability, transfer learning, and fairness.

Education

2017 - 2022

Ph.D. in Computer Science

MIT CSAIL
Advisor: Regina Barzilay
2016 - 2017

M.A. in Mathematics

University of Wisconsin–Madison
2012 - 2016

B.S. in Mathematics

Shanghai Jiao Tong University

Recent Work

Tinker: LLM Post-Training Platform

A post-training platform for customizing language models. Developing SFT and RL recipes, evals, and infrastructure, as well as documentation and community.

Cookbook Docs

AI Refinery: Enterprise Agentic Platform

Leading engineering and research for AI Refinery, enabling developers to build and govern complex agentic workflows.

Website SDK

Enterprise LLM Customization

Leading the development of customized LLMs for high-stakes business and media domains.

Tech Report Fortune Analytics

Publications

See Google Scholar for full list of publications.

2026

Ares: Adaptive Reasoning Effort Selection for Efficient LLM Agents

Jingbo Yang, Bairu Hou, Wei Wei, Yujia Bao, Shiyu Chang
arXiv:2603.07915
PaperCode

Adaptive Data Augmentation with Multi-armed Bandit: Sample-Efficient Embedding Calibration for Implicit Pattern Recognition

Minxue Tang, Yangyang Yu, Aolin Ding, Maziyar Baran Pouyan, Taha Belkhouja, Yujia Bao
Conference on Computer Vision and Pattern Recognition (CVPR)
Paper

WebDART: Dynamic Decomposition and Re-planning for Complex Web Tasks

Jingbo Yang, Bairu Hou, Wei Wei, Shiyu Chang, Yujia Bao
Findings of the Association for Computational Linguistics: ACL
PaperCode

Observations and Remedies for Large Language Model Bias in Self-Consuming Performative Loop

Yaxuan Wang, Zhongteng Cai, Yujia Bao, Xueru Zhang, Yang Liu
Annual Meeting of the Association for Computational Linguistics (ACL)
Paper

MCP-Bench: Benchmarking tool-using llm agents with complex real-world tasks via mcp servers

Zhenting Wang, Qi Chang, Hemani Patel, Shashank Biju, Cheng-En Wu, Quan Liu, Aolin Ding, Alireza Rezazadeh, Ankit Shah, Yujia Bao, Eugene Siow
International Conference on Learning Representations (ICLR)
PaperCode

DRAGON: Guard LLM Unlearning in Context via Negative Detection and Reasoning

Yaxuan Wang, Chris Yuhao Liu, Quan Liu, Jinglong Pang, Wei Wei, Yujia Bao, Yang Liu
International Conference on Learning Representations (ICLR)
PaperCode

2025

PromptBridge: Cross-Model Prompt Transfer for Large Language Models

Yaxuan Wang, Quan Liu, Zhenting Wang, Zichao Li, Wei Wei, Yang Liu, Yujia Bao
arXiv:2512.01420
Paper

SFT-GO: Supervised Fine-Tuning with Group Optimization for Large Language Models

Gyuhak Kim, Sumiran Singh Thakur, Su Min Park, Wei Wei, Yujia Bao
arXiv:2506.15021
Paper

Collaborative Memory: Multi-User Memory Sharing in LLM Agents with Dynamic Access Control

Alireza Rezazadeh, Zichao Li, Ange Lou, Yuying Zhao, Wei Wei, Yujia Bao
arXiv:2505.18279
Paper

Advertising in AI systems: Society must be vigilant

Menghua Wu, Yujia Bao
arXiv:2505.18425
Paper

Enhancing Retrieval Systems with Inference-Time Logical Reasoning

Felix Faltings, Wei Wei, Yujia Bao
Annual Meeting of the Association for Computational Linguistics (ACL)
Paper

KVLink: Accelerating Large Language Models via Efficient KV Cache Reuse

Jingbo Yang, Bairu Hou, Wei Wei, Yujia Bao, Shiyu Chang
Conference on Neural Information Processing Systems (NeurIPS)
PaperCode

H-CoT: Hijacking the Chain-of-Thought Safety Reasoning Mechanism to Jailbreak Large Reasoning Models

Martin Kuo, Jianyi Zhang, Aolin Ding, Qinsi Wang, Louis DiValentin, Yujia Bao, Wei Wei, Hai Li, Yiran Chen
arXiv:2502.12893
PaperCode

Improving Data Efficiency via Curating LLM-Driven Rating Systems

Jinlong Pang, Jiaheng Wei, Ankit Parag Shah, Zhaowei Zhu, Yaxuan Wang, Chen Qian, Yang Liu, Yujia Bao, Wei Wei
International Conference on Learning Representations (ICLR)
PaperCode

LLM Unlearning via Loss Adjustment with Only Forget Data

Yaxuan Wang, Jiaheng Wei, Chris Yuhao Liu, Jinlong Pang, Quan Liu, Ankit Parag Shah, Yujia Bao, Yang Liu, Wei Wei
International Conference on Learning Representations (ICLR)
PaperCode

From isolated conversations to hierarchical schemas: Dynamic tree memory representation for llms

Alireza Rezazadeh, Zichao Li, Wei Wei, Yujia Bao
International Conference on Learning Representations (ICLR)
Paper

Sample, estimate, aggregate: A recipe for causal discovery foundation models

Menghua Wu, Yujia Bao, Regina Barzilay, Tommi Jaakkola
Transactions on Machine Learning Research (TMLR)
PaperCode

2024

Harnessing business and media insights with large language models

Yujia Bao, Ankit Parag Shah, Neeru Narang, Jonathan Rivers, Rajeev Maksey, Lan Guan, Louise N Barrere, Shelley Evenson, Rahul Basole, Connie Miao, others
arXiv:2406.06559
Paper

Channel Vision Transformers: An Image Is Worth C x 16 x 16 Words

Yujia Bao, Srinivasan Sivanandan, Theofanis Karaletsos
International Conference on Learning Representations (ICLR)
PaperCode

2023

Contextual Vision Transformers for Robust Representation Learning

Yujia Bao, Theofanis Karaletsos
arXiv:2305.19402
PaperCode

2022

Learning to Split for Automatic Bias Detection

Yujia Bao, Regina Barzilay
arXiv:2204.13749
PaperCode

Learning Stable Classifiers by Transferring Unstable Features

Yujia Bao, Shiyu Chang, Regina Barzilay
International Conference on Machine Learning (ICML)
PaperCode

2021

Predict then Interpolate: A Simple Algorithm to Learn Stable Classifiers

Yujia Bao, Shiyu Chang, Regina Barzilay
International Conference on Machine Learning (ICML)
PaperCode

2020

Few-shot Text Classification with Distributional Signatures

Yujia Bao*, Menghua Wu*, Shiyu Chang, Regina Barzilay
International Conference on Learning Representations (ICLR)
PaperCode

2018

Deriving Machine Attention from Human Rationales

Yujia Bao, Shiyu Chang, Mo Yu, Regina Barzilay
Conference on Empirical Methods in Natural Language Processing (EMNLP)
PaperCode