Selected Publications
(Google Scholar Profile) (Full List at DBLP)
DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts
Jiading Gai*, Shuai Zhang*, Xiang Song, Bernie Wang, George Karypis.
Preprint, 2026.
HetRL: Efficient Reinforcement Learning for LLMs in Heterogeneous Environments
Yongjun He*, Shuai Zhang*, Jiading Gai, Xiyuan Zhang, Boran Han, Bernie Wang, Huzefa Rangwala, George Karypis.
MLSys, 2026. Ninth Annual Conference on Machine Learning and Systems
Tiny but Mighty: A Software-Hardware Co-Design Approach for Efficient Multimodal Inference on Battery-Powered Small Devices
Yilong Li, Shuai Zhang, Yijing Zeng, Hao Zhang, Xinmiao Xiong, Jingyu Liu, Pan Hu, Suman Banerjee.
ICLR, 2026. The Fourteenth International Conference on Learning Representations
Mitra: Mixed synthetic priors for enhancing tabular foundation models
Xiyuan Zhang, Danielle C Maddix, Junming Yin, Nick Erickson, Abdul Fatir Ansari, Boran Han, Shuai Zhang, Leman Akoglu, Christos Faloutsos, Michael W Mahoney, Cuixiong Hu, Huzefa Rangwala, George Karypis, Bernie Wang.
NeurIPS, 2025. The Thirty-ninth Annual Conference on Neural Information Processing Systems
PipeRAG: Fast Retrieval-Augmented Generation via Algorithm-System Co-design
Wenqi Jiang, Shuai Zhang, Boran Han, Jie Wang, Bernie Wang, Tim Kraska.
KDD, 2025. code
PalmBench: A Comprehensive Benchmark of Compressed Large Language Models on Mobile Platforms
Yilong Li, Jingyu Liu, Hao Zhang, M Badri Narayanan, Utkarsh Sharma, Shuai Zhang, Pan Hu, Yijing Zeng, Jayaram Raghuram, Suman Banerjee.
ICLR, 2025.
Unraveling the Gradient Descent Dynamics of Transformers
Bingqing Song, Boran Han, Shuai Zhang, Jie Ding, Mingyi Hong
NeurIPS 2024 . Thirty-seventh Conference on Neural Information Processing Systems.
Discovering Bias in Latent Space: An Unsupervised Debiasing Approach
Dyah Adila, Shuai Zhang, Boran Han, Bernie Wang.
ICML 2024. Forty-first International Conference on Machine Learning.
CaMML: Context-Aware Multimodal Learner for Large Models website
Yixin Chen*, Shuai Zhang*, Boran Han, Tong He, Bo Li.
ACL 2024. The 62nd Annual Meeting of the Association for Computational Linguistics. code
Area Chair Award.
CoMM: Collaborative Multi-Agent, Multi-Reasoning-Path Prompting for Complex Problem Solving
Pei (Patrick) Chen,Boran Han, Shuai Zhang (Corresponding Author).
NAACL 2024 Findings. 2024 Annual Conference of the North American Chapter of the Association for Computational Linguistics . code
Transferring Knowledge From Large Foundation Models to Small Downstream Models
Shikai Qiu, Boran Han, Danielle C. Maddix, Shuai Zhang, Bernie Wang, Andrew Gordon Wilson .
ICML 2024. Forty-first International Conference on Machine Learning. code
Bridging Sources in Geospatial Sensing with Cross Sensor Pretraining
Boran Han, Shuai Zhang, Xingjian Shi, Markus Reichstein
CVPR 2024. The IEEE / CVF Computer Vision and Pattern Recognition Conference. code
Prompt Pre-Training with Twenty-Thousand Classes for Open-Vocabulary Visual Recognition
Shuhuai Ren, Aston Zhang, Yi Zhu, Shuai Zhang, Shuai Zheng, Mu Li, Alex Smola, Xu Sun.
NeurIPS 2023 . Thirty-seventh Conference on Neural Information Processing Systems. code
Data-Informed Geometric Space Selection
Shuai Zhang, Wenqi Jiang.
NeurIPS 2023 . Thirty-seventh Conference on Neural Information Processing Systems.
xFraud: Explainable Fraud Transaction Detection
Susie Xi Rao, Shuai Zhang, Zhichao Han, Zitao Zhang, Wei Min, Zhiyao Chen, Yinan Shan, Yang Zhao, Ce Zhang.
VLDB 2022 . The Proceedings of the VLDB Endowment. code
Neural Methods for Logical Reasoning over Knowledge Graphs
Alfonso Amayuelas, Shuai Zhang, Susie Xi Rao, Ce Zhang.
ICLR 2022 . The International Conference on Learning Representations. code
Beyond Fully-Connected Layers with Quaternions: Parameterization of Hypercomplex Multiplications with 1/n Parameters
Aston Zhang, Yi Tay, Shuai Zhang, Alvin Chan, Anh Tuan Luu, Siu Hui, Jie Fu.
ICLR 2021. The International Conference on Learning Representations. code
Outstanding Paper Award.
Self-Instantiated Recurrent Units with Dynamic Soft Recursion
Aston Zhang, Yi Tay, Yikang Shen, Alvin Chan, Shuai Zhang
NeurIPS 2021 . Thirty-fifth Conference on Neural Information Processing Systems. code
MicroRec: Accelerating Deep Recommendation Systems to Microseconds by Hardware and Data Structure Solutions
Wenqi Jiang, Zhenhao He, Shuai Zhang, Thomas B. Preußer, Kai Zeng, Liang Feng, Jiansong Zhang, Tongxuan Liu, Yong Li, Jingren Zhou, Ce Zhang, Gustavo Alonso.
MLSys 2021. Fourth Conference on Machine Learning and Systems. code
HyperML: A Boosting Metric Learning Approach in Hyperbolic Space for Recommender Systems
Lucas Vinh Tran, Yi Tay, Shuai Zhang, Gao Cong, Xiaoli Li.
WSDM 2020. The 13th ACM International Conference on Web Search and Data Mining.
Best Paper Award Runner-up.
Quaternion Knowledge Graph Embeddings
Shuai Zhang, Yi Tay, Lina Yao, Qi Liu.
NeurIPS 2019. Thirty-third Conference on Neural Information Processing Systems code.
Lightweight and Efficient Neural Natural Language Processing with Quaternion Networks
Yi Tay, Aston Zhang, Anh Tuan Luu, Jinfeng Rao, Shuai Zhang, Shuohang Wang, Jie Fu and Siu Cheung Hui.
ACL 2019. The 57th Annual Meeting of the Association for Computational Linguistics code.
[* denotes equal contribution]