{"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/arxiv-2602-22681","title":"Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement","arxiv_id":"2602.22681","date":"2026-02-26","proceeding":null,"authors":["Shuchen Zhu","Rizhen Hu","Mingze Wang","Mou Sun","Xue Wang","Kun Yuan","Zaiwen Wen"],"abstract":"Pre-training Large Language Models requires immense computational resources, making optimizer efficiency essential. The optimization landscape is highly anisotropic, with loss reduction driven predominantly by progress along flat directions. While matrix-based optimizers such as Muon and SOAP leverage fine-grained curvature information to outperform AdamW, their updates tend toward isotropy -- relatively conservative along flat directions yet potentially aggressive along sharp ones. To address this limitation, we first establish a unified Riemannian Ordinary Differential Equation (ODE) framework that elucidates how common adaptive algorithms operate synergistically: the preconditioner induces a Riemannian geometry that mitigates ill-conditioning, while momentum serves as a Riemannian damping term that promotes convergence. Guided by these insights, we propose LITE, a generalized acceleration strategy that enhances training dynamics by applying larger Hessian damping coefficients and learning rates along flat trajectories. Extensive experiments demonstrate that LITE significantly accelerates both Muon and SOAP across diverse architectures (Dense, MoE), parameter scales (130M--1.3B), datasets (C4, Pile), and learning-rate schedules (cosine, warmup-stable-decay). Theoretical analysis confirms that LITE facilitates faster convergence along flat directions in anisotropic landscapes, providing a principled approach to efficient LLM pre-training. The code is available at https://github.com/SHUCHENZHU/LITE.","url_abs":"https://arxiv.org/abs/2602.22681","url_pdf":"https://arxiv.org/pdf/2602.22681","source":{"archive":null,"snapshot":"2025-07-28","note":"not in the Papers with Code archive (frozen at the snapshot)","row_kind":"graph","title_abstract_authors_date":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)"},"code_links":[],"tasks":[],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2602.22681","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2602.22681"}},"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":[{"mentioned_in_github":null,"is_official":null,"provenance":"deterministic:regex_extraction","mentioned_in_paper":null,"url":"https://github.com/SHUCHENZHU/LITE","reach":null}],"summary":{"ran_fixture":1,"ran_draft_wrong":1,"unverified":5},"by_repo_kind":{"found_in_text":{"samples":7,"ran":2,"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":7,"samples":[{"code_sha256_prefix":"f725bc2d76076485","entry":"apply_rotary_pos_emb","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"peft_pretraining/modeling_llama.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/peft_pretraining/modeling_llama.py","link_basis":"harvester_set","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"f725bc2d76076485"}},{"code_sha256_prefix":"b99eea6376d1e212","entry":"rotate_half","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"peft_pretraining/modeling_llama.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/peft_pretraining/modeling_llama.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"b99eea6376d1e212"}},{"code_sha256_prefix":"25f37cd0df62256a","entry":"beta_scheduler","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"muonlite.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/muonlite.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"25f37cd0df62256a"}},{"code_sha256_prefix":"f28cdb4832ce7cf3","entry":"cosine_lr_schedule","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"muonlite.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/muonlite.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"f28cdb4832ce7cf3"}},{"code_sha256_prefix":"0311c26d2bf832a0","entry":"get_cyclical_cosine_schedule_with_min_lr","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"peft_pretraining/training_utils.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/peft_pretraining/training_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"0311c26d2bf832a0"}},{"code_sha256_prefix":"59f17ef66274a3ca","entry":"get_scheculer","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"peft_pretraining/training_utils.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/peft_pretraining/training_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"59f17ef66274a3ca"}},{"code_sha256_prefix":"d277e7e92ab2f75b","entry":"get_wsd_schedule_with_min_lr","repo":"SHUCHENZHU/LITE","repo_kind":"found_in_text","path":"peft_pretraining/training_utils.py","file_url":"https://github.com/SHUCHENZHU/LITE/blob/HEAD/peft_pretraining/training_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"d277e7e92ab2f75b"}}]},"arxiv_metadata":{"licence":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)","fields":["title","abstract","authors","date"],"primary_category":"cs.LG","source":"arxiv_2026.jsonl"},"syntology_extracted_results":null}