{"id":13635801,"url":"https://github.com/cor3bit/somax","last_synced_at":"2025-04-19T04:31:31.916Z","repository":{"id":244966841,"uuid":"781098364","full_name":"cor3bit/somax","owner":"cor3bit","description":"Stochastic Second-Order Methods in JAX","archived":false,"fork":false,"pushed_at":"2024-09-23T13:20:23.000Z","size":456,"stargazers_count":2,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2024-11-09T05:34:50.354Z","etag":null,"topics":["gauss-newton-method","hessian-free","jax","machine-learning","natural-gradient","quasi-newton","second-order-optimization","stochastic-optimization"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cor3bit.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-04-02T18:39:28.000Z","updated_at":"2024-10-09T01:35:00.000Z","dependencies_parsed_at":"2024-06-18T18:06:24.774Z","dependency_job_id":"1d5bb1e5-71ba-4f1a-8897-685c67497c46","html_url":"https://github.com/cor3bit/somax","commit_stats":null,"previous_names":["cor3bit/somax"],"tags_count":1,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cor3bit%2Fsomax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cor3bit%2Fsomax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cor3bit%2Fsomax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cor3bit%2Fsomax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cor3bit","download_url":"https://codeload.github.com/cor3bit/somax/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":249606363,"owners_count":21298851,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["gauss-newton-method","hessian-free","jax","machine-learning","natural-gradient","quasi-newton","second-order-optimization","stochastic-optimization"],"created_at":"2024-08-02T00:00:52.100Z","updated_at":"2025-04-19T04:31:31.526Z","avatar_url":"https://github.com/cor3bit.png","language":"Python","readme":"\u003ch1 align='center'\u003eSomax\u003c/h1\u003e\n\n\u003cp align='center'\u003e\n  \u003cimg src=\"assets/somax_logo_mini.png\" alt=\"logo\" width=\"250px\"/\u003e\n\u003c/p\u003e\n\nSomax is a library of Second-Order Methods for stochastic optimization\nwritten in [JAX](https://github.com/google/jax).\nSomax is based on the [JAXopt](https://github.com/google/jaxopt) StochasticSolver API,\nand can be used as a drop-in\nreplacement for JAXopt as well as\n[Optax](https://github.com/google-deepmind/optax) solvers.\n\nCurrently supported methods:\n\n- Diagonal Scaling:\n    - [AdaHessian](https://ojs.aaai.org/index.php/AAAI/article/view/17275);\n    - [Sophia](https://arxiv.org/abs/2305.14342);\n- Hessian-free Optimization:\n    - [Newton-CG](https://epubs.siam.org/doi/10.1137/10079923X);\n- Quasi-Newton:\n    - [Stochastic Quasi-Newton Framework (SQN)](https://arxiv.org/abs/1606.04838);\n- Gauss-Newton:\n    - [Exact Gauss-Newton (EGN)](https://arxiv.org/abs/2405.14402);\n    - [Incremental Gauss-Newton Descent (IGND)](https://arxiv.org/abs/2408.05560);\n    - [Stochastic Gauss-Newton (SGN)](https://arxiv.org/abs/2006.02409);\n- Natural Gradient:\n    - [Natural Gradient with Sherman-Morrison-Woodbury formula (SWM-NG)](https://arxiv.org/abs/1906.02353).\n\nFuture releases:\n\n- Add support for separate \"gradient batches\"\n  and \"curvature batches\" for all solvers;\n- Add support for Optax rate schedules.\n\n⚠️ Since JAXopt is currently being merged into Optax,\nSomax at some point will switch to the Optax API as well.\n\n\n*The catfish in the logo is a nod to \"сом\", \nthe Belarusian word for \"catfish\", also pronounced as \"som\".\n\n\n## Installation\n\n```bash\npip install python-somax\n```\n\nRequires [JAXopt](https://github.com/google/jaxopt) 0.8.2+.\n\n## Quick example\n\n```py\nfrom somax import EGN\n\n# initialize the solver\nsolver = EGN(\n    predict_fun=model.apply,\n    loss_type='mse',\n    learning_rate=0.1,\n    regularizer=1.0,\n)\n\n# initialize the solver state\nopt_state = solver.init_state(params)\n\n# run the optimization loop\nfor i in range(10):\n    params, opt_state = solver.update(params, opt_state, batch_x, batch_y)\n```\n\nSee more in the [examples](examples) folder.\n\n## Citation\n\n```bibtex\n@misc{korbit2024somax,\n  author = {Nick Korbit},\n  title = {{SOMAX}: a library of second-order methods for stochastic optimization written in {JAX}},\n  year = {2024},\n  url = {https://github.com/cor3bit/somax},\n}\n```\n\n## See also\n\n**Optimization with JAX**  \n[Optax](https://github.com/google-deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  \n[JAXopt](https://github.com/google/jaxopt): deterministic second-order methods (e.g., Gauss-Newton, Levenberg\nMarquardt), stochastic first-order methods PolyakSGD, ArmijoSGD.\n\n**Awesome Projects**  \n[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of various JAX projects.  \n[Awesome SOMs](https://github.com/cor3bit/awesome-soms): a list\nof resources for second-order optimization methods in machine learning.\n\n## Acknowledgements\n\nSome of the implementation ideas are based on the following repositories:\n\n- Line Search in JAXopt: https://github.com/google/jaxopt/blob/main/jaxopt/_src/armijo_sgd.py#L48\n\n- L-BFGS Inverse Hessian-Gradient product in JAXopt: https://github.com/google/jaxopt/blob/main/jaxopt/_src/lbfgs.py#L44\n\n- AdaHessian (official implementation): https://github.com/amirgholami/adahessian\n\n- AdaHessian (Nestor Demeure's implementation): https://github.com/nestordemeure/AdaHessianJax\n\n- Sophia (official implementation): https://github.com/Liuhong99/Sophia\n\n- Sophia (levanter implementation): https://github.com/stanford-crfm/levanter/blob/main/src/levanter/optim/sophia.py\n\n","funding_links":[],"categories":["Implementation in JAX"],"sub_categories":["Other"],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcor3bit%2Fsomax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcor3bit%2Fsomax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcor3bit%2Fsomax/lists"}