TANDEM: Bi-Level Data Mixture Optimization with Twin Networks

The capabilities of large language models (LLMs) significantly depend on training data drawn from various domains. Optimizing domain-specific mixture ratios can be modeled as a bi-level optimization problem, which we simplify into a single-level penalized form and solve with twin networks: a proxy model trained on primary data and a dynamically updated reference model trained with additional data. Our proposed method, Twin Networks for bi-level DatA mixturE optiMization (TANDEM), measures the data efficacy through the difference between the twin models and up-weights domains that benefit more from the additional data. TANDEM provides theoretical guarantees and wider applicability, compared to prior approaches. Furthermore, our bi-level perspective suggests new settings to study domain reweighting such as data-restricted scenarios and supervised fine-tuning, where optimized mixture ratios significantly improve the performance. Extensive experiments validate TANDEM's effectiveness in all scenarios.

Authors

• Jiaxing Wang JD.com
• Deping Xiang
• Jin Xu Pazhou Lab, Guangzhou South China University of Technology
• Mingyang Yi
• Guoqiang Gong
• Zicheng Zhang University of the Chinese Academy of Sciences
• Haoran Li Institute of Automation, Chinese Academy of Sciences & School of Artificial Intelligence, University of Chinese Academy of Sciences
• Pengzhang Liu
• Zhen Chen The Chinese University of Hong Kong
• Ke Zhang Hangzhou Dianzi University
• Ju Fan Renmin University of China
• Qixia Jiang Tsinghua University

Keywords

No keywords are indexed for this paper.