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
|
"""Matched non-backprop adapters for the audited CIFAR ResNet topology.
These adapters reuse :class:`CIFARLocalResNet`'s forward tensors, BatchNorm,
option-A shortcuts, and local optimizer. They are kept separate from the
established SDIL/KP implementation so a crossover baseline cannot silently
change the already confirmed forward model.
"""
import torch
import torch.nn.functional as F
from .conv import CIFARLocalResNet
class CIFARDualPropResNet(CIFARLocalResNet):
"""Dual Propagation on the residual DAG with author DP-transpose updates.
``s_plus`` and ``s_minus`` are initialized to the ordinary forward states.
A ``fwK`` pass updates every residual-DAG node in topological order. The
feedforward drive uses the forward edge, while the difference of each
child state is transported through the exact transpose of that same edge.
This is intentional symmetric feedback in the Dual Propagation baseline,
not a claim of weight-transport-free learning.
"""
def __init__(self, *args, alpha=0.0, dp_beta=0.1,
inference_passes=16, **kwargs):
super().__init__(*args, **kwargs)
if not 0.0 <= alpha <= 1.0:
raise ValueError("Dual Propagation alpha must lie in [0, 1]")
if dp_beta <= 0 or inference_passes < 1:
raise ValueError("invalid Dual Propagation inference settings")
self.dp_alpha = float(alpha)
self.dp_beta = float(dp_beta)
self.dp_inference_passes = int(inference_passes)
self._node_kind = {0: ("stem",)}
self._outgoing = {index: [] for index in range(self.n_hidden)}
for block in self.blocks:
first = block["first"]
second = block["second"]
parent = first - 1
self._node_kind[first] = ("first", parent)
self._node_kind[second] = (
"second", first, parent, block["out_channels"],
block["stride"])
self._outgoing[parent].append((first, "conv"))
self._outgoing[first].append((second, "conv"))
self._outgoing[parent].append((second, "shortcut"))
@staticmethod
def _option_a_shortcut_transpose(value, input_shape, stride):
"""Adjoint of the parameter-free option-A shortcut."""
in_channels = input_shape[1]
out_channels = value.shape[1]
missing = out_channels - in_channels
if missing < 0:
raise ValueError("option-A transpose cannot increase input width")
before = missing // 2
selected = value[:, before:before + in_channels]
if stride == 1:
if tuple(selected.shape) != tuple(input_shape):
raise ValueError("option-A transpose shape mismatch")
return selected
result = value.new_zeros(input_shape)
result[:, :, ::2, ::2] = selected
return result
def _node_prediction(self, index, states, image):
"""Return the unrectified local prediction and its edge cache."""
kind = self._node_kind[index]
if kind[0] == "stem":
pre = image
shortcut = None
elif kind[0] == "first":
pre = states[kind[1]]
shortcut = None
else:
pre = states[kind[1]]
shortcut = self._option_a_shortcut(
states[kind[2]], kind[3], kind[4])
spec = self.layer_specs[index]
convolution = F.conv2d(
pre, self.W[index], stride=spec.stride, padding=spec.padding)
normalized, normalization = self._normalize(
index, convolution, training=True, update_stats=False)
prediction = (
shortcut + spec.branch_scale * normalized
if shortcut is not None else normalized)
return prediction, {
"pre": pre,
"normalization": normalization,
"shortcut": shortcut,
}
def _edge_transpose(self, child, edge_kind, field, states, image):
"""Apply one residual-DAG edge transpose to a child state field."""
kind = self._node_kind[child]
if edge_kind == "shortcut":
if kind[0] != "second":
raise AssertionError("only a second convolution has a shortcut")
return self._option_a_shortcut_transpose(
field, states[kind[2]].shape, kind[4])
_, cache = self._node_prediction(child, states, image)
spec = self.layer_specs[child]
local_field = field * spec.branch_scale
local_field, _, _ = self._normalization_backward(
child, local_field, cache["normalization"])
return torch.nn.grad.conv2d_input(
cache["pre"].shape, self.W[child], local_field,
stride=spec.stride, padding=spec.padding)
def _outgoing_feedback(self, index, deltas, states, image):
result = torch.zeros_like(states[index])
for child, edge_kind in self._outgoing[index]:
result.add_(self._edge_transpose(
child, edge_kind, deltas[child], states, image))
if index == self.n_hidden - 1:
spatial = states[index].shape[2] * states[index].shape[3]
result.add_(
(deltas[-1] @ self.W_out)[:, :, None, None] / spatial)
return result
def infer_dual_states(self, image, one_hot, clean_forward=None):
"""Run the author ``fwK`` DP-transpose state updates on the DAG."""
if clean_forward is None:
clean_forward = self.forward(
image, return_cache=True, training=True, update_stats=False)
plus = [
value.detach().clone() for value in clean_forward["hiddens"]]
minus = [value.detach().clone() for value in plus]
plus.append(clean_forward["logits"].detach().clone())
minus.append(clean_forward["logits"].detach().clone())
fixed_prediction = clean_forward["logits"].detach()
alpha = self.dp_alpha
for _ in range(self.dp_inference_passes):
for index in range(self.n_hidden):
states = [
alpha * positive + (1.0 - alpha) * negative
for positive, negative in zip(plus[:-1], minus[:-1])]
prediction, _ = self._node_prediction(index, states, image)
deltas = [
positive - negative
for positive, negative in zip(plus, minus)]
feedback = self._outgoing_feedback(
index, deltas, states, image)
plus[index] = F.relu(
prediction + (1.0 - alpha) * feedback)
minus[index] = F.relu(prediction - alpha * feedback)
states = [
alpha * positive + (1.0 - alpha) * negative
for positive, negative in zip(plus[:-1], minus[:-1])]
features = states[-1].mean(dim=(2, 3))
prediction = features @ self.W_out.t() + self.b_out
output_field = self.dp_beta * (
torch.softmax(fixed_prediction, dim=1) - one_hot)
plus[-1] = prediction - (1.0 - alpha) * output_field
minus[-1] = prediction + alpha * output_field
return plus, minus
@torch.no_grad()
def dualprop_ascent_directions(self, image, plus, minus):
"""Evaluate the local DP contrastive correlations without autograd."""
alpha = self.dp_alpha
beta = self.dp_beta
states = [
alpha * positive + (1.0 - alpha) * negative
for positive, negative in zip(plus[:-1], minus[:-1])]
deltas = [
(positive - negative) / beta
for positive, negative in zip(plus, minus)]
directions = []
gamma_directions = []
beta_directions = []
batch = image.shape[0]
for index in range(self.n_hidden):
_, cache = self._node_prediction(index, states, image)
spec = self.layer_specs[index]
local_field = deltas[index] * spec.branch_scale
local_field, gamma_direction, beta_direction = (
self._normalization_backward(
index, local_field, cache["normalization"]))
direction = torch.nn.grad.conv2d_weight(
cache["pre"], self.W[index].shape, local_field,
stride=spec.stride, padding=spec.padding)
directions.append(direction / batch)
if gamma_direction is not None:
gamma_directions.append(gamma_direction / batch)
beta_directions.append(beta_direction / batch)
features = states[-1].mean(dim=(2, 3))
output_weight = deltas[-1].t() @ features / batch
output_bias = deltas[-1].mean(dim=0)
return (
directions, gamma_directions, beta_directions,
output_weight, output_bias)
def dualprop_step(self, image, labels, eta, eta_output=None,
momentum=0.0, weight_decay=0.0):
"""One fully local DP-transpose update on a matched ResNet batch."""
one_hot = F.one_hot(labels, self.n_classes).to(image.dtype)
with torch.no_grad():
clean = self.forward(
image, return_cache=True, training=True, update_stats=True)
loss = F.cross_entropy(clean["logits"], labels)
plus, minus = self.infer_dual_states(
image, one_hot, clean_forward=clean)
(directions, gamma_directions, beta_directions,
output_weight, output_bias) = self.dualprop_ascent_directions(
image, plus, minus)
self.apply_ascent(
directions, output_weight, output_bias, eta,
eta_output=eta_output, momentum=momentum,
weight_decay=weight_decay,
gamma_directions=gamma_directions,
beta_directions=beta_directions)
return float(loss)
|
