Andrew (Ching-Yuan) Bai

Hello! Welcome to my personal website.

I am a 5th year PhD student in Computer Science at UCLA, where I work with Cho-Jui Hsieh. My research interest involves studying the memorization and forgetting mechanisms when training machine learning models.

Since the LLM boom, I focus on applications in LLM post-training. My latest project investigates why fine-tuning LLM with RLHF leads to less forgetting compared to supervised fine-tuning. Previously I worked on selecting the best instruction-tuning checkpoint for downstream preference alignment (RLHF) performance in LLMs. I also investigated the forgetting of context-awareness when pretrained LLMs are supervised fine-tuned on instruction data.

My favorite part of research is collaborating and engaging in insightful discussions with fellow reseachers. I have a diverse exposure to various topics through collaborations. My collaboration projects include reward modeling, long video generation, LLM agents, instruction selection, diffusion model data memorization, and prompt optimization.

Before the LLM boom I worked on data selection. I proposed a simple and easy-to-adopt sample selection schemes for prioritizing data samples during training. I also worked on developing practical interpretation methods for black-box models, allowing humans to better understand and trust machine learning models in real life (concept-based interpretability).

I was an undergraduate student in Computer Science at National Taiwan University. I worked with Hsuan-Tien Lin on generative modeling and time series forecasting. We held the first-ever generative modeling competition in collaboration with Kaggle. I also worked with Chung-Wei Lin on system verification and falsification.

Email: andrewbai [AT] cs.ucla.edu

Links: [Resume] [CV] [Github] [Google Scholar] [Linkedin]

Publications

2025

• Andrew Bai, Justin Cui, Ruochen Wang, Cho-Jui Hsieh. Concepts or Skills? Rethinking Instruction Selection for Multi-modal Models. Under submission review.
  [bib | arxiv]
• Yihan Wang*, Andrew Bai*, Nanyun Peng, Cho-Jui Hsieh. On the Loss of Context-awareness in General Instruction Fine-tuning. Under submission review.
  [bib | arxiv | code]
• Andrew Bai, Chih-Kuan Yeh, Cho-Jui Hsieh, Ankur Taly. An Efficient Rehearsal Scheme for Catastrophic Forgetting Mitigation during Multi-stage Fine-tuning. In Findings of the Association for Computational Linguistics: 2025 Annual Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics (NAACL), Apr 2025.
  [bib | arxiv | code]

2024

• Tong Xie*, Haoyu Li*, Andrew Bai, Cho-Jui Hsieh. Data Attribution for Diffusion Models: Timestep-induced Bias in Influence Estimation. In Transactions on Machine Learning Research (TMLR), Jun 2024.
  [bib | arxiv]

2023

• Andrew Bai, Chih-Kuan Yeh, Pradeep Ravikumar, Neil Y. C. Lin, Cho-Jui Hsieh. Concept Gradient: Concept-based Interpretation Without Linear Assumption. In Proceedings of the 11th International Conference on Learning Representations (ICLR), May 2023.
  [bib | arxiv | code]

2022

• Andrew Bai, Cho-Jui Hsieh, Wendy Chih-wen Kan, Hsuan-Tien Lin. Reducing Training Sample Memorization in GANs by Training with Memorization Rejection. Arxiv preprint.
  [bib | arxiv | code]

2021

• Ching-Yuan Bai, Hsuan-Tien Lin, Colin Raffel, and Wendy Chih-wen Kan. On training sample memorization: Lessons from benchmarking generative modeling with a large-scale competition. In Proceedings of the 27th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), August 2021.
  [bib | data | arxiv | code ]

2020

• Shih-Lun Wu*, Ching-Yuan Bai*, Kai-Chieh Chang, Yi-Ting Hsieh, Chao Huang, Chung-Wei Lin, Eunsuk Kang, and Qi Zhu. Efficient system verification with multiple weakly-hard constraints for runtime monitoring. In Proceedings of the International Conference on Runtime Verification (RV), October 2020.
  [bib | pdf ]
• Ching-Yuan Bai, Buo-Fu Chen, and Hsuan-Tien Lin. Benchmarking tropical cyclone rapid intensification with satellite images and attention-based deep models. In Proceedings of the European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML/PKDD), September 2020.
  [bib | data | arxiv | code ]