Zekun WangContinual Learning Diffusion Models
A geometric view of learning new tasks while preserving earlier abilities
Zekun Wang · Georgia Institute of Technology
CW10 task 1, hammer, before and after learning nine more tasks.A robot should still know how to use a hammer after learning nine more skills. An image generator that learns new animals should still be able to draw the vehicles it learned earlier. We want each new task to expand what a model can do. But when successive tasks change the same weights, learning can erase the abilities it should be building on.
Our two papers approach this problem through diffusion models. The first asks which parameter changes threaten earlier tasks and finds that diffusion gradients offer a surprisingly compact answer. The second combines that information with replay and uncovers a local connection to meta-learning: continual training can preserve a useful initialization from which earlier tasks recover quickly. Together, they connect rank-1 Fisher, generative replay, and trust regions to continual learning in image generation and robotic control.
The same weights have to serve every task
Diffusion models learn to remove noise from training examples. Repeated denoising can then produce an image or, for a diffusion policy, an action sequence conditioned on recent observations. Our experiments use separate image generators and robot policies. Each model learns its own sequence of tasks, such as new image classes or different manipulation skills.
The difficulty begins when the next task arrives. Its gradients change parameters that earlier tasks also depend on. The model must learn from the new data while continuing to perform earlier tasks, without returning to their full training datasets. Learning the current task and retaining earlier tasks are joint requirements. A model that succeeds at each task on arrival can still end the sequence unable to perform most of them.3, 4
The diagonal approximation misses how weights change together
Elastic weight consolidation, or EWC, protects earlier tasks by penalizing changes to their learned parameters. Fisher information determines which directions receive the strongest penalty. Our diffusion work estimates an empirical Fisher by averaging outer products of denoising loss gradients. These products retain correlations between parameters, including changes that reinforce or compensate for one another.
The full matrix is expensive to store, so EWC commonly keeps only its diagonal. Every weight receives an individual importance value, but all the off-diagonal entries disappear. In an image denoiser, several weights can work together to represent an edge or a texture. Their individual importance does not tell us how changing them together affects the model. The missing entries describe those relationships.
The comparison below starts with a Fisher matrix whose gradients largely agree. Its off-diagonal structure gives the loss a preferred geometry. Removing that structure changes the shape of the valley. A rank-1 estimate retains the shared pattern, so its geometry stays close to full Fisher. We use a richer illustrative loss surface to make the difference visible, with the same nonlinear detail in every panel. The next section explains why diffusion gradients can have this structure.5, 7
What diagonal Fisher misses
Diagonal Fisher removes parameter correlations. A rank-1 estimate preserves the dominant pattern and stays close to full Fisher in this illustrative comparison. View full-size figureIncreasing the strength of an inaccurate penalty cannot restore the correlations it discarded. We need a compact estimate that preserves the important direction. Diffusion models provide a way to find it.
Diffusion gradients reveal a common direction
Later in the forward noising process, noise dominates the input. The rank-1 paper studies this low signal-to-noise ratio regime near convergence using a locally linear score approximation. Its central observation is that per-example gradients align strongly with their mean. Their outer products therefore concentrate Fisher information along one direction.
The animation below multiplies each sample gradient by itself to form a matrix, then averages these matrices to build the empirical Fisher. When the gradients share a dominant direction, the scaled outer product of their mean recovers that pattern, including relationships between different weights.
We test this structure using 1,024 sampled gradients from a small MNIST diffusion model. Figure 3 reports a second-to-first eigenvalue ratio of 0.022 at diffusion timestep 700. The diagonal approximation has relative Frobenius error near 1.0 across the evaluated timesteps, while rank-1 error becomes much smaller at later steps. These are measured matrix approximation errors, separate from the illustrative parameter-space comparison.1, 7
Watch the Fisher matrix take shape
Averaging gradient outer products builds empirical Fisher. When gradients align, a scaled outer product of their mean keeps the dominant correlations. View full-size figureMean gradients also align across timesteps, supporting a practical estimate from gradients averaged over sampled noise levels. The resulting EWC penalty stores a direction and a scale, then penalizes the squared projection of each update onto that direction. Its representation scales linearly with parameter count per stored task, like diagonal EWC, while preserving the correlation structure that matters most.
The assumptions behind this connection
The analysis concerns variance-preserving diffusion with squared-error denoising, near convergence, under the paper’s score-model assumptions. Using the mean requires a nonzero mean gradient, and combining timesteps relies on their alignment. Each task can have its own protected direction.
Finding shared solutions and keeping them stable
A better Fisher does not solve continual learning by itself. EWC uses Fisher to define a penalty around earlier task solutions. The model can reach a shared solution for Tasks 1 and 2, then find that Task 3 has no compatible solution there. Learning the new task forces it away from the earlier overlap. Figure 1(a) illustrates this limitation. The model also needs a training signal that helps it find weights that work across tasks.
Generative distillation supplies the missing training signal. A frozen copy of the preceding image generator produces replay examples. The current model learns from new data while matching the teacher’s denoising predictions on replayed inputs. Earlier tasks now contribute directly to training. This encourages parameter sharing by asking the same weights to satisfy current and earlier tasks together.
Figure 1(b) shows how replay changes this problem. Training on current and generated earlier examples promotes parameter sharing and can reach a different basin. But generated examples are imperfect. Several parameter regions may fit the replay data without preserving the original tasks equally well. As replay errors accumulate, a model can keep fitting its generated examples while losing earlier abilities.
Figure 1(c) combines both ingredients in Trust Region. Replay guides learning toward shared parameters, while the rank-1 Fisher penalty limits changes to combinations of weights that earlier tasks depend on. The three views below show the roles separately and together. EWC provides local protection, replay provides earlier-task training data, and the combination favors a useful shared solution when one is nearby.8, 9
EWC, replay, and Trust Region
EWC protects earlier solutions locally. Replay promotes parameter sharing. Trust Region combines both to preserve a useful shared solution. View full-size figureThe trust region describes a neighborhood with a small Fisher-weighted penalty around earlier optima. Movement along less sensitive directions costs less. Training uses this soft penalty alongside current data and replay, which discourages harmful changes while allowing the model to adapt. A single rank-1 term constrains one direction, so this neighborhood need not have a closed boundary.
How quickly does an earlier skill return?
Fast recovery is a useful property of the model’s parameters. After learning a new task perturbs an earlier skill, a few further updates should bring that skill back. The model remains a good initialization for its earlier tasks. This is the connection to gradient-based meta-learning.
MAML learns such an initialization by taking an update on support data and optimizing the query loss after adaptation. We start with the continual-learning objective and compare its old-task contribution with that one-step meta-gradient. Replay supplies the query signal, while Fisher supplies the local curvature. The derivation below shows where those roles meet.2, 6, 10, 11, 12
The meta-learning structure inside Trust Region
Compare the contribution of one earlier task. The current-task gradient remains in the training update. Here δᵢ = θ − θᵢ* measures displacement from the earlier checkpoint.
Start with the old-task gradients
Trust Region combines replay with a Fisher penalty. MAML first adapts on support data, then differentiates the query loss through that step.
θi′=θ−αgis(θ)
Trust Region
One-step MAML
giMAML=(I−αHis)giq(θi′)
Eqs. 5–6
Replay supplies the query signal. Fisher supplies curvature.
Near an earlier optimum, the paper approximates the replay and adapted-query gradients by Fᵢδᵢ, and the support Hessian by Fᵢ.
gi≃giq(θi′)≃Fiδi,His≃Fi Trust Region
giTR≃(β+λ)Fiδi
One-step MAML
giMAML≃(Fi−αFi2)δi
Eqs. 7–12
Applying Fisher twice preserves its direction
For rank-1 Fisher, ρᵢ is the nonzero eigenvalue and uᵢ is a unit vector. The second curvature factor changes the scale, so both expressions reduce to the same Fᵢδᵢ direction.
Fi≃ρiuiui⊤⟹Fi2≃ρiFi
Trust Region
giTR≃scale(β+λ)shared directionFiδi One-step MAML
giMAML≃scale(1−αρi)shared directionFiδi Eqs. 11–13
1−αρi>0
Under these local approximations and a positive MAML coefficient, both old-task contributions descend along −Fᵢδᵢ. Their scales can differ across tasks. This connects the continual-learning objective to a MAML-style initialization without adding a support/query training loop.We checked the connection directly on the CW10 diffusion policy. Its old-task update reached a mean cosine similarity of 0.83 with an exact second-order, one-step MAML meta-gradient. Both were evaluated at the same parameters, with support and query samples from the same stream. This measures the directional agreement predicted by the analysis.
We then test the behavior that motivates the connection. We keep evaluating Task 1 while training later tasks with current data, replay, and the Fisher penalty. The curves track how quickly earlier performance returns during this ongoing training. A useful initialization should require only a small correction to recover the earlier skill.10, 11
Fast recovery during continual learning
ImageNet-500
Task 1 FID · lower is better

Continual World 10
Task 1 success rate · higher is better

Original Figure 2 from the paper.On CW10, Trust Region recovers 90% of initial Task 1 success within 2 to 45 updates at every later task transition. At the transition to Task 10, the counts are 2 updates with Trust Region and 8,000 with Replay. Replay ties at several transitions and is faster at one. Trust Region maintains short recovery times across the full sequence, alongside better final retention.
The derivation identifies a meta-learning structure inside a practical continual-learning update. The gradient comparison tests that connection directly. The recovery curves show its behavioral consequence: learning new tasks while keeping earlier abilities easy to recover.
The assumptions behind this connection
The derivation uses locality around approximate earlier optima, agreement between replay and adapted-query gradients, similar support/query curvature, and a rank-1 Fisher approximation. The empirical Fisher is a curvature proxy. The per-task directions agree when 1 − αρᵢ is positive and Fᵢδᵢ is nonzero. The scales can differ by task, so this does not make the summed updates or complete algorithms identical. Recovery is measured on previously learned tasks during continual training. It is not a separate adaptation test on unseen tasks.
Retaining images and manipulation skills
The rank-1 paper evaluates four image datasets. On ImageNet-1k, twenty tasks introduce fifty classes each. Adding rank-1 EWC to generative distillation reduces final average FID from 69.0 to 48.5 and forgetting from 46.2 to 15.2. FID measures generation quality, with lower values indicating a closer match to real images. Forgetting measures its increase between first learning a task and the final evaluation.
On this same stream, diagonal EWC with distillation ends at 73.8 FID. Rank-1 EWC without distillation ends at 74.3. Both comparisons matter: useful curvature and a cross-task training signal are needed together.
The trust-region paper uses a separate ten-task ImageNet-500 sequence. The combined method improves final average FID from 53.4 to 44.5 and reduces forgetting from 18.2 to 10.6, compared with replay. These comparisons belong to different protocols and should be read within each experiment.
On Continual-World-10, diffusion policies learn from offline expert demonstrations and use state observations for simulated manipulation. Trust Region reaches 88.3% average success versus 85.3% with Replay. Average forgetting falls from 8.2 to 4.4 percentage points. The same principle helps retain action-generation skills across distinct tasks.13, 14
Measured results
Final performance and forgetting on ImageNet-1k, ImageNet-500, and CW10, with one-standard-error bars. Compare methods within each benchmark. View full-size figureMeasured ImageNet-1k results, Rank-1 Fisher Table 1: 1,000 classes at 32 × 32 pixels over 20 tasks. Trust Region reports final average FID 48.5 ± 1.9 and forgetting 15.2 ± 4.8 FID points, compared with 69.0 ± 2.2 and 46.2 ± 12.9 for distillation alone. Lower is better. Errors are standard errors over 3 seeds. Source: https://arxiv.org/html/2509.23593v2#S5.T1
Measured ImageNet-500 results, Trust Region Table 1: the first 500 ImageNet classes at 32 × 32 pixels over 10 tasks. Trust Region reports final average FID 44.5 ± 2.3 and forgetting 10.6 ± 3.0 FID points, compared with Replay at 53.4 ± 6.0 and 18.2 ± 4.6. Lower is better. Errors are standard errors. This task sequence is separate from the Rank-1 Fisher experiment. Source: https://arxiv.org/html/2602.02417v1#S4.T1
Measured diffusion-policy results on Continual World 10, Trust Region Table 1: Trust Region reports average success 88.3% ± 0.4 and forgetting 4.4 ± 0.9 percentage points. Replay reports 85.3% ± 2.0 and 8.2 ± 2.0. Higher success and lower forgetting are better. Errors are standard errors in percentage points. The policy trains on offline demonstrations and is evaluated in simulation. Source: https://arxiv.org/html/2602.02417v1#S4.T1
Reported Task 1 recovery on CW10, Trust Region Table 2: at 90% of initial success while learning Task 10, Trust Region takes 2 gradient updates and Replay 8,000. While learning Task 5, Replay is faster at 20 versus 30 updates. Lower is better. These are particular transitions at a relative target, with no reported errors. Source: https://arxiv.org/html/2602.02417v1#S4.T2
Important limits: Rank-1 Fisher Table 1 keeps a stronger non-continual reference visible at 11.7 ± 0.1 FID. Rank-1 EWC without distillation forgets more than diagonal EWC, 41.3 ± 1.8 versus 34.2 ± 3.6 FID points. In the trust-region paper, the CW10 configuration table lists 4 seeds while Figure 2(b) lists 3, so the seed count is unresolved here. Missing values stay unreported. Sources: https://arxiv.org/html/2509.23593v2#S5.T1 and https://arxiv.org/html/2602.02417v1#A3.T4 and https://arxiv.org/html/2602.02417v1#S4.F2.sf2
Learning should keep adding to what a model can do
The two papers address complementary requirements for learning that accumulates. We need a useful account of which parameter changes matter, and training signals that help the model satisfy tasks together. Diffusion gradients make the first tractable. Replay makes the second possible, while a better constraint limits the drift that replay can introduce.
Our broader aim is continual learning that supports adaptation and generalization at deployment time. Larger generators, longer task streams, and physical robots remain important tests. Success means learning new abilities efficiently, retaining earlier knowledge, and using both when the next unfamiliar task arrives.1, 2
Read the full work

ICLR 2026
Zekun Wang, Anant Gupta, Zihan Dong, and Christopher J. MacLellan
Read PDF

NeurIPS 2026
Zekun Wang, Anant Gupta, and Christopher J. MacLellan
Read PDF
Further reading
- Avoid Catastrophic Forgetting with Rank-1 Fisher from Diffusion Models
ICLR 2026. Section 3.1 derives the gradient-alignment argument.
- Trust Region Continual Learning as an Implicit Meta-Learner
NeurIPS 2026. Section 3.3, Equations 5–13, derives the local, task-wise connection to the MAML meta-gradient.
- Denoising Diffusion Probabilistic Models
Ho, Jain, and Abbeel. NeurIPS 2020. The denoising objective behind the image models.
- Diffusion Policy: Visuomotor Policy Learning via Action Diffusion
Chi and colleagues. Robotics: Science and Systems 2023. Action generation through diffusion.
- Overcoming catastrophic forgetting in neural networks
Kirkpatrick and colleagues. PNAS 2017. Elastic weight consolidation.
- Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
Finn, Abbeel, and Levine. ICML 2017. Learning an initialization for rapid adaptation.
- Measured Fisher spectrum and approximation errors
Rank-1 Fisher paper, Figure 3. Full Fisher, diagonal, and rank-1 comparisons across diffusion timesteps.
- Promoting parameter sharing through generative distillation
Rank-1 Fisher paper, Section 3.4. The frozen teacher and the combined training objective.
- EWC, replay, and trust-region continual learning
Trust-region paper, Figure 1 and Section 3.2. Task overlap, replay drift, and the hybrid approach.
- Direct comparison with an exact MAML meta-gradient
Trust Region paper, Appendix D.1 and Table 9. The CW10 old-task update has mean cosine similarity 0.83 with an exact second-order, one-step MAML meta-gradient.
- Task 1 performance throughout continual training
Trust Region paper, Figure 2. Original ImageNet-500 and CW10 learning curves across ten tasks.
- Task 1 recovery while learning later tasks
Trust-region paper, Table 2. Update counts at specified FID and success-rate thresholds.
- Continual image generation with rank-1 Fisher
Rank-1 Fisher paper, Table 1. Final FID and forgetting across four image datasets.
- ImageNet-500 and Continual-World-10 results
Trust-region paper, Table 1. Final image quality, manipulation success, and forgetting.