{"id":15805257,"url":"https://github.com/evanatyourservice/kron_torch","last_synced_at":"2025-05-05T20:03:44.642Z","repository":{"id":257806129,"uuid":"864743147","full_name":"evanatyourservice/kron_torch","owner":"evanatyourservice","description":"An implementation of PSGD Kron second-order optimizer for PyTorch","archived":false,"fork":false,"pushed_at":"2025-03-29T22:35:22.000Z","size":306,"stargazers_count":89,"open_issues_count":1,"forks_count":6,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-25T06:03:42.925Z","etag":null,"topics":["machine-learning","optimization","pytorch","second-order-optimization"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"cc-by-4.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/evanatyourservice.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,"zenodo":null}},"created_at":"2024-09-29T03:28:51.000Z","updated_at":"2025-04-17T20:30:01.000Z","dependencies_parsed_at":"2024-10-01T17:20:44.726Z","dependency_job_id":"43aefbb2-925f-4cbc-8f71-4cd4d36ac8f7","html_url":"https://github.com/evanatyourservice/kron_torch","commit_stats":{"total_commits":36,"total_committers":2,"mean_commits":18.0,"dds":"0.13888888888888884","last_synced_commit":"eedf42c386dcae624c3da2c059b70d3e98bf745d"},"previous_names":["evanatyourservice/kron_torch"],"tags_count":22,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/evanatyourservice%2Fkron_torch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/evanatyourservice%2Fkron_torch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/evanatyourservice%2Fkron_torch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/evanatyourservice%2Fkron_torch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/evanatyourservice","download_url":"https://codeload.github.com/evanatyourservice/kron_torch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252569638,"owners_count":21769516,"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":["machine-learning","optimization","pytorch","second-order-optimization"],"created_at":"2024-10-05T02:05:44.477Z","updated_at":"2025-05-05T20:03:44.588Z","avatar_url":"https://github.com/evanatyourservice.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# PSGD Kron\n\nFor Xi-Lin's original PSGD repo, see [psgd_torch](https://github.com/lixilinx/psgd_torch).\n\nFor JAX versions, see [psgd_jax](https://github.com/evanatyourservice/psgd_jax) and \n[distributed_kron](https://github.com/evanatyourservice/distributed_kron).\n\nImplementation of PSGD Kron for PyTorch. \nPSGD is a second-order optimizer originally created by Xi-Lin Li that uses either a hessian-based \nor whitening-based (gg^T) preconditioner and lie groups to improve training convergence, \ngeneralization, and efficiency. I highly suggest taking a look at Xi-Lin's PSGD repo's readme linked\nto above for interesting details on how PSGD works and experiments using PSGD. There are also \npaper resources listed near the bottom of this readme.\n\n### `kron`:\n\nThe most versatile and easy-to-use PSGD optimizer is `kron`, which uses a Kronecker-factored \npreconditioner. It has less hyperparameters that need tuning than adam, and can generally act as a \ndrop-in replacement.\n\n### Thanks\n\nShoutout to @ClashLuke for developing efficiency improvements for PSGD Kron in the \n[heavyball](https://github.com/HomebrewML/HeavyBall) repo, and for the design of 'smart_one_diag' memory\nsave mode, which is a method to improve memory usage and speed with almost no cost to the optimizer's\neffectiveness. In Xi-Lin's repo, the equivalent is setting `preconditioner_max_skew=1`.\n\n## Installation\n\n```bash\npip install kron-torch\n```\n\n## Basic Usage (Kron)\n\nKron schedules the preconditioner update probability by default to start at 1.0 and anneal to 0.03 \nat the beginning of training, so training will be slightly slower at the start but will speed up \nby around 4k steps.\n\nFor basic usage, use `kron` optimizer like any other pytorch optimizer:\n\n```python\nfrom kron_torch import Kron\n\noptimizer = Kron(params)\n\noptimizer.zero_grad()\nloss.backward()\noptimizer.step()\n```\n\n**Basic hyperparameters:**\n\nTLDR: Start with a learning rate around 3x smaller than adam's, and a weight decay 3-10x larger. \nThere is no b2 or epsilon.\n\nThese next 3 settings control whether a dimension's preconditioner is diagonal or triangular. \nFor example, for a layer with shape (256, 128), triagular preconditioners would be shapes (256, 256)\nand (128, 128), and diagonal preconditioners would be shapes (256,) and (128,). Depending on how \nthese settings are chosen, `kron` can balance between memory/speed and effectiveness. Defaults lead\nto most precoditioners being triangular except for 1-dimensional layers and very large dimensions.\n\n`max_size_triangular`: Any dimension with size above this value will have a diagonal preconditioner.\n\n`min_ndim_triangular`: Any tensor with less than this number of dims will have all diagonal \npreconditioners. Default is 2, so single-dim layers like bias and scale will use diagonal\npreconditioners.\n\n`memory_save_mode`: Can be None, 'smart_one_diag', 'one_diag', or 'all_diag'. None is default and lets all \npreconditioners be triangular. 'smart_one_diag' sets the largest dim to diagonal only if it's larger than\nthe second largest dim (if it stands out). 'one_diag' sets the largest or last dim per layer as diagonal using \n`np.argsort(shape)[::-1][0]`. 'all_diag' sets all preconditioners to be diagonal.\n\n`preconditioner_update_probability`: Preconditioner update probability uses a schedule by default \nthat works well for most cases. It anneals from 1 to 0.03 at the beginning of training, so training \nwill be slightly slower at the start but will speed up by around 4k steps. PSGD generally benefits\nfrom more preconditioner updates at the start of training, but once the preconditioner is learned \nit's okay to do them less often. An easy way to adjust update frequency is to define your own schedule\nusing the `precond_update_prob_schedule` function in kron.py (just changing the `min_prob` value \nis easiest) and pass this into kron through the `preconditioner_update_probability` hyperparameter.\n\nThis is the default schedule defined in the `precond_update_prob_schedule` function at the top of kron.py:\n\n\u003cimg src=\"assets/default_schedule.png\" alt=\"Default Schedule\" width=\"800\" style=\"max-width: 100%; height: auto;\" /\u003e\n\n\n## Resources\n\nPSGD papers and resources listed from Xi-Lin's repo\n\n1) Xi-Lin Li. Preconditioned stochastic gradient descent, [arXiv:1512.04202](https://arxiv.org/abs/1512.04202), 2015. (General ideas of PSGD, preconditioner fitting losses and Kronecker product preconditioners.)\n2) Xi-Lin Li. Preconditioner on matrix Lie group for SGD, [arXiv:1809.10232](https://arxiv.org/abs/1809.10232), 2018. (Focus on preconditioners with the affine Lie group.)\n3) Xi-Lin Li. Black box Lie group preconditioners for SGD, [arXiv:2211.04422](https://arxiv.org/abs/2211.04422), 2022. (Mainly about the LRA preconditioner. See [these supplementary materials](https://drive.google.com/file/d/1CTNx1q67_py87jn-0OI-vSLcsM1K7VsM/view) for detailed math derivations.)\n4) Xi-Lin Li. Stochastic Hessian fittings on Lie groups, [arXiv:2402.11858](https://arxiv.org/abs/2402.11858), 2024. (Some theoretical works on the efficiency of PSGD. The Hessian fitting problem is shown to be strongly convex on set ${\\rm GL}(n, \\mathbb{R})/R_{\\rm polar}$.)\n5) Omead Pooladzandi, Xi-Lin Li. Curvature-informed SGD via general purpose Lie-group preconditioners, [arXiv:2402.04553](https://arxiv.org/abs/2402.04553), 2024. (Plenty of benchmark results and analyses for PSGD vs. other optimizers.)\n\n\n## License\n\n[![CC BY 4.0][cc-by-image]][cc-by]\n\nThis work is licensed under a [Creative Commons Attribution 4.0 International License][cc-by].\n\n2024 Evan Walters, Omead Pooladzandi, Xi-Lin Li\n\n\n[cc-by]: http://creativecommons.org/licenses/by/4.0/\n[cc-by-image]: https://licensebuttons.net/l/by/4.0/88x31.png\n[cc-by-shield]: https://img.shields.io/badge/License-CC%20BY%204.0-lightgrey.svg\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fevanatyourservice%2Fkron_torch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fevanatyourservice%2Fkron_torch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fevanatyourservice%2Fkron_torch/lists"}