Meta-Learning: Learning to Learn

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

import torch
import torch.nn as nn
from torch.nn import functional as F
from typing import List, Tuple, Dict
import copy
 
class MAML:
    """
    Model-Agnostic Meta-Learning (MAML) implementation.
    
    Key insight: Learn an initialization from which a few gradient steps
    produce excellent task-specific parameters.
    """
    
    def __init__(
        self,
        model: nn.Module,
        inner_lr: float = 0.01,      # α: task adaptation learning rate
        outer_lr: float = 0.001,     # β: meta-learning rate
        inner_steps: int = 5,        # Number of gradient steps for adaptation
        first_order: bool = False    # Use first-order approximation?
    ):
        self.model = model
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr
        self.inner_steps = inner_steps
        self.first_order = first_order
        
        # Meta-optimizer updates the initialization
        self.meta_optimizer = torch.optim.Adam(model.parameters(), lr=outer_lr)
    
    def inner_loop(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        params: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Inner loop: Adapt to a specific task.
        
        Takes K gradient steps on the support set, starting from 
        the current meta-parameters.
        
        Args:
            support_x: Support set inputs [n_support, ...]
            support_y: Support set labels [n_support]
            params: Current parameters (shared initialization)
        
        Returns:
            adapted_params: Task-specific parameters after adaptation
        """
        adapted_params = {k: v.clone() for k, v in params.items()}
        
        for step in range(self.inner_steps):
            # Forward pass with current adapted parameters
            logits = self.functional_forward(support_x, adapted_params)
            loss = F.cross_entropy(logits, support_y)
            
            # Compute gradients w.r.t. adapted parameters
            grads = torch.autograd.grad(
                loss, 
                adapted_params.values(),
                create_graph=not self.first_order  # Need graph for 2nd order
            )
            
            # Gradient descent update
            adapted_params = {
                k: adapted_params[k] - self.inner_lr * g
                for (k, _), g in zip(adapted_params.items(), grads)
            }
        
        return adapted_params
    
    def outer_loop(
        self, 
        task_batch: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]
    ) -> float:
        """
        Outer loop: Update initialization for better adaptability.
        
        For each task:
        1. Adapt parameters using support set (inner loop)
        2. Evaluate adapted parameters on query set
        3. Accumulate meta-gradients
        
        Key: Gradient flows THROUGH the inner loop adaptation.
        
        Args:
            task_batch: List of (support_x, support_y, query_x, query_y) tuples
        
        Returns:
            Mean meta-loss across tasks
        """
        self.meta_optimizer.zero_grad()
        meta_loss = 0.0
        
        # Get current parameters as a dictionary
        params = dict(self.model.named_parameters())
        
        for support_x, support_y, query_x, query_y in task_batch:
            # Inner loop: task-specific adaptation
            adapted_params = self.inner_loop(support_x, support_y, params)
            
            # Evaluate adapted parameters on query set
            query_logits = self.functional_forward(query_x, adapted_params)
            task_loss = F.cross_entropy(query_logits, query_y)
            
            meta_loss += task_loss
        
        # Average over task batch
        meta_loss = meta_loss / len(task_batch)
        
        # Compute meta-gradients and update
        # Gradients flow through inner_loop adaptation!
        meta_loss.backward()
        self.meta_optimizer.step()
        
        return meta_loss.item()
    
    def functional_forward(
        self, 
        x: torch.Tensor, 
        params: Dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """
        Forward pass using specified parameters.
        
        This allows using adapted parameters without modifying model state.
        Implementation depends on model architecture.
        """
        # Example for a simple CNN with named parameters
        # Actual implementation requires model-specific handling
        
        # For Linear layers: F.linear(x, params['layer.weight'], params['layer.bias'])
        # For Conv layers: F.conv2d(x, params['conv.weight'], params['conv.bias'], ...)
        
        # Many implementations use higher-level libraries like 'higher' or 'learn2learn'
        raise NotImplementedError("Implement based on model architecture")
    
    def adapt_and_evaluate(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        query_x: torch.Tensor,
        query_y: torch.Tensor
    ) -> Tuple[float, float]:
        """
        Adapt to a new task and evaluate.
        
        Used at test time to evaluate few-shot performance.
        """
        params = dict(self.model.named_parameters())
        
        # Adapt to task
        adapted_params = self.inner_loop(support_x, support_y, params)
        
        # Evaluate on query
        with torch.no_grad():
            query_logits = self.functional_forward(query_x, adapted_params)
            loss = F.cross_entropy(query_logits, query_y)
            predictions = query_logits.argmax(dim=1)
            accuracy = (predictions == query_y).float().mean()
        
        return loss.item(), accuracy.item()

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

import torch
import torch.nn as nn
from typing import List, Tuple
import copy
 
class Reptile:
    """
    Reptile: A simpler alternative to MAML.
    
    Key insight: Instead of computing meta-gradients through adaptation,
    simply move the initialization toward adapted parameters.
    
    Update rule: θ ← θ + ε(θ'_i - θ)
    
    This is equivalent to averaging across task-specific fine-tuned
    parameters, providing an implicit regularization effect.
    """
    
    def __init__(
        self,
        model: nn.Module,
        inner_lr: float = 0.01,       # Learning rate for task adaptation
        outer_lr: float = 0.001,      # Meta-learning rate (step size toward θ')
        inner_steps: int = 5,         # SGD steps per task
    ):
        self.model = model
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr
        self.inner_steps = inner_steps
    
    def train_step(
        self,
        task_batch: List[Tuple[torch.Tensor, torch.Tensor]]
    ) -> float:
        """
        Reptile training step.
        
        For each task:
        1. Clone current parameters
        2. Take K gradient steps on task data
        3. Compute θ' - θ (direction toward task optimum)
        
        Then: Move θ toward average task-adapted parameters
        """
        # Store original parameters
        original_weights = {
            name: param.clone() 
            for name, param in self.model.named_parameters()
        }
        
        # Accumulate (θ' - θ) across tasks
        weight_diffs = {
            name: torch.zeros_like(param) 
            for name, param in self.model.named_parameters()
        }
        
        total_loss = 0.0
        
        for support_x, support_y in task_batch:
            # Reset to original weights before each task
            for name, param in self.model.named_parameters():
                param.data.copy_(original_weights[name])
            
            # Task adaptation: K gradient descent steps
            optimizer = torch.optim.SGD(
                self.model.parameters(), 
                lr=self.inner_lr
            )
            
            for step in range(self.inner_steps):
                optimizer.zero_grad()
                logits = self.model(support_x)
                loss = nn.functional.cross_entropy(logits, support_y)
                loss.backward()
                optimizer.step()
                
                if step == self.inner_steps - 1:
                    total_loss += loss.item()
            
            # Accumulate difference: θ'_i - θ
            for name, param in self.model.named_parameters():
                weight_diffs[name] += param.data - original_weights[name]
        
        # Average across tasks
        n_tasks = len(task_batch)
        
        # Update original weights: θ ← θ + ε * mean(θ'_i - θ)
        for name, param in self.model.named_parameters():
            param.data.copy_(
                original_weights[name] + 
                self.outer_lr * weight_diffs[name] / n_tasks
            )
        
        return total_loss / n_tasks
    
    def adapt(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        adapt_steps: int = None
    ):
        """
        Adapt to a new task at test time.
        
        Simply fine-tune for K steps (same as inner loop during training).
        """
        steps = adapt_steps or self.inner_steps
        optimizer = torch.optim.SGD(
            self.model.parameters(),
            lr=self.inner_lr
        )
        
        for _ in range(steps):
            optimizer.zero_grad()
            logits = self.model(support_x)
            loss = nn.functional.cross_entropy(logits, support_y)
            loss.backward()
            optimizer.step()
 
# Theoretical connection between Reptile and MAML
"""
Reptile approximates MAML in expectation:
 
For a single inner step:
    Reptile update: θ ← θ + ε(θ - α∇L(θ) - θ) = θ - εα∇L(θ)
    
This is just SGD! The magic happens with multiple inner steps.
 
With K inner steps, Reptile's update direction includes:
1. Average gradient (like SGD)
2. Curvature-aware terms (like second-order methods)
 
The key insight: Moving toward task-adapted parameters implicitly
captures second-order information about the loss landscape.
 
Empirically, Reptile performs comparably to FOMAML and sometimes
approaches full MAML, while being simpler to implement.
"""

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

import torch
import torch.nn as nn
from typing import Dict, List
 
class MAMLPlusPlus:
    """
    MAML++ with practical improvements for stability and performance.
    """
    
    def __init__(
        self,
        model: nn.Module,
        inner_lr_init: float = 0.01,
        outer_lr: float = 0.001,
        inner_steps: int = 5,
        learn_inner_lr: bool = True,
        per_layer_lr: bool = True,
        multi_step_loss: bool = True,
        msl_decay: float = 0.5,        # Weight decay for earlier steps
        anneal_second_order: bool = True,
        anneal_epochs: int = 50,
    ):
        self.model = model
        self.inner_steps = inner_steps
        self.multi_step_loss = multi_step_loss
        self.msl_decay = msl_decay
        self.anneal_second_order = anneal_second_order
        self.anneal_epochs = anneal_epochs
        self.current_epoch = 0
        
        # Learnable per-layer inner learning rates
        if per_layer_lr and learn_inner_lr:
            # One learning rate per layer
            self.inner_lrs = nn.ParameterDict({
                name.replace('.', '_'): nn.Parameter(torch.tensor(inner_lr_init))
                for name, _ in model.named_parameters()
            })
        elif learn_inner_lr:
            # Single learnable learning rate
            self.inner_lrs = nn.Parameter(torch.tensor(inner_lr_init))
        else:
            self.inner_lrs = inner_lr_init
        
        # Meta-optimizer includes inner learning rates if learnable
        meta_params = list(model.parameters())
        if isinstance(self.inner_lrs, nn.ParameterDict):
            meta_params.extend(self.inner_lrs.values())
        elif isinstance(self.inner_lrs, nn.Parameter):
            meta_params.append(self.inner_lrs)
            
        self.meta_optimizer = torch.optim.Adam(meta_params, lr=outer_lr)
    
    def get_inner_lr(self, param_name: str) -> torch.Tensor:
        """Get the inner learning rate for a specific parameter."""
        if isinstance(self.inner_lrs, nn.ParameterDict):
            key = param_name.replace('.', '_')
            return self.inner_lrs[key]
        elif isinstance(self.inner_lrs, nn.Parameter):
            return self.inner_lrs
        else:
            return torch.tensor(self.inner_lrs)
    
    @property
    def use_second_order(self) -> bool:
        """Determine whether to use second-order gradients."""
        if not self.anneal_second_order:
            return True
        # Linear annealing from 0 to 1 over anneal_epochs
        return self.current_epoch >= self.anneal_epochs
    
    def multi_step_loss_weights(self) -> List[float]:
        """
        Compute weights for multi-step loss optimization.
        
        Earlier steps get lower weights, final step gets highest.
        This encourages the model to optimize adaptation trajectory.
        """
        weights = []
        for k in range(self.inner_steps):
            # Exponential decay: later steps weighted more
            weight = self.msl_decay ** (self.inner_steps - k - 1)
            weights.append(weight)
        
        # Normalize
        total = sum(weights)
        return [w / total for w in weights]
    
    def inner_loop_with_msl(
        self,
        support_x: torch.Tensor,
        support_y: torch.Tensor,
        params: Dict[str, torch.Tensor]
    ) -> tuple:
        """
        Inner loop with multi-step loss tracking.
        
        Returns adapted parameters AND list of losses at each step.
        """
        adapted_params = {k: v.clone() for k, v in params.items()}
        step_losses = []
        
        for step in range(self.inner_steps):
            logits = self.functional_forward(support_x, adapted_params)
            loss = nn.functional.cross_entropy(logits, support_y)
            step_losses.append(loss)
            
            grads = torch.autograd.grad(
                loss,
                adapted_params.values(),
                create_graph=self.use_second_order
            )
            
            # Per-parameter learning rates
            adapted_params = {}
            for (name, param), g in zip(params.items(), grads):
                lr = self.get_inner_lr(name)
                adapted_params[name] = param - lr * g
        
        return adapted_params, step_losses
    
    def outer_loop(self, task_batch: List) -> float:
        """
        Outer loop with multi-step loss and other MAML++ improvements.
        """
        self.meta_optimizer.zero_grad()
        meta_loss = 0.0
        
        msl_weights = self.multi_step_loss_weights()
        params = dict(self.model.named_parameters())
        
        for support_x, support_y, query_x, query_y in task_batch:
            adapted_params, step_losses = self.inner_loop_with_msl(
                support_x, support_y, params
            )
            
            # Query loss
            query_logits = self.functional_forward(query_x, adapted_params)
            query_loss = nn.functional.cross_entropy(query_logits, query_y)
            
            if self.multi_step_loss:
                # Weighted combination of inner step losses + query loss
                task_loss = query_loss
                for weight, step_loss in zip(msl_weights, step_losses):
                    task_loss = task_loss + 0.1 * weight * step_loss
            else:
                task_loss = query_loss
            
            meta_loss += task_loss
        
        meta_loss = meta_loss / len(task_batch)
        meta_loss.backward()
        self.meta_optimizer.step()
        
        return meta_loss.item()
    
    def functional_forward(self, x, params):
        """Model-specific functional forward pass."""
        raise NotImplementedError

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

import torch
import learn2learn as l2l
from learn2learn.algorithms import MAML
 
def train_maml(
    model,
    train_dataset,
    val_dataset,
    n_epochs: int = 100,
    tasks_per_batch: int = 32,
    n_way: int = 5,
    k_shot: int = 5,
    q_query: int = 15,
    inner_lr: float = 0.01,
    outer_lr: float = 0.001,
    inner_steps: int = 5,
    first_order: bool = False,
    grad_clip: float = 10.0,
):
    """
    Complete MAML training loop with best practices.
    
    Uses learn2learn library for clean implementation.
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    # Wrap model with MAML
    maml = MAML(model, lr=inner_lr, first_order=first_order)
    
    # Meta-optimizer
    optimizer = torch.optim.Adam(maml.parameters(), lr=outer_lr)
    
    # Task samplers
    train_tasks = l2l.data.TaskDataset(
        train_dataset,
        task_transforms=[
            l2l.data.transforms.NWays(train_dataset, n_way),
            l2l.data.transforms.KShots(train_dataset, k_shot + q_query),
            l2l.data.transforms.LoadData(train_dataset),
        ],
        num_tasks=-1  # Infinite
    )
    
    best_val_acc = 0.0
    patience = 10
    patience_counter = 0
    
    for epoch in range(n_epochs):
        model.train()
        epoch_loss = 0.0
        epoch_acc = 0.0
        
        for task_idx in range(tasks_per_batch):
            # Sample task
            batch = train_tasks.sample()
            
            # Split into support and query
            support_x = batch[0][:n_way * k_shot].to(device)
            support_y = batch[1][:n_way * k_shot].to(device)
            query_x = batch[0][n_way * k_shot:].to(device)
            query_y = batch[1][n_way * k_shot:].to(device)
            
            # Clone for task-specific adaptation
            learner = maml.clone()
            
            # Inner loop: adapt to task
            for step in range(inner_steps):
                support_logits = learner(support_x)
                support_loss = torch.nn.functional.cross_entropy(
                    support_logits, support_y
                )
                learner.adapt(support_loss)
            
            # Outer loop: evaluate on query, accumulate gradients
            query_logits = learner(query_x)
            query_loss = torch.nn.functional.cross_entropy(
                query_logits, query_y
            )
            
            epoch_loss += query_loss.item()
            
            # Accuracy
            preds = query_logits.argmax(dim=1)
            epoch_acc += (preds == query_y).float().mean().item()
            
            # Accumulate gradients
            query_loss.backward()
        
        # Meta-update with gradient clipping
        torch.nn.utils.clip_grad_norm_(maml.parameters(), grad_clip)
        optimizer.step()
        optimizer.zero_grad()
        
        # Logging
        avg_loss = epoch_loss / tasks_per_batch
        avg_acc = epoch_acc / tasks_per_batch
        print(f"Epoch {epoch + 1}: Loss={avg_loss:.4f}, Acc={avg_acc:.2%}")
        
        # Validation
        if (epoch + 1) % 5 == 0:
            val_acc = evaluate_maml(maml, val_dataset, n_way, k_shot, 
                                    q_query, inner_steps, device)
            print(f"  Validation Acc: {val_acc:.2%}")
            
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(maml.state_dict(), 'best_maml.pt')
                patience_counter = 0
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print("Early stopping!")
                    break
    
    return maml