{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/aioli-a-unified-optimization-framework-for","title":"Aioli: A Unified Optimization Framework for Language Model Data Mixing","arxiv_id":"2411.05735","date":"2024-11-08","proceeding":null,"authors":["Mayee F. Chen","Michael Y. Hu","Nicholas Lourie","Kyunghyun Cho","Christopher Ré"],"abstract":"Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over training runs to dynamically updating proportions throughout training. Surprisingly, we find that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity. To understand this inconsistency, we unify existing methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law -- an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of a method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we observe. Using this insight, we derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.27 test perplexity points, whereas existing methods fail to consistently beat stratified sampling, doing up to 6.9 points worse. Moreover, in a practical setting where proportions are learned on shorter runs due to computational constraints, Aioli can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.012 test perplexity points.","url_abs":"https://arxiv.org/abs/2411.05735v2","url_pdf":"https://arxiv.org/pdf/2411.05735v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"aioli-a-unified-optimization-framework-for","repo_url":"https://github.com/hazyresearch/aioli","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[{"task_slug":"language-modeling","task_name":"Language Modeling"},{"task_slug":"language-modelling","task_name":"Language Modelling"},{"task_slug":"math","task_name":"Math"}],"methods":[{"method_slug":"set","method_name":"SET"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2411.05735","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2411.05735"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"deterministic:regex_extraction","url":"https://github.com/HazyResearch/aioli","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/hazyresearch/aioli","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran":5,"ran_draft_wrong":4,"ran_honours":1,"unverified":21},"by_repo_kind":{"official":{"samples":31,"ran":10,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"aa9e619f3278e4bc","entry":"AbstractTrainer","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"aa9e619f3278e4bc"}},{"code_sha256_prefix":"3cae68a5efb29625","entry":"AioliTrainer","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"3cae68a5efb29625"}},{"code_sha256_prefix":"58161c1847700499","entry":"LinearWarmupCosineLR","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"58161c1847700499"}},{"code_sha256_prefix":"a6ada8ea3d1c8d08","entry":"StringDataCollator","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a6ada8ea3d1c8d08"}},{"code_sha256_prefix":"ac3ed26036cd8353","entry":"create_optimizer_scheduler","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ac3ed26036cd8353"}},{"code_sha256_prefix":"7ecff0f37380fa48","entry":"get_scheduler_extended","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"7ecff0f37380fa48"}},{"code_sha256_prefix":"2cc66e9c3b61d3ad","entry":"get_steps","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran_honours","verification_level":1,"contract_check":"HONOURS","metamorphic_tier":"well_formed","behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2cc66e9c3b61d3ad"}},{"code_sha256_prefix":"ceda5669cbbae013","entry":"get_tokenized_train_dataset","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ceda5669cbbae013"}},{"code_sha256_prefix":"34a317ad4275d585","entry":"get_tokenized_val_dataset","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"34a317ad4275d585"}},{"code_sha256_prefix":"2832a5cae93655cc","entry":"get_train_dataloader","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2832a5cae93655cc"}},{"code_sha256_prefix":"d251e579d9c1e9b5","entry":"ExtendedSchedulerType","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"d251e579d9c1e9b5"}},{"code_sha256_prefix":"838c84352eff715d","entry":"aggregate_task_category","repo":"HazyResearch/aioli","repo_kind":"official","path":"evaluator/utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/evaluator/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"838c84352eff715d"}},{"code_sha256_prefix":"cc3ad484cd8dd56c","entry":"calculate_r_squared","repo":"HazyResearch/aioli","repo_kind":"official","path":"analysis/law.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/analysis/law.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"cc3ad484cd8dd56c"}},{"code_sha256_prefix":"92cb091d2427e662","entry":"collate","repo":"HazyResearch/aioli","repo_kind":"official","path":"mixing_law_fitting.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/mixing_law_fitting.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"92cb091d2427e662"}},{"code_sha256_prefix":"9e72f9c47ff11092","entry":"cover_simplex","repo":"HazyResearch/aioli","repo_kind":"official","path":"dirichlet_cover.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/dirichlet_cover.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"9e72f9c47ff11092"}},{"code_sha256_prefix":"4ac03e5149b1ef69","entry":"fit_multi_obj_scaling_laws","repo":"HazyResearch/aioli","repo_kind":"official","path":"analysis/law.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/analysis/law.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4ac03e5149b1ef69"}},{"code_sha256_prefix":"ae58a6abc7dd4a7e","entry":"fit_scaling_laws","repo":"HazyResearch/aioli","repo_kind":"official","path":"analysis/law.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/analysis/law.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ae58a6abc7dd4a7e"}},{"code_sha256_prefix":"d0f3b712c535ab79","entry":"get_flat_grad","repo":"HazyResearch/aioli","repo_kind":"official","path":"trainer/doge_trainer.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/trainer/doge_trainer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"d0f3b712c535ab79"}},{"code_sha256_prefix":"e4fac505ef94582c","entry":"get_grad_dict","repo":"HazyResearch/aioli","repo_kind":"official","path":"trainer/doge_trainer.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/trainer/doge_trainer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"e4fac505ef94582c"}},{"code_sha256_prefix":"0cb39a42612d17ad","entry":"get_logger","repo":"HazyResearch/aioli","repo_kind":"official","path":"utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"0cb39a42612d17ad"}},{"code_sha256_prefix":"f262463f78c9ed6e","entry":"get_model_type","repo":"HazyResearch/aioli","repo_kind":"official","path":"utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"f262463f78c9ed6e"}},{"code_sha256_prefix":"f9c6848b0a94945a","entry":"get_slimpajama_6b","repo":"HazyResearch/aioli","repo_kind":"official","path":"dataset/slimpj_dataset.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/dataset/slimpj_dataset.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"f9c6848b0a94945a"}},{"code_sha256_prefix":"fb3cbc008efc4356","entry":"interleave_datasets","repo":"HazyResearch/aioli","repo_kind":"official","path":"dataset/slimpj_dataset.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/dataset/slimpj_dataset.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"fb3cbc008efc4356"}},{"code_sha256_prefix":"1f02cd5ca2dc84d5","entry":"interleave_per_batch_mapped","repo":"HazyResearch/aioli","repo_kind":"official","path":"dataset/slimpj_dataset.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/dataset/slimpj_dataset.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"1f02cd5ca2dc84d5"}},{"code_sha256_prefix":"9b58281c06357914","entry":"load_weight_diff_from_folder","repo":"HazyResearch/aioli","repo_kind":"official","path":"mixing_law_fitting.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/mixing_law_fitting.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"9b58281c06357914"}},{"code_sha256_prefix":"7a0383d468979671","entry":"load_weights_from_folder","repo":"HazyResearch/aioli","repo_kind":"official","path":"mixing_law_fitting.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/mixing_law_fitting.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"7a0383d468979671"}},{"code_sha256_prefix":"6b3f2ccf676d8314","entry":"make_output_dir","repo":"HazyResearch/aioli","repo_kind":"official","path":"utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"6b3f2ccf676d8314"}},{"code_sha256_prefix":"f919e9b193bdb7f9","entry":"parse_resumed_model_cpkt","repo":"HazyResearch/aioli","repo_kind":"official","path":"evaluator/utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/evaluator/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"f919e9b193bdb7f9"}},{"code_sha256_prefix":"ed16693f655ad564","entry":"save_aioli_matrices","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ed16693f655ad564"}},{"code_sha256_prefix":"2dec393904abb40d","entry":"save_loss","repo":"HazyResearch/aioli","repo_kind":"official","path":"evaluator/utils.py","file_url":"https://github.com/HazyResearch/aioli/blob/HEAD/evaluator/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2dec393904abb40d"}},{"code_sha256_prefix":"30c6dd51f43019b0","entry":"save_model_and_optimizer","repo":"hazyresearch/aioli","repo_kind":"official","path":"trainer/aioli_trainer.py","file_url":"https://github.com/hazyresearch/aioli/blob/HEAD/trainer/aioli_trainer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"30c6dd51f43019b0"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}