{"id":21135253,"url":"https://github.com/cloneofsimo/karras-power-ema-tutorial","last_synced_at":"2025-04-13T09:38:08.541Z","repository":{"id":215794046,"uuid":"739789942","full_name":"cloneofsimo/karras-power-ema-tutorial","owner":"cloneofsimo","description":null,"archived":false,"fork":false,"pushed_at":"2024-01-06T15:19:26.000Z","size":845,"stargazers_count":51,"open_issues_count":1,"forks_count":1,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-03-27T01:11:11.678Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cloneofsimo.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"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-01-06T14:54:52.000Z","updated_at":"2024-11-21T10:37:19.000Z","dependencies_parsed_at":null,"dependency_job_id":"fa9f8acb-3148-4e83-ae09-3186ac929139","html_url":"https://github.com/cloneofsimo/karras-power-ema-tutorial","commit_stats":null,"previous_names":["cloneofsimo/karras-power-ema-tutorial"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fkarras-power-ema-tutorial","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fkarras-power-ema-tutorial/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fkarras-power-ema-tutorial/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fkarras-power-ema-tutorial/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cloneofsimo","download_url":"https://codeload.github.com/cloneofsimo/karras-power-ema-tutorial/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248691842,"owners_count":21146452,"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":[],"created_at":"2024-11-20T06:46:56.328Z","updated_at":"2025-04-13T09:38:08.535Z","avatar_url":"https://github.com/cloneofsimo.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Karras Power Function EMA (Post-training EMA synthesis)\n\nThis tutorial-repo implements the Karras's Power function EMA, quite incredible trick introduced in the paper [Analyzing and Improving the Training Dynamics of Diffusion Models](https://arxiv.org/abs/2312.02696v1) by Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine.\n\n\n\n\n# So What is Karras's Power function EMA?\n\nI recommend you to read the paper for full detail, but here is the big picture.\n\nRecall that EMA'ing checkpoint is about keeping track of smooth-version of model parameters, $\\theta_\\beta$, where $\\theta_\\beta(t) = \\beta \\theta_\\beta(t-1) + (1-\\beta) \\theta(t)$\n\n, where $\\beta$ is the decaying factor close to 1. Using EMA typically makes the model more robust, and it is a common practice in training deep neural networks.\n\n\nYou want to use EMA, but...\n\n1. You *don't* want the ema to be too slow, because it will make random initialization's contribution to the final model too big.\n2. You definitely want the decaying factor to be self-similar, because you should be able to *increase-time* of the training.\n3. You want to set decaying factor *post-hoc*, because you don't want to retrain the model from scratch with different decaying factor.\n\nKarras's Power function EMA is the answer to all of these problems. He first uses power-function version of EMA where instead of keeping beta constant, he uses $\\beta(t) = (1 - 1/t)^{1 + \\gamma}$, where $\\gamma$ is the hyperparameter. This makes the contribution of historical parameters self-similar, and you can increase the time of the training and it will not change how you expect the EMA to behave (i.e., if you want first 10% of the training to contribute x% of the final model, increasing/decreasing training time will not change that).\n\n# Overall Algorithm and Implementation\n\nSo there is two main part of the algorithm.\n\n1. Saving two copies of the EMA-model, each with different width.\n2. Recovering arbitrary-width EMA\n\nThink of width as decaying factor. Larger width means it will be *smoother*.\n\n![Alt text](image.png)\n\n## First, save two copies of the EMA, with different width\n\nThis is the easy part. You just need to save two copies of the EMA, each with different width (different $\\gamma$).\n\n```python\ngamma_1 = 5\ngamma_2 = 10\nmodel = Model()\nmodel_ema_1 = copy.deepcopy(model).cpu()\nmodel_ema_2 = copy.deepcopy(model).cpu()\n\n\nfor i, batch in enumerate(data_loader):\n    beta_1 = (1 - 1/(i+1)) ** (1 + gamma_1)\n    beta_2 = (1 - 1/(i+1)) ** (1 + gamma_2)\n    # train model\n    loss.backward()\n    optimizer.step()\n    for p, p_ema_1, p_ema_2 in zip(model.parameters(), model_ema_1.parameters(), model_ema_2.parameters()):\n        p_ema_1.data = p_ema_1.data * beta_1 + p.data * (1 - beta_1)\n        p_ema_2.data = p_ema_2.data * beta_2 + p.data * (1 - beta_2)\n    \n    if i % save_freq == 0:\n        torch.save(model_ema_1.state_dict(), f'./model_ema_1_{i}.pth')\n        torch.save(model_ema_2.state_dict(), f'./model_ema_2_{i}.pth')\n\n```\n\n\n## Second, recover arbitrary-decay EMA after training.\n\nNow what if you want to recover EMA with $\\gamma_3$? Incredibly, you can do this with all the saved checkpoints. The math behind this in the paper is bit... *not straightforward* so here is my version of the explanation.\n\nEMA, by definition, can be considered as integral of trajectory of the model parameters. So if you have some weighting function $w(t)$, such that\n\n$$\\theta_e(T) = \\int_0^t w(t) \\theta(t) dt$$\n\nFor a fixed training runs, $t \\in [0, T]$, because we saved two copies of EMA for every, say, total of $n$ checkpoints for different $\\gamma$ this means we know the integral value of the trajectory of the\n\n$$\\theta_{i,j} = \\int_0^T w_{i, j}(t) \\theta(t) dt$$\n\nfor $i = 1, 2$ and $j = 1, 2, \\cdots, n$. $i$ correponds to different width, $j$ corresponds to $j$ th checkpoint. Notice how\n\n$$\nw_{i, j}(t) = \\begin{cases} \nt^{\\gamma_i} / g_{i,j} \u0026 \\text{if } t \u003c j \\\\\n0 \u0026 \\text{otherwise}\n\\end{cases} $$\n\nwhere $g_{i,j}$ is simply the normalization constant to make $\\int_0^T w_{i, j}(t) dt = 1$.\n\nOur goal is then to \n\n1. find a approximate $\\hat{w}_3(t)$ that will give us the EMA that corresponds with $\\gamma_3$.\n\n2. find the correpsonding $\\theta_{3,T}$\n\nSee where this is going? Our goal is to approximate $w_3(t)$ as linear combination of $w_{1, j}(t)$ and $w_{2, j}(t)$, i.e.,\n\n$$w_3(t) = \\sum_{j=1}^n \\alpha_j w_{1, j}(t) + \\beta_j w_{2, j}(t)$$\n\nwhere $\\alpha_j$ and $\\beta_j$ are the coefficients we need to find. This way,\n\n\n$$\\theta_{3,T} = \\int_0^T w_3(t) \\theta(t) dt = \\sum_{j=1}^n \\alpha_j \\theta_{1, j} + \\beta_j \\theta_{2, j}$$\n\nAha! Now we can find $\\alpha_j$ and $\\beta_j$ by solving the linear system of equations. Let's just take this one step further.\n\n## Goal let us project $w_3(t)$ onto the subspace spanned by $w_{1, j}(t)$ and $w_{2, j}(t)$\n\nWe have $K$ functions $f_k(t)$, we have a target $g(t)$ and we want to find $k$ coefficients $\\alpha_i$ such that\n\n$$\\min \\int_0^T \\left( g(t) - \\sum_{i=1}^K \\alpha_i f_i(t) \\right)^2 dt$$\n\nHow would you solve this? \n\nDefine inner product as\n\n$$\\langle f, g \\rangle = \\int_0^T f(t) g(t) dt$$\n\nThen we can rewrite the problem as\n\n$$\\min \\left \\| g - \\sum_{i=1}^K \\alpha_i f_i \\right \\|_2^2$$\n\nif we define $\\| f \\|_2 = \\sqrt{\\langle f, f \\rangle}$, expanding the norm, we get\n\n```math\n\\min {\\left \\| g \\right \\|_2}^2 - 2 \\sum_{i=1}^K \\alpha_i \\langle g, f_i \\rangle + \\sum_{i=1}^K \\sum_{j=1}^K \\alpha_i \\alpha_j \\langle f_i, f_j \\rangle\n```\n\nHa, so substituting $A_{i,j} = \\langle f_i, f_j \\rangle$ and $b_i = \\langle g, f_i \\rangle$, we actually just had linear least square problem!\n\n$$\\min \\left \\| g \\right \\|_2^2 - 2 \\alpha^T b + \\alpha^T A \\alpha$$\n\nwhere $\\alpha = (\\alpha_1, \\cdots, \\alpha_K)^T$.\n\nSo the solution is simply\n\n$$\\alpha = A^{+} b$$\n\nwhere $A^{+}$ is the pseudo-inverse of $A$. We are left to just use the $\\alpha$ to get $\\theta_{3,T}$.\n\n\u003e Note : Well if you ever studied functional analysis, you realize hey, *there exists unique solution* to this problem, via *Hilbert's Projection Theorem*. The above is simply finding the projection of $g$ onto the subspace spanned by $f_i$, in $L^2$ space.\n\nSo thing you learned:\n\n1. The level of approximation is determined by the number of checkpoints you saved. More checkpoints, better approximation.\n2. This doesn't have to be power-function EMA. You can use any weighting function $w(t)$, as long as you can compute the integral of the trajectory of the model parameters.\n\n# I don't care about the math just give me the code?\n\nOk, but reminder this is just for a power-function EMA. You can use this for any weighting function $w(t)$.\n\nIn the above code, you saved two copies of EMA, each with different $\\gamma$. Now you want to recover EMA with $\\gamma_3$. Suppose you saved $n$ checkpoints, at iteration $i_1, i_2, \\cdots, i_n$. Then you can do the following.\n\n```python\nt_checkpoint = t[checkpoint_index]\n    \nts = np.concatenate((t_checkpoint, t_checkpoint))\ngammas = np.concatenate(\n    (\n        np.ones_like(checkpoint_index) * gamma_1,\n        np.ones_like(checkpoint_index) * gamma_2,\n    )\n)\n\nx = solve_weights(ts, gammas, last_index, gamma_3)\nemapoints = np.concatenate((y_t_ema1[checkpoint_index], y_t_ema2[checkpoint_index]))\n\ny_t_ema3 = np.dot(x, emapoints)\n```\n\nwhere `solve_weights` is the function that solves the linear least square problem. You can find the implementation in `ema_eq.py`.\n\nThe result is the EMA with $\\gamma_3$.\n\n![Alt text](ema_eq.png)","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fkarras-power-ema-tutorial","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcloneofsimo%2Fkarras-power-ema-tutorial","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fkarras-power-ema-tutorial/lists"}