publications

Pretraining & Scaling

1. Preprint

   What Is a Repeated Token Worth? The Scaling Geometry of Multi-Epoch Pretraining

   Yekun Chai, and Haoyi Xiong

   arXiv preprint arXiv:2610.05591, Oct 2026

   Abs TL;DR HTML PDF

   The cost of repeating data follows one variable: extra epochs divided by unique tokens per parameter. Repeats keep most of their value until a critical epoch count set by tokens per parameter, not model size; with unique data fixed, the compute-optimal run grows model size and epochs together up to that point.

   As pretraining increasingly repeats data, every run faces three questions: how many epochs to take, how that number should change with model size, and whether anything besides the epoch count matters. We answer them by pricing a repeated token against two references: one epoch on the same data, which gives its value, and fresh data at equal compute, which gives its cost. Against fresh data, the cost of repetition follows a single variable, the number of extra epochs divided by the unique tokens per parameter. Against the same data, a second epoch is worth nearly as much as a fresh one, and repeated tokens fall to half the value of fresh ones after a critical epoch count that grows with the training budget per parameter but hardly with model size. With unique data fixed, the predicted compute-optimal run grows model size and epochs together until loss stops improving, near the critical epoch count. The same variable accounts for the direction of size trends that appear to conflict: larger models tolerate fewer epochs when the corpus is fixed, from about 15 at 127M to 4 at 2B parameters, but not when unique data grow with the model. Counts alone do not determine loss: at identical counts, replaying shards consecutively raises loss by up to 0.46 bits per byte, concentrating repeats on fewer samples also raises it, lower-entropy sources degrade faster with repetition, and re-tokenizing repeats helps only under heavy repetition. These results offer an empirical guide to pretraining when unique data, rather than compute, are the binding constraint.

2. Tech Report

   ERNIE 5.0 Technical Report

   Baidu ERNIE

   arXiv, Feb 2026

   TL;DR PDF

   Text, image, audio and video trained from scratch as one autoregressive model; generation stays in discrete token space instead of handing off to a separate decoder.

3. Tech Report

   ERNIE 4.5 Technical Report

   Baidu ERNIE

   Jun 2025

   TL;DR PDF Models

   A multimodal MoE that shares parameters across modalities while keeping dedicated ones for each, adding vision without paying for it in text quality.

4. COLING-Industry

   Aurora-M: Open Source Continual Pre-training for Multilingual Language and Code

   Taishi Nakamura, Mayank Mishra, Simone Tedeschi, Yekun Chai, Jason T. Stillerman, Felix Friedrich, and 35 more authors

   In Proceedings of the 31st International Conference on Computational Linguistics: Industry Track, Jan 2025

   Abs TL;DR HTML PDF Blog Code

   Continual pretraining past 2T tokens without catastrophic forgetting, and the first open multilingual model tuned on human-reviewed safety instructions.

   Pretrained language models are integral part of AI applications, but their high computational cost for training limits accessibility. Initiatives such as Bloom and StarCoder aim to democratize access to pretrained models for collaborative community development. Despite these efforts, such models encounter challenges such as limited multilingual capabilities, risks of catastrophic forgetting during continual pretraining, and the high costs of training models from scratch, alongside the need to align with AI safety standards and regulatory frameworks. This paper presents Aurora-M, a 15B parameter multilingual open-source model trained on English, Finnish, Hindi, Japanese, Vietnamese, and code. Continually pretrained from StarCoderPlus on 435B additional tokens, Aurora-M surpasses 2T tokens in total training token count. It is the first open-source multilingual model fine-tuned on human-reviewed safety instructions, thus aligning its development not only with conventional red-teaming considerations, but also with the specific concerns articulated in the Biden-Harris Executive Order on the Safe, Secure, and Trustworthy Development and Use of Artificial Intelligence. We evaluate Aurora-M across a wide range of tasks and languages, showcasing its robustness against catastrophic forgetting and its superior performance in multilingual settings, particularly in safety evaluations. We open-source Aurora-M and its variants to encourage responsible open-source development of large language models at https://huggingface.co/aurora-m.

5. EMNLP

   Autoregressive Pre-Training on Pixels and Texts

   Yekun Chai, Qingyi Liu^{^}, Jingwu Xiao^{^}, Shuohuan Wang, Yu Sun, and Hua Wu

   In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, Nov 2024

   Abs TL;DR HTML PDF Code Poster Slides

   Tokenizer-free language modeling: autoregressive pretraining on rendered text pixels, alone or mixed with text, and next-patch prediction scales without a vocabulary.

   The integration of visual and textual information represents a promising direction in the advancement of language models. In this paper, we explore the dual modality of language—both visual and textual—within an autoregressive framework, pre-trained on both document images and texts. Our method employs a multimodal training strategy, utilizing visual data through next patch prediction with a regression head and/or textual data through next token prediction with a classification head. We focus on understanding the interaction between these two modalities and their combined impact on model performance. Our extensive evaluation across a wide range of benchmarks shows that incorporating both visual and textual data significantly improves the performance of pixel-based language models. Remarkably, we find that a unidirectional pixel-based model trained solely on visual data can achieve comparable results to state-of-the-art bidirectional models on several language understanding tasks. This work uncovers the untapped potential of integrating visual and textual modalities for more effective language modeling. We release our code, data, and model checkpoints at https://github.com/ernie-research/pixelgpt.

6. EMNLP Oral

   On Training Data Influence of GPT Models

   Yekun Chai, Qingyi Liu^{^}, Shuohuan Wang, Yu Sun, Qiwei Peng, and Hua Wu

   In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, Nov 2024

   Abs TL;DR HTML PDF Code Poster Slides

   Simulates how individual training examples shape loss and metric trajectories during training, and generalizes to training data it has not seen.

   Oral

   Amidst the rapid advancements in generative language models, the investigation of how training data shapes the performance of GPT models is still emerging. This paper presents GPTfluence, a novel approach that leverages a featurized simulation to assess the impact of training examples on the training dynamics of GPT models. Our approach not only traces the influence of individual training instances on performance trajectories, such as loss and other key metrics, on targeted test points but also enables a comprehensive comparison with existing methods across various training scenarios in GPT models, ranging from 14 million to 2.8 billion parameters, across a range of downstream tasks. Contrary to earlier methods that struggle with generalization to new data, GPTfluence introduces a parameterized simulation of training dynamics, demonstrating robust generalization capabilities to unseen training data. This adaptability is evident across both fine-tuning and instruction-tuning scenarios, spanning tasks in natural language understanding and generation. We make our code and data publicly available at https://github.com/ernie-research/gptfluence.

7. Tech Report

   StarCoder 2 and The Stack v2: The Next Generation

   Anton Lozhkov , Raymond Li, Loubna Ben Allal, Federico Cassano, Joel Lamy-Poirier, Nouamane Tazi, and 60 more authors

   Feb 2024

   TL;DR PDF Blog Code

   Treats repository context, provenance, and licensing as part of the code-pretraining recipe for trillion-token training, rather than metadata discarded during curation.

8. ACL-Findings

   ERNIE-Code: Beyond English-Centric Cross-lingual Pretraining for Programming Languages

   Yekun Chai, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, and Hua Wu

   In Findings of the Association for Computational Linguistics: ACL 2023, Jul 2023

   TL;DR HTML PDF Code

   Pretraining code and natural language jointly before that was standard: 116 natural languages and 6 programming languages in one model.

Post-training & Agents

1. Preprint

   Finding the Move Is Not Winning the Game: XiangqiBench for Closed-Loop Evaluation of LLM Agents

   Yekun Chai, Qiwei Peng, and Haoyi Xiong

   arXiv preprint arXiv:2610.02425, Oct 2026

   Abs TL;DR HTML PDF Code

   Naming the right move is not converting it: across 8,568 endgame trajectories against an engine defender, frontier LLM agents show large gaps in conversion, consistency, and simulation.

   Static evaluations credit a language model for naming the right move, but an agent must carry a plan through to a verified outcome while an opponent responds. We introduce XiangqiBench, an executable benchmark that measures this difference in Chinese chess: starting from 119 tactical endgames with forced mates supported by engine or checks-only search, an LLM agent must deliver checkmate against an engine defender. An interactive REPL interface separates real moves, state queries, and forward simulation, and we record 8,568 multi-turn trajectories from 12 frontier LLMs under two observation protocols. Three signals that look like competence each overstate closed-loop success. (i) The Conversion Gap: models play the stored reference first move in 26.1% of Sighted trials, yet only 13.9% of these trials end in a win. (ii) The Consistency Gap: the leading model reaches 38.7% pass@3 but only 5.9% pass^3, winning all three trials on 7 of the 46 positions it ever wins. (iii) The Simulation Gap: 32.3% of accepted simulation calls stop on an illegal move, and in 49.3% of comparable cases the real defender replies differently from the line the agent simulated; self-authored rollouts check legality but cannot anticipate the opponent. Finding the move is not winning the game: agent evaluations should score closed-loop outcomes and report reliability alongside coverage.

2. ACL

   Curiosity-Driven Reinforcement Learning from Human Feedback

   Haoran Sun^{*^}, Yekun Chai^*†, Shuohuan Wang, Yu Sun, Hua Wu , and Haifeng Wang

   In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Jul 2025

   Abs TL;DR HTML PDF Code

   Intrinsic-reward exploration for RLHF that preserves output diversity without sacrificing alignment.

   Reinforcement learning from human feedback (RLHF) has proven effective in aligning large language models (LLMs) with human preferences, but often at the cost of reduced output diversity. This trade-off between diversity and alignment quality remains a significant challenge. Drawing inspiration from curiosity-driven exploration in reinforcement learning, we introduce curiosity-driven RLHF (CD-RLHF), a framework that incorporates intrinsic rewards for novel states, alongside traditional sparse extrinsic rewards, to optimize both output diversity and alignment quality. We demonstrate the effectiveness of CD-RLHF through extensive experiments on a range of tasks, including text summarization and instruction following. Our approach achieves significant gains in diversity on multiple diversity-oriented metrics while maintaining alignment with human preferences comparable to standard RLHF. We will make our code publicly available.

3. ICLR

   MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions

   Yekun Chai^*† , Haoran Sun^{*^}, Huang Fang, Shuohuan Wang, Yu Sun, and Hua Wu

   In The Thirteenth International Conference on Learning Representations, Apr 2025

   TL;DR HTML PDF Code

   Temporal abstraction for RLHF: treating token sequences as macro-actions improves credit assignment in long-horizon generation.

4. ICLR Spotlight

   Tool-Augmented Reward Modeling

   Lei Li^{*^}, Yekun Chai^*†, Shuohuan Wang, Yu Sun, Hao Tian, Ningyu Zhang, and 1 more author

   In The Twelfth International Conference on Learning Representations, May 2024

   TL;DR HTML PDF Code Poster Slides

   Introduced the first tool-augmented reward model, which calls external tools to verify factual claims before assigning a reward.

   Spotlight

Analysis, Evaluation & Efficiency

1. EMNLP

   CodeMixBench: Evaluating Code-Mixing Capabilities of LLMs Across 18 Languages

   Yilun Yang^{^}, and Yekun Chai^†

   In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, Nov 2025

   Abs TL;DR HTML PDF Code

   Eighteen languages and eight tasks for code-mixing. Models still drop when the switch crosses language families.

   Code-mixing, the practice of switching between languages within a conversation, poses unique challenges for traditional NLP. Existing benchmarks are limited by their narrow language pairs and tasks, failing to adequately assess large language models’ (LLMs) code-mixing abilities. Despite the recognized importance of code-mixing for multilingual users, research on LLMs in this context remains sparse. Additionally, current techniques for synthesizing code-mixed data are underdeveloped to generate code-mixing. In response, we introduce CodeMixBench, a comprehensive benchmark covering eight tasks, including three specific to LLMs and five traditional NLP tasks, and 18 languages across seven language families. We also propose a new method for generating large-scale synthetic code-mixed texts by combining word substitution with GPT-4 prompting. Our evaluation reveals consistent underperformance of LLMs on code-mixed datasets involving different language families. Enhancements in training data size, model scale, and few-shot learning could improve their performance. The code and dataset are available at https://github.com/Jeromeyluck/CodeMixBench.

2. EMNLP

   Understanding Subword Compositionality of Large Language Models

   Qiwei Peng, Yekun Chai, and Anders Søgaard

   In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, Nov 2025

   TL;DR HTML PDF

   Five LLM families fall into three groups by how they compose subwords into words, visible in layer-wise structural similarity and in sensitivity to word form.

3. EMNLP

   Debiasing Multilingual LLMs in Cross-lingual Latent Space

   Qiwei Peng, Guimin Hu, Yekun Chai, and Anders Søgaard

   In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, Nov 2025

   Abs TL;DR HTML PDF

   Debiasing in an autoencoder-learned cross-lingual latent space, rather than on raw LLM representations, reduces bias by up to 65% and transfers better across languages.

   Debiasing techniques such as SentDebias aim to reduce bias in large language models (LLMs). Previous studies have evaluated their cross-lingual transferability by directly applying these methods to LLM representations, revealing their limited effectiveness across languages. In this work, we therefore propose to perform debiasing in a joint latent space rather than directly on LLM representations. We construct a well-aligned cross-lingual latent space using an autoencoder trained on parallel TED talk scripts. Our experiments with Aya-expanse and two debiasing techniques across four languages (English, French, German, Dutch) demonstrate that a) autoencoders effectively construct a well-aligned cross-lingual latent space, and b) applying debiasing techniques in the learned cross-lingual latent space significantly improves both the overall debiasing performance and cross-lingual transferability.

4. EMNLP-Findings

   EvolKV: Evolutionary KV Cache Compression for LLM Inference

   Bohan Yu^{^}, and Yekun Chai^†

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

   Abs TL;DR HTML PDF

   Layer-wise evolutionary search for KV-cache budgets. Task-driven compression can beat the full cache at a fraction of the memory.

   Existing key-value (KV) cache compression methods typically rely on heuristics, such as uniform cache allocation across layers or static eviction policies, however, they ignore the critical interplays among layer-specific feature patterns and task performance, which can lead to degraded generalization. In this paper, we propose EvolKV, an adaptive framework for layer-wise, task-driven KV cache compression that jointly optimizes the memory efficiency and task performance. By reformulating cache allocation as a multi-objective optimization problem, EvolKV leverages evolutionary search to dynamically configure layer budgets while directly maximizing downstream performance. Extensive experiments on 11 tasks demonstrate that our approach outperforms all baseline methods across a wide range of KV cache budgets on long-context tasks and surpasses heuristic baselines by up to 7 percentage points on GSM8K. Notably, EvolKV achieves superior performance over the full KV cache setting on code completion while utilizing only 1.5% of the original budget, suggesting the untapped potential in learned compression strategies for KV cache budget allocation.

5. EMNLP-Findings

   Tokenization Falling Short: On Subword Robustness in Large Language Models

   Yekun Chai , Yewei Fang, Qiwei Peng, and Xuhong Li

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

   Abs TL;DR HTML PDF Code Poster Slides

   Scaling reduces the curse of tokenization but never removes it, and subword regularization such as BPE-dropout buys more typo robustness than more parameters.

   Language models typically tokenize raw text into sequences of subword identifiers from a predefined vocabulary, a process inherently sensitive to typographical errors, length variations, and largely oblivious to the internal structure of tokens—issues we term *the curse of tokenization*. In this study, we delve into these drawbacks and demonstrate that large language models (LLMs) remain susceptible to these problems. This study systematically investigates these challenges and their impact on LLMs through three critical research questions: (1) complex problem solving, (2) token structure probing, and (3) resilience to typographical variation. Our findings reveal that scaling model parameters can mitigate the issue of tokenization; however, LLMs still suffer from biases induced by typos and other text format variations. Our experiments show that subword regularization such as BPE-dropout can mitigate this issue. We release our evaluation code and data at https://github.com/FloatAI/TKEval.

6. ICML

   GiLOT: Interpreting Generative Language Models via Optimal Transport

   Xuhong Li^*, Jiamin Chen^*, Yekun Chai^*, and Haoyi Xiong

   In Forty-first International Conference on Machine Learning, Jul 2024

   TL;DR HTML PDF Code

   Attribution over the full next-token distribution via optimal transport, rather than scoring one sampled continuation.

7. LREC-COLING

   HumanEval-XL: A Multilingual Code Generation Benchmark for Cross-lingual Natural Language Generalization

   Qiwei Peng^*, Yekun Chai^*, and Xuhong Li

   In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation (LREC-COLING 2024), May 2024

   Abs TL;DR PDF Code Poster Slides

   23 natural languages paired with 12 programming languages in parallel, so understanding the prompt language can be told apart from writing the target language.

   Large language models (LLMs) have made significant progress in generating codes from textual prompts. However, existing benchmarks have mainly concentrated on translating English prompts to multilingual codes or have been constrained to very limited natural languages (NLs). These benchmarks have overlooked the vast landscape of massively multilingual NL to multilingual code, leaving a critical gap in the evaluation of multilingual LLMs. In response, we introduce HumanEval-XL, a massively multilingual code generation benchmark specifically crafted to address this deficiency. HumanEval-XL establishes connections between 23 NLs and 12 programming languages (PLs), and comprises of a collection of 22,080 prompts with an average of 8.33 test cases. By ensuring parallel data across multiple NLs and PLs, HumanEval-XL offers a comprehensive evaluation platform for multilingual LLMs, allowing the assessment of the understanding of different NLs. Our work serves as a pioneering step towards filling the void in evaluating NL generalization in the area of multilingual code generation. We make our evaluation code and data publicly available at https://github.com/FloatAI/HumanEval-XL.

8. NeurIPS Datasets and Benchmarks

   M⁴: A Unified XAI Benchmark for Faithfulness Evaluation of Feature Attribution Methods across Metrics, Modalities and Models

   Xuhong Li, Mengnan Du, Jiamin Chen, Yekun Chai, Himabindu Lakkaraju, and Haoyi Xiong

   In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, Dec 2023

   TL;DR PDF Code Poster Slides

   Scores feature-attribution methods with one set of faithfulness metrics across images and text, CNNs and Transformers, with metrics grouped by the ground truth they require.

   Datasets and Benchmarks

Earlier Work

1. ICASSP Oral

   Improved Training of Mixture-of-Experts Language GANs

   Yekun Chai, Qiyue Yin , and Junge Zhang

   In 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Jun 2023

   TL;DR PDF

   A mixture-of-experts generator with feature-statistics alignment gives language GANs a richer training signal than real-or-fake discrimination alone.

   Oral

2. EMNLP-Findings

   Clip-Tuning: Towards Derivative-free Prompt Learning with a Mixture of Rewards

   Yekun Chai, Shuohuan Wang, Yu Sun, Hao Tian, Hua Wu , and Haifeng Wang

   In Findings of the Association for Computational Linguistics: EMNLP 2022, Dec 2022

   Abs TL;DR PDF

   Derivative-free prompt learning with a mixture of rewards from dropout-thinned subnetworks of one frozen PLM, matching gradient-based prompt tuning in few-shot settings.

   Derivative-free prompt learning has emerged as a lightweight alternative to prompt tuning, which only requires model inference to optimize the prompts. However, existing work did not take full advantage of the over-parameterized characteristics of large pre-trained language models (PLMs). In this paper, we propose Clip-Tuning, a simple yet effective method that adopts diverse frozen “thinned” networks of PLMs to obtain *a mixture of rewards* and thus advance the derivative-free prompt learning. The thinned networks consist of all the hidden units that survive a stationary dropout strategy, whose inference predictions reflect an ensemble of partial views over prompted training samples. Our method outperforms previous gradient-free prompt learning methods and achieves parity with gradient-based counterparts on seven language understanding benchmarks under few-shot settings.

3. EMNLP-Findings

   Counter-Contrastive Learning for Language GANs

   Yekun Chai , Haidong Zhang, Qiyue Yin , and Junge Zhang

   In Findings of the Association for Computational Linguistics: EMNLP 2021, Nov 2021

   Abs TL;DR PDF Poster Slides

   Replaces the binary real-or-fake signal in language GANs with a counter-contrastive objective that pulls generated and real representations together and keeps the discriminator from overtraining.

   Generative Adversarial Networks (GANs) have achieved great success in image synthesis, but have proven to be difficult to generate natural language. Challenges arise from the uninformative learning signals passed from the discriminator. In other words, the poor learning signals limit the learning capacity for generating languages with rich structures and semantics. In this paper, we propose to adopt the counter-contrastive learning (CCL) method to support the generator’s training in language GANs. In contrast to standard GANs that adopt a simple binary classifier to discriminate whether a sample is real or fake, we employ a counter-contrastive learning signal that advances the training of language synthesizers by (1) pulling the language representations of generated and real samples together and (2) pushing apart representations of real samples to compete with the discriminator and thus prevent the discriminator from being overtrained. We evaluate our method on both synthetic and real benchmarks and yield competitive performance compared to previous GANs for adversarial sequence generation.

4. ACL

   Highway Transformer: Self-Gating Enhanced Self-Attentive Networks

   Yekun Chai, Shuo Jin, and Xinwen Hou

   In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, Jul 2020

   Abs TL;DR PDF Code Slides

   Adds LSTM-style self-gating units beside self-attention and feed-forward layers, speeding up convergence, with gates in shallow layers helping most.

   Self-attention mechanisms have made striking state-of-the-art (SOTA) progress in various sequence learning tasks, standing on the multi-headed dot product attention by attending to all the global contexts at different locations. Through a pseudo information highway, we introduce a gated component self-dependency units (SDU) that incorporates LSTM-styled gating units to replenish internal semantic importance within the multi-dimensional latent space of individual representations. The subsidiary content-based SDU gates allow for the information flow of modulated latent embeddings through skipped connections, leading to a clear margin of convergence speed with gradient descent algorithms. We may unveil the role of gating mechanism to aid in the context-based Transformer modules, with hypothesizing that SDU gates, especially on shallow layers, could push it faster to step towards suboptimal points during the optimization process.

* equal contribution    † corresponding author / project lead    ^ mentored