Research

Reinforcement Learning

We designed a semi-supervised paradigm to use a light-weight scoring method as reward signal from large models to improve math reasoning on small models using different reinforcement learning algorithms, achieving knowledge distillation from an 8B model to 125M-600M models with a 10% improvement on grade school math problems and 1% improvement on competition-level math problems. Right now, we are working on expanding this knowledge distillation framework to more general tasks and non-autoregressive (such as discrete diffusion) models.

Optimization and its applications

We proposed a new first-order penalty-based method to solve bilevel optimization problems with convex constrained lower-level minimization problem. We applied this algorithm to perform hyperparameter tuning on variants of group distributionally robust optimization (worse group loss minimization) problems, achieving both empirically sound and theoretically grounded classification results.

Multi-modal generative modeling

Many real-world datasets are by nature multi-modal. Using information-bottleneck principles, we have developed information-theoretic unsupervised learning methods for multi-modal generation and have shown that our model renders meaningful representations on challenging visual-audio, biological, and synthetic datasets. We further extended this method to a semi-supervised setup where only a small percentage (0.1%-1%) of labels are available with sustained prediction and generation performance. Right now, we are working on incorporating synergy between multi-modal generation and multi-modal classification to improve both tasks.