Continual Learning: Learning Without Forgetting

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from typing import Dict, List, Tuple
import copy
 
class EWC:
    """
    Elastic Weight Consolidation implementation.
    
    EWC protects important parameters by penalizing changes proportional
    to their Fisher Information, which measures parameter sensitivity.
    
    Mathematical formulation:
        L(θ) = L_new(θ) + (λ/2) Σᵢ Fᵢ(θᵢ - θ*ᵢ)²
    
    where F is the diagonal of the Fisher Information Matrix.
    """
    
    def __init__(self, model: nn.Module, lambda_ewc: float = 1000):
        """
        Args:
            model: The neural network model
            lambda_ewc: Regularization strength (higher = more protection)
        """
        self.model = model
        self.lambda_ewc = lambda_ewc
        
        # Storage for task-specific parameters and Fisher information
        self.saved_params: Dict[str, torch.Tensor] = {}
        self.fisher_diag: Dict[str, torch.Tensor] = {}
        
    def compute_fisher(
        self, 
        dataloader: DataLoader, 
        num_samples: int = 2000
    ) -> None:
        """
        Compute the diagonal of the Fisher Information Matrix.
        
        The Fisher Information measures how much the log-likelihood changes
        when parameters change. Higher Fisher = more important parameter.
        
        Mathematical definition:
            F = E[∇log p(y|x,θ) ∇log p(y|x,θ)ᵀ]
        
        Args:
            dataloader: DataLoader for the completed task
            num_samples: Number of samples for Fisher estimation
        """
        self.model.eval()
        
        # Initialize Fisher to zeros
        fisher = {n: torch.zeros_like(p) 
                  for n, p in self.model.named_parameters() 
                  if p.requires_grad}
        
        samples_seen = 0
        
        for inputs, targets in dataloader:
            if samples_seen >= num_samples:
                break
                
            inputs = inputs.to(next(self.model.parameters()).device)
            batch_size = inputs.size(0)
            
            # Forward pass
            outputs = self.model(inputs)
            
            # For classification: sample from output distribution
            # This is the key insight - we use the model's own predictions
            probs = F.softmax(outputs, dim=1)
            sampled_labels = probs.multinomial(1).squeeze()
            
            # Compute log-likelihood for sampled labels
            log_probs = F.log_softmax(outputs, dim=1)
            log_likelihood = log_probs.gather(1, sampled_labels.unsqueeze(1))
            
            # Compute gradients for each sample
            for i in range(min(batch_size, num_samples - samples_seen)):
                self.model.zero_grad()
                log_likelihood[i].backward(retain_graph=(i < batch_size - 1))
                
                # Accumulate squared gradients (diagonal of outer product)
                for n, p in self.model.named_parameters():
                    if p.requires_grad and p.grad is not None:
                        fisher[n] += p.grad.data.clone() ** 2
                
                samples_seen += 1
                
        # Normalize by number of samples
        for n in fisher:
            fisher[n] /= samples_seen
            
        self.fisher_diag = fisher
        
        # Save current parameters as the optimal for this task
        self.saved_params = {
            n: p.data.clone() 
            for n, p in self.model.named_parameters()
            if p.requires_grad
        }
        
        self.model.train()
        
    def penalty(self) -> torch.Tensor:
        """
        Compute the EWC penalty term.
        
        Returns:
            Scalar tensor: (λ/2) Σᵢ Fᵢ(θᵢ - θ*ᵢ)²
        """
        if not self.fisher_diag:
            return torch.tensor(0.0)
            
        loss = torch.tensor(0.0).to(next(self.model.parameters()).device)
        
        for n, p in self.model.named_parameters():
            if n in self.fisher_diag:
                # Importance-weighted quadratic penalty
                fisher = self.fisher_diag[n]
                optimal = self.saved_params[n]
                
                loss += (fisher * (p - optimal) ** 2).sum()
                
        return (self.lambda_ewc / 2) * loss
    
    def training_step(
        self, 
        inputs: torch.Tensor, 
        targets: torch.Tensor,
        criterion: nn.Module,
        optimizer: torch.optim.Optimizer
    ) -> float:
        """
        Perform one EWC training step.
        
        Returns:
            Combined loss value
        """
        optimizer.zero_grad()
        
        # Forward pass
        outputs = self.model(inputs)
        
        # Task loss + EWC penalty
        task_loss = criterion(outputs, targets)
        ewc_loss = self.penalty()
        total_loss = task_loss + ewc_loss
        
        # Backward and optimize
        total_loss.backward()
        optimizer.step()
        
        return total_loss.item()
 
 
# Example usage
def demonstrate_ewc():
    """Demonstrate EWC on sequential tasks."""
    
    # Simple MLP for demonstration
    model = nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10)
    )
    
    ewc = EWC(model, lambda_ewc=1000)
    
    # After training on Task A:
    # ewc.compute_fisher(task_a_dataloader)
    
    # Training on Task B with EWC protection:
    # for inputs, targets in task_b_dataloader:
    #     loss = ewc.training_step(inputs, targets, criterion, optimizer)
    
    print("EWC ready for continual learning!")
    
demonstrate_ewc()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

import torch
import torch.nn as nn
from typing import Dict
 
class OnlineEWC:
    """
    Online Elastic Weight Consolidation.
    
    Maintains a running Fisher estimate and reference point, avoiding
    linear storage growth with number of tasks.
    
    Update rules:
        F_cumulative = γ * F_old + F_new
        θ* = θ_current (updated after each task)
    
    This approach trades some precision for constant memory usage.
    """
    
    def __init__(
        self, 
        model: nn.Module, 
        lambda_ewc: float = 1000,
        gamma: float = 0.9
    ):
        """
        Args:
            model: Neural network model
            lambda_ewc: Regularization strength  
            gamma: Decay factor for old Fisher (0 = only new, 1 = cumulative)
        """
        self.model = model
        self.lambda_ewc = lambda_ewc
        self.gamma = gamma
        
        self.cumulative_fisher: Dict[str, torch.Tensor] = {}
        self.reference_params: Dict[str, torch.Tensor] = {}
        
    def update_fisher(self, new_fisher: Dict[str, torch.Tensor]) -> None:
        """
        Update cumulative Fisher with new task's Fisher.
        
        Uses exponential moving average:
            F_cumulative = γ * F_old + F_new
        """
        if not self.cumulative_fisher:
            # First task: just store the Fisher
            self.cumulative_fisher = {k: v.clone() for k, v in new_fisher.items()}
        else:
            # Subsequent tasks: weighted combination
            for n in new_fisher:
                self.cumulative_fisher[n] = (
                    self.gamma * self.cumulative_fisher[n] + 
                    new_fisher[n]
                )
                
    def update_reference(self) -> None:
        """Update reference parameters to current values."""
        self.reference_params = {
            n: p.data.clone()
            for n, p in self.model.named_parameters()
            if p.requires_grad
        }
        
    def consolidate(self, new_fisher: Dict[str, torch.Tensor]) -> None:
        """
        Consolidate knowledge after completing a task.
        
        Call this after training on each task:
        1. Update Fisher with new task's importance
        2. Store current parameters as reference point
        """
        self.update_fisher(new_fisher)
        self.update_reference()
        
    def penalty(self) -> torch.Tensor:
        """
        Compute Online EWC penalty.
        
        Same formulation as standard EWC but uses cumulative Fisher
        and single reference point.
        """
        if not self.cumulative_fisher:
            return torch.tensor(0.0)
            
        device = next(self.model.parameters()).device
        loss = torch.tensor(0.0).to(device)
        
        for n, p in self.model.named_parameters():
            if n in self.cumulative_fisher:
                fisher = self.cumulative_fisher[n]
                ref = self.reference_params[n]
                loss += (fisher * (p - ref) ** 2).sum()
                
        return (self.lambda_ewc / 2) * loss
 
 
class ProgressiveEWC:
    """
    Progressive Network variant with EWC for forward transfer.
    
    Freezes old task columns while allowing new capacity.
    Uses EWC only for shared representations.
    """
    
    def __init__(self, base_model: nn.Module, lambda_ewc: float = 1000):
        self.base_model = base_model
        self.lambda_ewc = lambda_ewc
        self.task_columns = [base_model]
        self.ewc_modules = [OnlineEWC(base_model, lambda_ewc)]
        
    def add_task_column(self, new_column: nn.Module) -> None:
        """Add a new column for a new task with lateral connections."""
        # Freeze previous columns
        for param in self.task_columns[-1].parameters():
            param.requires_grad = False
            
        self.task_columns.append(new_column)
        self.ewc_modules.append(OnlineEWC(new_column, self.lambda_ewc))
        
    def forward(self, x: torch.Tensor, task_id: int) -> torch.Tensor:
        """Forward through appropriate column."""
        return self.task_columns[task_id](x)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

import torch
import torch.nn as nn
from typing import Dict, Optional
import copy
 
class SynapticIntelligence:
    """
    Synaptic Intelligence for continual learning.
    
    SI measures parameter importance by tracking the contribution 
    of each parameter to the loss reduction during training.
    
    Key equations:
        ω_k = Σ_t (∂L/∂θ_k)_t * Δθ_k^t  (online importance)
        Ω_k = ω_k / ((Δθ_k)² + ξ)        (normalized importance)
    
    Reference: Zenke et al., "Continual Learning Through Synaptic Intelligence"
    """
    
    def __init__(
        self, 
        model: nn.Module,
        c: float = 0.1,           # Regularization strength
        xi: float = 1e-3          # Damping constant
    ):
        """
        Args:
            model: Neural network model
            c: Coefficient for importance regularization (like λ in EWC)
            xi: Small constant for numerical stability
        """
        self.model = model
        self.c = c
        self.xi = xi
        
        # Parameter tracking
        self.params: Dict[str, torch.Tensor] = {}
        self.prev_params: Dict[str, torch.Tensor] = {}
        
        # Importance accumulator (ω in the paper)
        self.omega_accumulator: Dict[str, torch.Tensor] = {}
        
        # Consolidated importance (Ω in the paper)
        self.importance: Dict[str, torch.Tensor] = {}
        
        # Reference parameters from previous task
        self.reference_params: Dict[str, torch.Tensor] = {}
        
        self._initialize_tracking()
        
    def _initialize_tracking(self) -> None:
        """Initialize parameter tracking dictionaries."""
        for n, p in self.model.named_parameters():
            if p.requires_grad:
                self.params[n] = p
                self.prev_params[n] = p.data.clone()
                self.omega_accumulator[n] = torch.zeros_like(p.data)
                
    def update_omega(self) -> None:
        """
        Update importance accumulator after each training step.
        
        Called after optimizer.step() to track:
            ω_k += -gradient_k * Δθ_k
        
        Note: We use negative gradient because parameter update is
        in the direction of -gradient, so contribution to loss reduction
        is -gradient * Δθ.
        """
        for n, p in self.model.named_parameters():
            if n in self.omega_accumulator and p.grad is not None:
                # Parameter change in this step
                delta = p.data - self.prev_params[n]
                
                # Accumulate contribution: -gradient * delta
                # Gradient points toward increasing loss, so contribution
                # to decreasing loss is in the negative gradient direction
                self.omega_accumulator[n] += (-p.grad.data * delta).clamp(min=0)
                
                # Update previous parameter value
                self.prev_params[n] = p.data.clone()
                
    def consolidate(self) -> None:
        """
        Consolidate importance after completing a task.
        
        Computes normalized importance:
            Ω_k = ω_k / ((Δθ_k)² + ξ)
        
        And adds to cumulative importance from previous tasks.
        """
        for n, p in self.model.named_parameters():
            if n in self.omega_accumulator:
                # Total parameter change during this task
                delta_squared = (p.data - self.reference_params.get(
                    n, self.prev_params[n]
                )) ** 2
                
                # Compute normalized importance for this task
                task_importance = self.omega_accumulator[n] / (delta_squared + self.xi)
                
                # Add to cumulative importance
                if n in self.importance:
                    self.importance[n] += task_importance
                else:
                    self.importance[n] = task_importance
                    
                # Reset accumulator for next task
                self.omega_accumulator[n].zero_()
                
        # Update reference parameters to current values
        self.reference_params = {
            n: p.data.clone() 
            for n, p in self.model.named_parameters()
            if p.requires_grad
        }
        
    def penalty(self) -> torch.Tensor:
        """
        Compute SI regularization penalty.
        
        Returns:
            c * Σ_k Ω_k * (θ_k - θ*_k)²
        """
        if not self.importance:
            return torch.tensor(0.0)
            
        device = next(self.model.parameters()).device
        loss = torch.tensor(0.0).to(device)
        
        for n, p in self.model.named_parameters():
            if n in self.importance:
                ref = self.reference_params[n]
                loss += (self.importance[n] * (p - ref) ** 2).sum()
                
        return self.c * loss
    
    def training_step(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        criterion: nn.Module,
        optimizer: torch.optim.Optimizer
    ) -> float:
        """
        Single SI training step with importance tracking.
        
        IMPORTANT: Call update_omega() after optimizer.step()!
        """
        optimizer.zero_grad()
        
        outputs = self.model(inputs)
        task_loss = criterion(outputs, targets)
        si_loss = self.penalty()
        total_loss = task_loss + si_loss
        
        total_loss.backward()
        optimizer.step()
        
        # CRITICAL: Update omega after parameter update
        self.update_omega()
        
        return total_loss.item()
 
 
def compare_ewc_si():
    """Side-by-side comparison of importance computation."""
    
    print("EWC vs SI Importance Computation")
    print("=" * 50)
    print()
    print("EWC (Fisher Information):")
    print("  - Computed AFTER training on task")
    print("  - Based on: curvature of likelihood")
    print("  - Formula: F = E[∇log p(y|x)²]")
    print("  - Requires: separate pass over data")
    print()
    print("SI (Path Integral):")
    print("  - Computed DURING training")
    print("  - Based on: contribution to loss reduction")
    print("  - Formula: ω = Σ (-∇L) · Δθ")
    print("  - Requires: tracking at each step")
    
compare_ewc_si()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from typing import Dict
 
class MemoryAwareSynapses:
    """
    Memory Aware Synapses (MAS) for continual learning.
    
    MAS computes importance based on output sensitivity:
    how much the network output changes when parameters change.
    
    Key equation:
        Ω_k = (1/N) Σ_n || ∂F(x_n)/∂θ_k ||
    
    This is label-free, making MAS suitable for unsupervised learning.
    
    Reference: Aljundi et al., "Memory Aware Synapses"
    """
    
    def __init__(
        self,
        model: nn.Module,
        lambda_reg: float = 1.0
    ):
        """
        Args:
            model: Neural network model
            lambda_reg: Regularization strength
        """
        self.model = model
        self.lambda_reg = lambda_reg
        
        self.importance: Dict[str, torch.Tensor] = {}
        self.reference_params: Dict[str, torch.Tensor] = {}
        
    def compute_importance(
        self,
        dataloader: DataLoader,
        num_samples: int = 2000
    ) -> None:
        """
        Compute MAS importance weights.
        
        For each parameter, we compute the average gradient magnitude
        of the output with respect to that parameter.
        
        Args:
            dataloader: Data to compute importance over
            num_samples: Maximum samples to use
        """
        self.model.eval()
        
        # Initialize importance accumulators
        importance = {
            n: torch.zeros_like(p)
            for n, p in self.model.named_parameters()
            if p.requires_grad
        }
        
        samples_seen = 0
        
        for inputs, _ in dataloader:  # Labels not needed!
            if samples_seen >= num_samples:
                break
                
            inputs = inputs.to(next(self.model.parameters()).device)
            batch_size = inputs.size(0)
            
            # Forward pass
            outputs = self.model(inputs)
            
            # Compute gradient of output magnitude w.r.t. parameters
            # We use L2 norm of output as a scalar measure
            output_norm = outputs.norm(dim=-1).sum()
            
            self.model.zero_grad()
            output_norm.backward()
            
            # Accumulate absolute gradients
            for n, p in self.model.named_parameters():
                if p.requires_grad and p.grad is not None:
                    importance[n] += p.grad.data.abs()
                    
            samples_seen += batch_size
            
        # Normalize by number of samples
        for n in importance:
            importance[n] /= samples_seen
            
        # Update cumulative importance
        if not self.importance:
            self.importance = importance
        else:
            for n in importance:
                self.importance[n] += importance[n]
                
        # Save reference parameters
        self.reference_params = {
            n: p.data.clone()
            for n, p in self.model.named_parameters()
            if p.requires_grad
        }
        
        self.model.train()
        
    def penalty(self) -> torch.Tensor:
        """
        Compute MAS regularization penalty.
        
        Returns:
            λ * Σ_k Ω_k * (θ_k - θ*_k)²
        """
        if not self.importance:
            return torch.tensor(0.0)
            
        device = next(self.model.parameters()).device
        loss = torch.tensor(0.0).to(device)
        
        for n, p in self.model.named_parameters():
            if n in self.importance:
                ref = self.reference_params[n]
                loss += (self.importance[n] * (p - ref) ** 2).sum()
                
        return self.lambda_reg * loss
 
 
class GradientEpisodicMemory:
    """
    Gradient Episodic Memory (GEM) - a constraint-based approach.
    
    Instead of soft regularization, GEM projects gradients to ensure
    they don't increase loss on previous tasks.
    
    This is a fundamentally different approach: rather than penalizing
    weight changes, GEM directly modifies gradients.
    """
    
    def __init__(
        self,
        model: nn.Module,
        memory_per_task: int = 256,
        margin: float = 0.5
    ):
        self.model = model
        self.memory_per_task = memory_per_task
        self.margin = margin
        
        self.episodic_memory: Dict[int, tuple] = {}  # task_id -> (inputs, targets)
        self.reference_gradients: Dict[int, Dict[str, torch.Tensor]] = {}
        
    def store_memory(self, task_id: int, inputs: torch.Tensor, targets: torch.Tensor):
        """Store exemplars for gradient projection."""
        self.episodic_memory[task_id] = (inputs, targets)
        
    def project_gradient(self, current_grad: Dict[str, torch.Tensor]):
        """
        Project current gradient to not conflict with previous tasks.
        
        Uses quadratic programming to find closest gradient that
        doesn't increase loss on any previous task.
        """
        # Implement gradient projection logic
        # This modifies gradients in-place
        pass

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict
import copy
 
class LearningWithoutForgetting:
    """
    Learning without Forgetting (LwF) for continual learning.
    
    LwF preserves behavior on old tasks through knowledge distillation.
    Before training on a new task, it records the model's outputs and
    uses them as soft targets during training.
    
    Key insight: Soft targets (logits/probabilities) contain more 
    information than hard labels, preserving inter-class relationships.
    
    Loss = L_new(input, label) + λ * L_distill(current_output_old, frozen_output_old)
    """
    
    def __init__(
        self,
        model: nn.Module,
        temperature: float = 2.0,
        lambda_distill: float = 1.0
    ):
        """
        Args:
            model: Neural network model
            temperature: Softmax temperature for distillation (higher = softer)
            lambda_distill: Weight for distillation loss
        """
        self.model = model
        self.temperature = temperature
        self.lambda_distill = lambda_distill
        
        self.old_model: Optional[nn.Module] = None
        self.old_task_output_dim = 0
        
    def prepare_for_new_task(self) -> None:
        """
        Prepare for training on a new task.
        
        Creates a frozen copy of the current model to serve as
        the 'teacher' for knowledge distillation.
        """
        # Create frozen copy of current model
        self.old_model = copy.deepcopy(self.model)
        self.old_model.eval()
        
        # Freeze all parameters
        for param in self.old_model.parameters():
            param.requires_grad = False
            
        # Track output dimension for old tasks
        # (Assumes model has an 'output_dim' or similar attribute)
        self.old_task_output_dim = self._get_output_dim()
        
    def _get_output_dim(self) -> int:
        """Get output dimension of the model."""
        # Implementation depends on model architecture
        for module in self.model.modules():
            if isinstance(module, nn.Linear):
                last_linear = module
        return last_linear.out_features
    
    def distillation_loss(
        self,
        current_logits: torch.Tensor,
        old_logits: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute knowledge distillation loss.
        
        Uses softmax with temperature to create soft targets,
        then computes KL divergence.
        
        Temperature > 1 creates softer probability distributions,
        which transfers more information about relative confidences.
        """
        T = self.temperature
        
        # Soft probabilities from old model (teacher)
        soft_targets = F.softmax(old_logits / T, dim=1)
        
        # Log-soft probabilities from current model (student)
        log_soft_current = F.log_softmax(current_logits / T, dim=1)
        
        # KL divergence loss (scaled by T²)
        # The T² scaling is important for gradient magnitudes
        loss = F.kl_div(
            log_soft_current,
            soft_targets,
            reduction='batchmean'
        ) * (T ** 2)
        
        return loss
    
    def compute_loss(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        criterion: nn.Module
    ) -> torch.Tensor:
        """
        Compute combined LwF loss.
        
        Returns:
            L_task + λ * L_distillation
        """
        # Forward through current model
        outputs = self.model(inputs)
        
        # Task loss for new task (only on new output heads)
        task_loss = criterion(outputs, targets)
        
        # Distillation loss for old tasks
        if self.old_model is not None:
            with torch.no_grad():
                old_outputs = self.old_model(inputs)
                
            # Only distill on old task output dimensions
            current_old = outputs[:, :self.old_task_output_dim]
            distill_loss = self.distillation_loss(current_old, old_outputs)
            
            total_loss = task_loss + self.lambda_distill * distill_loss
        else:
            total_loss = task_loss
            
        return total_loss
 
 
class LwFMultiTask:
    """
    LwF variant for class-incremental learning with shared head.
    
    In class-incremental learning, new classes are added over time.
    This changes the output layer architecture, requiring careful
    handling of distillation.
    """
    
    def __init__(
        self,
        feature_extractor: nn.Module,
        initial_classes: int,
        temperature: float = 2.0,
        lambda_distill: float = 1.0
    ):
        self.feature_extractor = feature_extractor
        self.temperature = temperature
        self.lambda_distill = lambda_distill
        
        # Get feature dimension
        with torch.no_grad():
            dummy = torch.randn(1, 3, 32, 32)
            feat_dim = feature_extractor(dummy).shape[1]
            
        # Output layer (grows with new classes)
        self.classifier = nn.Linear(feat_dim, initial_classes)
        self.num_classes = initial_classes
        
        # Frozen copy for distillation
        self.old_feature_extractor = None
        self.old_classifier = None
        self.old_num_classes = 0
        
    def add_classes(self, new_classes: int) -> None:
        """
        Expand classifier for new classes while preserving old weights.
        """
        # Store old model
        self.old_feature_extractor = copy.deepcopy(self.feature_extractor)
        self.old_classifier = copy.deepcopy(self.classifier)
        self.old_num_classes = self.num_classes
        
        # Freeze old model
        for p in self.old_feature_extractor.parameters():
            p.requires_grad = False
        for p in self.old_classifier.parameters():
            p.requires_grad = False
            
        # Create expanded classifier
        old_weight = self.classifier.weight.data
        old_bias = self.classifier.bias.data
        
        new_total = self.num_classes + new_classes
        new_classifier = nn.Linear(old_weight.shape[1], new_total)
        
        # Copy old weights
        new_classifier.weight.data[:self.num_classes] = old_weight
        new_classifier.bias.data[:self.num_classes] = old_bias
        
        self.classifier = new_classifier
        self.num_classes = new_total
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.feature_extractor(x)
        return self.classifier(features)
    
    def compute_loss(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        criterion: nn.Module
    ) -> torch.Tensor:
        """Combined task and distillation loss."""
        # Current model predictions
        features = self.feature_extractor(inputs)
        outputs = self.classifier(features)
        
        # Task loss
        task_loss = criterion(outputs, targets)
        
        # Distillation loss on old class outputs
        if self.old_feature_extractor is not None:
            with torch.no_grad():
                old_features = self.old_feature_extractor(inputs)
                old_outputs = self.old_classifier(old_features)
                
            current_old = outputs[:, :self.old_num_classes]
            distill_loss = F.kl_div(
                F.log_softmax(current_old / self.temperature, dim=1),
                F.softmax(old_outputs / self.temperature, dim=1),
                reduction='batchmean'
            ) * (self.temperature ** 2)
            
            return task_loss + self.lambda_distill * distill_loss
            
        return task_loss