publications

2026

1. ICML

   MineDraft: A Framework for Batch Parallel Speculative Decoding

   Zhenwei Tang, Arun Verma, Zijian Zhou, and 4 more authors

   In International Conference on Machine Learning (ICML), 2026

   Abs arXiv Bib HTML Code

   Speculative decoding (SD) accelerates large language model inference by using a smaller draft model to propose draft tokens that are subsequently verified by a larger target model. We introduce MineDraft, a framework that overlaps drafting with verification across two batches of requests, achieving improvements of up to 75% in throughput and 39% in end-to-end latency compared to standard approaches.

   @inproceedings{tang2026minedraft,
     title = {MineDraft: A Framework for Batch Parallel Speculative Decoding},
     author = {Tang, Zhenwei and Verma, Arun and Zhou, Zijian and Wu, Zhaoxuan and Prakash, Alok and Rus, Daniela and Low, Bryan Kian Hsiang},
     booktitle = {International Conference on Machine Learning (ICML)},
     year = {2026},
     url = {https://arxiv.org/abs/2603.18016},
   }

2. COLM

   

   CORAL: Towards Autonomous Multi-Agent Evolution for Open-Ended Discovery

   Ao Qu, Han Zheng, Zijian Zhou, and 14 more authors

   In Conference on Language Modeling (COLM), 2026

   Abs arXiv Bib HTML

   We present CORAL, a framework enabling autonomous multi-agent evolution for open-ended discovery tasks. Rather than relying on predetermined exploration strategies, CORAL employs long-running agents that explore, reflect, and collaborate through shared persistent memory, asynchronous multi-agent execution, and heartbeat-based interventions. Across mathematical, algorithmic, and optimization problems, CORAL achieves 3-10 times higher improvement rates with far fewer evaluations than fixed evolutionary search baselines. Notably, four co-evolving agents improve kernel engineering performance from 1363 to 1103 cycles.

   @inproceedings{qu2026coral,
     title = {CORAL: Towards Autonomous Multi-Agent Evolution for Open-Ended Discovery},
     author = {Qu, Ao and Zheng, Han and Zhou, Zijian and Yan, Yihao and Tang, Yihong and Ong, Shao Yong and Hong, Fenglu and Zhou, Kaichen and Jiang, Chonghe and Kong, Minwei and Zhu, Jiacheng and Jiang, Xuan and Li, Sirui and Wu, Cathy and Low, Bryan Kian Hsiang and Zhao, Jinhua and Liang, Paul Pu},
     booktitle = {Conference on Language Modeling (COLM)},
     year = {2026},
     archiveprefix = {arXiv},
     primaryclass = {cs.AI},
     url = {https://arxiv.org/abs/2604.01658},
   }

3. ICLR

   

   MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents

   Zijian Zhou, Ao Qu, Zhaoxuan Wu, and 6 more authors

   In International Conference on Learning Representations (ICLR), 2026

   Abs arXiv Bib HTML

   Modern language agents must operate over long-horizon, multi-turn interactions, where they retrieve external information, adapt to observations, and answer interdependent queries. Yet, most LLM systems rely on full-context prompting, appending all past turns regardless of their relevance. This leads to unbounded memory growth, increased computational costs, and degraded reasoning performance on out-of-distribution input lengths. We introduce MEM1, an end-to-end reinforcement learning framework that enables agents to operate with constant memory across long multi-turn tasks. At each turn, MEM1 updates a compact shared internal state that jointly supports memory consolidation and reasoning. Experiments across three domains show that MEM1-7B improves performance by 3.5x while reducing memory usage by 3.7x compared to Qwen2.5-14B-Instruct on a 16-objective multi-hop QA task, and generalizes beyond the training horizon.

   @inproceedings{zhou2025mem1,
     title = {MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents},
     author = {Zhou, Zijian and Qu, Ao and Wu, Zhaoxuan and Kim, Sunghwan and Prakash, Alok and Rus, Daniela and Zhao, Jinhua and Low, Bryan Kian Hsiang and Liang, Paul Pu},
     booktitle = {International Conference on Learning Representations (ICLR)},
     year = {2026},
     url = {https://arxiv.org/abs/2506.15841},
   }

2025

1. ACL

   

   TETRIS: Optimal Draft Token Selection for Batch Speculative Decoding

   Zhaoxuan Wu, Zijian Zhou, Arun Verma, and 3 more authors

   In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL), 2025

   Abs arXiv Bib HTML

   We propose TETRIS, a novel method that optimizes the total throughput of batch speculative decoding in multi-request settings. Unlike existing methods that optimize for a single request or a group of requests as a whole, TETRIS actively selects the most promising draft tokens (for every request in a batch) to be accepted when verified in parallel, resulting in fewer rejected tokens and hence less wasted computing resources. We show theoretically and empirically that TETRIS outperforms baseline speculative decoding and existing methods that dynamically select draft tokens, leading to a more efficient batch inference in LLMs.

   @inproceedings{wu-etal-2025-tetris,
     title = {TETRIS: Optimal Draft Token Selection for Batch Speculative Decoding},
     author = {Wu, Zhaoxuan and Zhou, Zijian and Verma, Arun and Prakash, Alok and Rus, Daniela and Low, Bryan Kian Hsiang},
     booktitle = {Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL)},
     address = {Vienna, Austria},
     year = {2025},
     publisher = {Association for Computational Linguistics},
     url = {https://arxiv.org/abs/2502.15197},
   }

2. EMNLP

   

   Uncovering Scaling Laws for Large Language Models via Inverse Problems

   Arun Verma, Zhaoxuan Wu, Zijian Zhou, and 15 more authors

   In Findings of the Association for Computational Linguistics: EMNLP, 2025

   Abs DOI Bib HTML

   Large Language Models (LLMs) are large-scale pretrained models that have achieved remarkable success across diverse domains. These successes have been driven by unprecedented complexity and scale in both data and computations. However, due to the high costs of training such models, brute-force trial-and-error approaches to improve LLMs are not feasible. Inspired by the success of inverse problems in uncovering fundamental scientific laws, this position paper advocates that inverse problems can also be used to efficiently uncover scaling laws that guide the building of LLMs to achieve a desirable performance with significantly better cost-effectiveness.

   @inproceedings{verma-etal-2025-uncovering,
     title = {Uncovering Scaling Laws for Large Language Models via Inverse Problems},
     author = {Verma, Arun and Wu, Zhaoxuan and Zhou, Zijian and Lin, Xiaoqiang and Chen, Zhiliang and Sim, Rachael Hwee Ling and Qiao, Rui and Wang, Jingtan and Bui, Nhung and Niu, Xinyuan and Hu, Wenyang and Lau, Gregory Kang Ruey and Khoo, Zi-Yu and Zhao, Zitong and Xu, Xinyi and Hemachandra, Apivich and Ng, See-Kiong and Low, Bryan Kian Hsiang},
     booktitle = {Findings of the Association for Computational Linguistics: EMNLP},
     address = {Suzhou, China},
     year = {2025},
     publisher = {Association for Computational Linguistics},
     pages = {25197--25211},
     doi = {10.18653/v1/2025.findings-emnlp.1373},
     url = {https://aclanthology.org/2025.findings-emnlp.1373/},
   }

2024

1. NeurIPS

   

   DETAIL: Task DEmonsTration Attribution for Interpretable In-context Learning

   Zijian Zhou, Xiaoqiang Lin, Xinyi Xu, and 3 more authors

   In Advances in Neural Information Processing Systems 37 (NeurIPS), 2024

   Abs DOI arXiv Bib HTML

   In-context learning (ICL) allows transformer-based language models to quickly learn a specific task with a few "task demonstrations" without updating their parameters. Taking the viewpoint of recent works showing that transformers learn in context by formulating an internal optimizer, we propose an influence function-based attribution technique, DETAIL, that addresses the specific characteristics of ICL. We show how DETAIL can help improve model performance through demonstration reordering and curation, and prove its wide applicability by showing attribution scores obtained on white-box models are transferable to black-box models.

   @inproceedings{zhou-etal-2024-detail,
     title = {DETAIL: Task DEmonsTration Attribution for Interpretable In-context Learning},
     author = {Zhou, Zijian and Lin, Xiaoqiang and Xu, Xinyi and Prakash, Alok and Rus, Daniela and Low, Bryan Kian Hsiang},
     booktitle = {Advances in Neural Information Processing Systems 37 (NeurIPS)},
     address = {Vancouver, Canada},
     year = {2024},
     publisher = {Curran Associates, Inc.},
     doi = {10.48550/arXiv.2405.14899},
     url = {https://arxiv.org/abs/2405.14899},
   }

2. EMNLP

   

   Position Paper: Data-Centric AI in the Age of Large Language Models

   Xinyi Xu, Zhaoxuan Wu, Rui Qiao, and 16 more authors

   In Findings of the Association for Computational Linguistics: EMNLP, 2024

   Abs DOI arXiv Bib HTML

   This position paper proposes a data-centric viewpoint of AI research, focusing on large language models (LLMs). We identify four specific scenarios centered around data, covering data-centric benchmarks and data curation, data attribution, knowledge transfer, and inference contextualization. In each scenario, we underscore the importance of data, highlight promising research directions, and articulate the potential impacts on the research community and, where applicable, the society as a whole.

   @inproceedings{xu-etal-2024-position,
     title = {Position Paper: Data-Centric AI in the Age of Large Language Models},
     author = {Xu, Xinyi and Wu, Zhaoxuan and Qiao, Rui and Verma, Arun and Shu, Yao and Wang, Jingtan and Niu, Xinyuan and He, Zhenfeng and Chen, Jiangwei and Zhou, Zijian and Lau, Gregory Kang Ruey and Dao, Hieu and Agussurja, Lucas and Sim, Rachael Hwee Ling and Lin, Xiaoqiang and Hu, Wenyang and Dai, Zhongxiang and Koh, Pang Wei and Low, Bryan Kian Hsiang},
     booktitle = {Findings of the Association for Computational Linguistics: EMNLP},
     address = {Miami, Florida, USA},
     year = {2024},
     publisher = {Association for Computational Linguistics},
     pages = {11895--11913},
     doi = {10.18653/v1/2024.findings-emnlp.695},
     url = {https://aclanthology.org/2024.findings-emnlp.695/},
   }

3. arXiv

   Data Value Estimation on Private Gradients

   Zijian Zhou, Xinyi Xu, Daniela Rus, and 1 more author

   2024

   Abs DOI arXiv Bib HTML

   For gradient-based machine learning methods, the de facto differential privacy technique is perturbing the gradients with random Gaussian noise. Can existing data valuation methods still be used when DP is enforced via gradient perturbations? We show that the answer is no with the default approach of injecting i.i.d. random noise: the estimation uncertainty of the data value estimation paradoxically linearly scales with more estimation budget. We propose to instead inject carefully correlated noise to provably remove the linear scaling of estimation uncertainty w.r.t. the budget.

   @misc{zhou2024data,
     title = {Data Value Estimation on Private Gradients},
     author = {Zhou, Zijian and Xu, Xinyi and Rus, Daniela and Low, Bryan Kian Hsiang},
     year = {2024},
     archiveprefix = {arXiv},
     primaryclass = {cs.LG},
     doi = {10.48550/arXiv.2412.17008},
     url = {https://arxiv.org/abs/2412.17008},
   }

2023

1. AAAI

   

   Probably Approximate Shapley Fairness with Applications in Machine Learning

   Zijian Zhou, Xinyi Xu, Rachael Hwee Ling Sim, and 2 more authors

   In Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence (AAAI), Oral Presentation, 2023

   Oral Abs DOI arXiv Bib HTML

   Oral presentation at AAAI 2023

   The Shapley value (SV) is adopted in various scenarios in machine learning (ML), including data valuation, agent valuation, and feature attribution, as it satisfies their fairness requirements. However, as exact SVs are infeasible to compute in practice, SV estimates are approximated instead. This approximation step raises an important question: do the SV estimates preserve the fairness guarantees of exact SVs? We generalise Shapley fairness to probably approximate Shapley fairness and propose fidelity score, a metric to measure the variation of SV estimates. We also propose a novel greedy active estimation (GAE) algorithm that achieves a better fairness guarantee than the de facto Monte-Carlo estimation.

   @inproceedings{zhou-etal-2023-probably,
     title = {Probably Approximate Shapley Fairness with Applications in Machine Learning},
     author = {Zhou, Zijian and Xu, Xinyi and Sim, Rachael Hwee Ling and Foo, Chuan Sheng and Low, Bryan Kian Hsiang},
     booktitle = {Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence (AAAI), Oral Presentation},
     pages = {5910--5918},
     year = {2023},
     address = {Washington, DC, USA},
     publisher = {Association for the Advancement of Artificial Intelligence},
     doi = {10.1609/aaai.v37i5.25732},
     url = {https://ojs.aaai.org/index.php/AAAI/article/view/25732},
   }