File size: 5,326 Bytes
f9abec3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
Synthetic continual-learning benchmark — fully offline (no downloads).

Conflicting-teacher stream (the setting where forgetting is strong and real):

  - Inputs x ~ N(0, I) in R^dim, shared across all tasks.
  - Each task t has its OWN random teacher hyperplane w_t: label y = 1[ w_t·x > 0 ].
  - A SINGLE shared head is used for every task.

Because each task demands a DIFFERENT decision boundary, learning task t+1
genuinely overwrites the weights that solved task t → catastrophic forgetting.
Consolidation (EWC / saṃskāra) protects the weights that mattered for earlier
teachers, trading some new-task fit for retained old-task accuracy.

Deliberately small — the whole study runs on CPU in well under a minute.
"""
from __future__ import annotations

from typing import List, Tuple

import torch
from torch.utils.data import TensorDataset, DataLoader


def make_tasks(
    n_tasks: int = 6,
    dim: int = 32,
    n_train: int = 1000,
    n_test: int = 400,
    margin: float = 0.15,   # drop points too close to the boundary (cleaner labels)
    seed: int = 0,
) -> List[Tuple[TensorDataset, TensorDataset]]:
    g = torch.Generator().manual_seed(seed)
    tasks = []
    for _ in range(n_tasks):
        w = torch.randn(dim, generator=g)
        w = w / w.norm()
        train = _teacher(w, n_train, margin, dim, g)
        test = _teacher(w, n_test, margin, dim, g)
        tasks.append((train, test))
    return tasks


def _teacher(w, n, margin, dim, g) -> TensorDataset:
    xs, ys = [], []
    while sum(len(y) for y in ys) < n:
        x = torch.randn(n, dim, generator=g)
        s = x @ w
        keep = s.abs() > margin                # enforce a margin → clean labels
        xs.append(x[keep]); ys.append((s[keep] > 0).long())
    x = torch.cat(xs)[:n]; y = torch.cat(ys)[:n]
    return TensorDataset(x, y)


def loaders(tasks, batch_size: int = 64):
    out = []
    for train, test in tasks:
        out.append((
            DataLoader(train, batch_size=batch_size, shuffle=True),
            DataLoader(test, batch_size=256, shuffle=False),
        ))
    return out


# ---------------------------------------------------------------------------
# Capacity-headroom stream: tasks share a common nonlinear feature map φ (one
# fixed random teacher MLP), and differ only in a per-task linear readout. A
# backbone that learns φ can solve ALL tasks jointly (so avg accuracy can be
# high — there is real headroom), but training a new task can overwrite the
# parts of φ the old tasks relied on (forgetting). This is the setting where
# consolidation/replay/decay/tapas can actually differentiate.
# Use with a MULTI-HEAD model: evaluate task k with head k.
# ---------------------------------------------------------------------------
def make_shared_feature_tasks(
    n_tasks: int = 8,
    dim: int = 24,
    teacher_hidden: int = 32,
    n_train: int = 800,
    n_test: int = 400,
    margin: float = 0.2,
    seed: int = 0,
):
    g = torch.Generator().manual_seed(seed)
    # one shared, fixed nonlinear teacher  φ: R^dim -> R^teacher_hidden
    W1 = torch.randn(dim, teacher_hidden, generator=g) / (dim ** 0.5)
    b1 = torch.randn(teacher_hidden, generator=g) * 0.1

    def phi(x):
        return torch.tanh(x @ W1 + b1)

    tasks = []
    for _ in range(n_tasks):
        u = torch.randn(teacher_hidden, generator=g)          # per-task readout
        u = u / u.norm()
        train = _shared(phi, u, n_train, dim, margin, g)
        test = _shared(phi, u, n_test, dim, margin, g)
        tasks.append((train, test))
    return tasks, dim


def make_difficulty_tasks(difficulty: float, n_tasks: int = 5, dim: int = 24,
                          n_train: int = 400, n_test: int = 200, margin: float = 0.1,
                          seed: int = 0):
    """Single-head continual stream with a tunable FORGETTING knob.

    Each task's teacher is w_t = normalize((1-d)·w0 + d·r_t):
      d=0  → every task identical (no forgetting; protection only hurts)
      d=1  → every task a fresh random boundary (severe forgetting; protection helps)
    The controller must read the resulting forgetting and adapt. Returns task list
    + dim; use with a SINGLE-head model."""
    g = torch.Generator().manual_seed(seed)
    w0 = torch.randn(dim, generator=g); w0 = w0 / w0.norm()
    tasks = []
    for _ in range(n_tasks):
        r = torch.randn(dim, generator=g)
        w = (1 - difficulty) * w0 + difficulty * r
        w = w / w.norm()
        tr = _hyper(w, n_train, dim, margin, g)
        te = _hyper(w, n_test, dim, margin, g)
        tasks.append((tr, te))
    return tasks, dim


def _hyper(w, n, dim, margin, g) -> TensorDataset:
    xs, ys = [], []
    while sum(len(y) for y in ys) < n:
        x = torch.randn(n, dim, generator=g)
        s = x @ w
        keep = s.abs() > margin
        xs.append(x[keep]); ys.append((s[keep] > 0).long())
    return TensorDataset(torch.cat(xs)[:n], torch.cat(ys)[:n])


def _shared(phi, u, n, dim, margin, g) -> TensorDataset:
    xs, ys = [], []
    while sum(len(y) for y in ys) < n:
        x = torch.randn(n, dim, generator=g)
        s = phi(x) @ u
        keep = s.abs() > margin
        xs.append(x[keep]); ys.append((s[keep] > 0).long())
    x = torch.cat(xs)[:n]; y = torch.cat(ys)[:n]
    return TensorDataset(x, y)