Distributionally-Robust Gradient Routing: A Bilevel Sparse Optimization Problem for Compute-Aware Mixture-of-Experts Training

Distributionally-Robust Gradient Routing: A Bilevel Sparse Optimization Problem for Compute-Aware Mixture-of-Experts Training | Research Square window.SnipcartSettings = { analytics: { enabled: false } }; (function() { var accessVector = localStorage.getItem('access_vector') || ''; window.dataLayer = window.dataLayer || []; if (accessVector) { window.dataLayer.push({ user: { profile: { profileInfo: { snid: accessVector } } } }); } })(); (function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start':new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0],j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src='https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f);})(window,document,'script','dataLayer','GTM-K279D39R'); Browse Preprints In Review Journals COVID-19 Preprints AJE Video Bytes Research Tools Research Promotion AJE Professional Editing AJE Rubriq About Preprint Platform In Review Editorial Policies Our Team Advisory Board Help Center Sign In Submit a Preprint Cite Share Download PDF Research Article Distributionally-Robust Gradient Routing: A Bilevel Sparse Optimization Problem for Compute-Aware Mixture-of-Experts Training Xian Rui Chen, Bei Ling Gao, Yu Qiao Fang, An Zhen Liu, Zhou Shen This is a preprint; it has not been peer reviewed by a journal. https://doi.org/ 10.21203/rs.3.rs-8381882/v1 This work is licensed under a CC BY 4.0 License Status: Posted Version 1 posted You are reading this latest preprint version Abstract Distributionally-Robust Gradient Routing (DRGR) is a bilevel sparse optimization framework for training compute-aware Mixture-of-Experts (MoE) models under domain uncertainty and strict resource budgets. DRGR jointly optimizes model parameters and sparse routing policies by minimizing worst-case generalization loss over an f-divergence ambiguity set while explicitly regularizing gradient-traffic and enforcing per-token and per-batch compute constraints. We derive a convex-concave dual reformulation of the inner adversary that yields stable low-dimensional optimization and closed-form adversarial weights, and we propose a proximal alternating minimization algorithm that combines group-sparse proximal updates with exact projection onto budget constraints. To address bilevel sensitivity we develop a Jacobian-free hypergradient estimator using Hessian-vector products implemented via conjugate-gradient or damped Neumann series; this estimator is amenable to distributed expert-parallel settings and is proven to guarantee descent on a natural merit function under controlled inexactness. We provide existence and stationarity guarantees, derive bounds linking routing sparsity and robustness radius to excess risk and communication cost, and propose an evaluation protocol measuring robustness-to-domain-shift, token-level fairness, and FLOPs/latency trade-offs. Empirical studies across synthetic and realistic multi-domain benchmarks demonstrate that DRGR substantially reduces worst-case error and gradient congestion while respecting strict compute budgets. Code and benchmark scripts are provided to facilitate reproducibility and wider adoption. bilevel optimization distributionally robust optimization mixture-of-experts sparse routing implicit differentiation compute-aware learning Full Text Additional Declarations The authors declare no competing interests. Cite Share Download PDF Status: Posted Version 1 posted You are reading this latest preprint version Research Square lets you share your work early, gain feedback from the community, and start making changes to your manuscript prior to peer review in a journal. As a division of Research Square Company, we’re committed to making research communication faster, fairer, and more useful. We do this by developing innovative software and high quality services for the global research community. Our growing team is made up of researchers and industry professionals working together to solve the most critical problems facing scientific publishing. Also discoverable on Platform About Our Team In Review Editorial Policies Advisory Board Help Center Resources Author Services Accessibility API Access RSS feed Manage Cookie Preferences © Research Square 2026 | ISSN 2693-5015 (online) Privacy Policy Terms of Service Do Not Sell My Personal Information {"props":{"pageProps":{"initialData":{"identity":"rs-8381882","acceptedTermsAndConditions":true,"allowDirectSubmit":true,"archivedVersions":[],"articleType":"Research Article","associatedPublications":[],"authors":[{"id":561542291,"identity":"9205bc41-bc2e-4c87-a073-650d1e3881e4","order_by":0,"name":"Xian Rui Chen","email":"","orcid":"","institution":"","correspondingAuthor":false,"prefix":"","firstName":"Xian","middleName":"Rui","lastName":"Chen","suffix":""},{"id":561542292,"identity":"4c4e1ac5-cc98-4e65-97f2-c3305c9d9671","order_by":1,"name":"Bei Ling Gao","email":"","orcid":"","institution":"","correspondingAuthor":false,"prefix":"","firstName":"Bei","middleName":"Ling","lastName":"Gao","suffix":""},{"id":561542293,"identity":"e1644e36-ac72-4409-9072-b1a9c452f3ea","order_by":2,"name":"Yu Qiao Fang","email":"","orcid":"","institution":"","correspondingAuthor":false,"prefix":"","firstName":"Yu","middleName":"Qiao","lastName":"Fang","suffix":""},{"id":561542294,"identity":"c800a685-1184-4fc7-b18a-89ed223c9841","order_by":3,"name":"An Zhen Liu","email":"","orcid":"","institution":"","correspondingAuthor":false,"prefix":"","firstName":"An","middleName":"Zhen","lastName":"Liu","suffix":""},{"id":561542295,"identity":"84d87fc9-ff76-43dd-8ec5-4fb6a01d049d","order_by":4,"name":"Zhou Shen","email":"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAZAAAAAyAQMAAABI0h/eAAAABlBMVEX///8AAABVwtN+AAAACXBIWXMAAA7EAAAOxAGVKw4bAAABFElEQVRIiWNgGAWjYDACCSBmbEASMGAH8QwsSNHCcwBESpCiRSIBJo4dyM9uPvbw5w6bxAbpw4c/fKi5J28u+fzqhh8FEgz87d0J2LQY3DmWbiB5Ji2xgS8tTXLGsWLDnbNzym72AB0mcebsBqxaJHLMJAzbDic28PCYMfOwJSQY3M5Ju8ED1GIgkYtVi/yM/G8SiWAt/J8///kH1HLzTNrNP3i0MNzIYZM4CLGFQZqxDajlBvux2/hsMbiRZibZ2JZm3MbDZibZ25dguOFMDtttGQMJHlx+kZ+R/EzyZ5uNbD8P8+MPP74lyBscP/7s5ps/NnL87b3YHQYDbAgmjwGYxKscDbA/IEX1KBgFo2AUDH8AAA8DYpsNUdNBAAAAAElFTkSuQmCC","orcid":"","institution":"","correspondingAuthor":true,"prefix":"","firstName":"Zhou","middleName":"","lastName":"Shen","suffix":""}],"badges":[],"createdAt":"2025-12-17 05:49:18","currentVersionCode":1,"declarations":{"humanSubjects":false,"vertebrateSubjects":false,"conflictsOfInterestStatement":false,"humanSubjectEthicalGuidelines":false,"humanSubjectConsent":false,"humanSubjectClinicalTrial":false,"humanSubjectCaseReport":false,"vertebrateSubjectEthicalGuidelines":false},"doi":"10.21203/rs.3.rs-8381882/v1","doiUrl":"https://doi.org/10.21203/rs.3.rs-8381882/v1","draftVersion":[],"editorialEvents":[],"editorialNote":"","failedWorkflow":false,"files":[{"id":98624382,"identity":"897f493e-cb30-4dd6-9e8a-ae3c67249aeb","added_by":"auto","created_at":"2025-12-19 17:08:23","extension":"pdf","order_by":1,"title":"","display":"","copyAsset":false,"role":"manuscript-pdf","size":413516,"visible":true,"origin":"","legend":"","description":"","filename":"paper.pdf","url":"https://assets-eu.researchsquare.com/files/rs-8381882/v1_covered_4859b1c0-e255-4019-8392-0f93098c9638.pdf"}],"financialInterests":"The authors declare no competing interests.","formattedTitle":"\u003cp\u003eDistributionally-Robust Gradient Routing: A Bilevel Sparse Optimization Problem for Compute-Aware Mixture-of-Experts Training\u003c/p\u003e","fulltext":[],"fulltextSource":"","fullText":"","funders":[],"hasAdminPriorityOnWorkflow":false,"hasManuscriptDocX":false,"hasOptedInToPreprint":true,"hasPassedJournalQc":"","hasAnyPriority":true,"hideJournal":true,"highlight":"","institution":"","isAcceptedByJournal":false,"isAuthorSuppliedPdf":true,"isDeskRejected":"","isHiddenFromSearch":false,"isInQc":false,"isInWorkflow":false,"isPdf":true,"isPdfUpToDate":true,"isWithdrawnOrRetracted":false,"journal":{"display":true,"email":"[email protected]","identity":"researchsquare","isNatureJournal":false,"hasQc":true,"allowDirectSubmit":true,"externalIdentity":"","sideBox":"","snPcode":"","submissionUrl":"/submission","title":"Research Square","twitterHandle":"researchsquare","acdcEnabled":true,"dfaEnabled":false,"editorialSystem":"","reportingPortfolio":"","inReviewEnabled":false,"inReviewRevisionsEnabled":true},"keywords":"bilevel optimization, distributionally robust optimization, mixture-of-experts, sparse routing, implicit differentiation, compute-aware learning","lastPublishedDoi":"10.21203/rs.3.rs-8381882/v1","lastPublishedDoiUrl":"https://doi.org/10.21203/rs.3.rs-8381882/v1","license":{"name":"CC BY 4.0","url":"https://creativecommons.org/licenses/by/4.0/"},"manuscriptAbstract":"\u003cp\u003eDistributionally-Robust Gradient Routing (DRGR) is a bilevel sparse optimization framework for training compute-aware Mixture-of-Experts (MoE) models under domain uncertainty and strict resource budgets. DRGR jointly optimizes model parameters and sparse routing policies by minimizing worst-case generalization loss over an f-divergence ambiguity set while explicitly regularizing gradient-traffic and enforcing per-token and per-batch compute constraints. We derive a convex-concave dual reformulation of the inner adversary that yields stable low-dimensional optimization and closed-form adversarial weights, and we propose a proximal alternating minimization algorithm that combines group-sparse proximal updates with exact projection onto budget constraints. To address bilevel sensitivity we develop a Jacobian-free hypergradient estimator using Hessian-vector products implemented via conjugate-gradient or damped Neumann series; this estimator is amenable to distributed expert-parallel settings and is proven to guarantee descent on a natural merit function under controlled inexactness. We provide existence and stationarity guarantees, derive bounds linking routing sparsity and robustness radius to excess risk and communication cost, and propose an evaluation protocol measuring robustness-to-domain-shift, token-level fairness, and FLOPs/latency trade-offs. Empirical studies across synthetic and realistic multi-domain benchmarks demonstrate that DRGR substantially reduces worst-case error and gradient congestion while respecting strict compute budgets. Code and benchmark scripts are provided to facilitate reproducibility and wider adoption.\u003c/p\u003e","manuscriptTitle":"Distributionally-Robust Gradient Routing: A Bilevel Sparse Optimization Problem for Compute-Aware Mixture-of-Experts Training","msid":"","msnumber":"","nonDraftVersions":[{"code":1,"date":"2025-12-18 04:44:40","doi":"10.21203/rs.3.rs-8381882/v1","editorialEvents":[{"type":"communityComments","content":0}],"status":"published","journal":{"display":true,"email":"[email protected]","identity":"researchsquare","isNatureJournal":false,"hasQc":true,"allowDirectSubmit":true,"externalIdentity":"","sideBox":"","snPcode":"","submissionUrl":"/submission","title":"Research Square","twitterHandle":"researchsquare","acdcEnabled":true,"dfaEnabled":false,"editorialSystem":"","reportingPortfolio":"","inReviewEnabled":false,"inReviewRevisionsEnabled":true}}],"origin":"","ownerIdentity":"ed81d30f-bcb1-4766-8ad1-01209863f3ff","owner":[],"postedDate":"December 18th, 2025","published":true,"recentEditorialEvents":[],"rejectedJournal":[],"revision":"","amendment":"","status":"posted","subjectAreas":[],"tags":[],"updatedAt":"2025-12-18T04:44:40+00:00","versionOfRecord":[],"versionCreatedAt":"2025-12-18 04:44:40","video":"","vorDoi":"","vorDoiUrl":"","workflowStages":[]},"version":"v1","identity":"rs-8381882","journalConfig":"researchsquare"},"__N_SSP":true},"page":"/article/[identity]/[[...version]]","query":{"redirect":"/article/rs-8381882","identity":"rs-8381882","version":["v1"]},"buildId":"XKTyCvWXoU3ODBz1xrDgd","isFallback":false,"isExperimentalCompile":false,"dynamicIds":[84888],"gssp":true,"scriptLoader":[]}