File size: 10,245 Bytes
4462dc6 | 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 | --- mainline-source/src/delta_ngram/qwen38_writer.py
+++ memory-source/src/delta_ngram/qwen38_writer.py
@@ -21,6 +21,40 @@
@property
def targets(self):
return self.target
+
+
+@dataclass(frozen=True)
+class NativeWriterFeatures:
+ """The original event features, without expanding shared tokens per head.
+
+ ``dense`` is populated only for the legacy single-head execution. The
+ multi-head execution keeps ``evidence`` [B,T,3*input_dim] separate from
+ ``head_features`` [events,head_dim]. ``as_tensor`` is for small diagnostics;
+ training should project these two pieces before gathering event rows.
+ """
+
+ rows: Tensor
+ indices: Tensor
+ head_features: Tensor
+ evidence: Tensor | None
+ dense: Tensor | None
+
+ def as_tensor(self) -> Tensor:
+ if self.dense is not None:
+ return self.dense
+ batch, token, _ = self.indices.unbind(dim=1)
+ return torch.cat((self.evidence[batch, token], self.head_features), dim=-1)
+
+
+@dataclass(frozen=True)
+class WriteEventComponents:
+ """Original scaled targets and pre-sigmoid gate logits for composition."""
+
+ rows: Tensor
+ target: Tensor
+ gate_logits: Tensor
+ features: NativeWriterFeatures
+
class NativeForwardWriter(nn.Module):
"""Map a complete, already-arrived evidence block to D1 write events.
@@ -84,6 +118,18 @@
def events(self, hidden: Tensor, rows: Tensor, mask: Tensor,
state: SparseMemoryState) -> WriteEvents:
"""Expose targets/gates for diagnostics; no write math is repeated here."""
+ components = self.event_components(hidden, rows, mask, state)
+ beta = torch.sigmoid(components.gate_logits).float()
+ return WriteEvents(components.rows, components.target, beta)
+
+ def event_components(self, hidden: Tensor, rows: Tensor, mask: Tensor,
+ state: SparseMemoryState) -> WriteEventComponents:
+ """Keep the original arithmetic while exposing logits and features.
+
+ Freezing the writer's parameters does not detach the direct state
+ inputs to the gate. No inverse sigmoid or alternate projection is
+ used to recover the base logits.
+ """
if hidden.ndim != 3 or hidden.shape[-1] != self.input_dim:
raise ValueError("hidden must be [B, T, input_dim]")
if rows.shape != (*hidden.shape[:2], self.num_heads) or mask.shape != hidden.shape[:2]:
@@ -93,40 +139,17 @@
if not (hidden.device == rows.device == mask.device == state.rows.device
== self.pool_score.weight.device):
raise ValueError("writer, evidence, and state must share a device")
- valid = mask.bool()
- events = valid[:, :, None] & rows.ne(0)
- indices = events.nonzero(as_tuple=False)
- if not indices.numel():
- return WriteEvents(rows.new_empty(0), state.values.new_empty((0, self.width)),
- state.counts.new_empty(0))
-
- hidden = hidden.detach().to(dtype=self.pool_score.weight.dtype)
- scores = self.pool_score(hidden).squeeze(-1)
- scores = scores.masked_fill(~valid, torch.finfo(scores.dtype).min)
- weights = torch.softmax(scores, dim=-1) * valid.to(dtype=scores.dtype)
- weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-12)
- pooled = (weights[:, :, None] * hidden).sum(dim=1)
-
- right_sum = torch.zeros_like(hidden)
- right_count = hidden.new_zeros(hidden.shape[:2])
- for offset in range(1, min(self.right_window, hidden.shape[1] - 1) + 1):
- # Only valid positions to the right, within this arrived block.
- active = valid[:, offset:].to(dtype=hidden.dtype)
- right_sum[:, :-offset] = right_sum[:, :-offset] + hidden[:, offset:] * active[:, :, None]
- right_count[:, :-offset] = right_count[:, :-offset] + active
- right = right_sum / right_count.clamp_min(1.0)[:, :, None]
-
- batch, token, head = indices.unbind(dim=1)
- head_features = self.head_embedding(head)
+ features = self.event_features(hidden, rows, mask)
+ if not features.rows.numel():
+ return WriteEventComponents(
+ features.rows, state.values.new_empty((0, self.width)),
+ state.counts.new_empty(0), features)
+ batch, token, head = features.indices.unbind(dim=1)
+ head_features = features.head_features
if self.heads_per_order == 1:
- # Preserve the executed single-head path, including floating-point order.
- features = torch.cat((hidden[batch, token], right[batch, token], pooled[batch],
- head_features), dim=-1)
- target = self.target_head(features).float()
+ target = self.target_head(features.dense).float()
else:
- # Both first layers are affine in the shared evidence. Project once
- # per token rather than copying 3*input_dim features to every head.
- evidence = torch.cat((hidden, right, pooled[:, None].expand_as(hidden)), dim=-1)
+ evidence = features.evidence
evidence_dim = 3 * self.input_dim
target_first = self.target_head[0]
shared = F.linear(evidence, target_first.weight[:, :evidence_dim],
@@ -136,16 +159,17 @@
self.target_head[1](shared[batch, token] + per_head)).float()
if self.target_scale is not None:
target = target * self.target_scale[head, None].float()
- event_rows = rows[batch, token, head]
+ event_rows = features.rows
state_features = None
if not self.operator.evidence_only_gate:
old_values, old_counts = state.lookup(event_rows)
- state_features = torch.cat((old_values.to(hidden.dtype),
- (target - old_values).to(hidden.dtype),
- old_counts.log1p()[:, None].to(hidden.dtype)), dim=-1)
+ dtype = self.pool_score.weight.dtype
+ state_features = torch.cat((old_values.to(dtype),
+ (target - old_values).to(dtype),
+ old_counts.log1p()[:, None].to(dtype)), dim=-1)
if self.heads_per_order == 1:
- gate_features = features if state_features is None else torch.cat(
- (features, state_features), dim=-1)
+ gate_features = features.dense if state_features is None else torch.cat(
+ (features.dense, state_features), dim=-1)
logits = self.gate_head(gate_features)
else:
gate_first = self.gate_head[0]
@@ -156,5 +180,54 @@
if state_features is not None:
gate = gate + F.linear(state_features, gate_first.weight[:, head_end:])
logits = self.gate_head[2](self.gate_head[1](gate))
- beta = torch.sigmoid(logits.squeeze(-1)).float()
- return WriteEvents(event_rows, target, beta)
+ return WriteEventComponents(event_rows, target, logits.squeeze(-1), features)
+
+ def event_features(self, hidden: Tensor, rows: Tensor, mask: Tensor) -> NativeWriterFeatures:
+ """Build the same detached evidence features used by targets and gates."""
+ if hidden.ndim != 3 or hidden.shape[-1] != self.input_dim:
+ raise ValueError("hidden must be [B, T, input_dim]")
+ if rows.shape != (*hidden.shape[:2], self.num_heads) or mask.shape != hidden.shape[:2]:
+ raise ValueError("rows/mask must have shapes [B, T, K] and [B, T]")
+ if rows.dtype != torch.long:
+ raise ValueError("row dtype must be int64")
+ if not (hidden.device == rows.device == mask.device == self.pool_score.weight.device):
+ raise ValueError("writer and evidence must share a device")
+ valid = mask.bool()
+ events = valid[:, :, None] & rows.ne(0)
+ indices = events.nonzero(as_tuple=False)
+ if not indices.numel():
+ return NativeWriterFeatures(
+ rows.new_empty(0), indices,
+ self.head_embedding.weight.new_empty((0, self.head_embedding.embedding_dim)),
+ None, self.pool_score.weight.new_empty((0, self.target_head[0].in_features)))
+
+ hidden = hidden.detach().to(dtype=self.pool_score.weight.dtype)
+ scores = self.pool_score(hidden).squeeze(-1)
+ scores = scores.masked_fill(~valid, torch.finfo(scores.dtype).min)
+ weights = torch.softmax(scores, dim=-1) * valid.to(dtype=scores.dtype)
+ weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+ pooled = (weights[:, :, None] * hidden).sum(dim=1)
+
+ right_sum = torch.zeros_like(hidden)
+ right_count = hidden.new_zeros(hidden.shape[:2])
+ for offset in range(1, min(self.right_window, hidden.shape[1] - 1) + 1):
+ # Only valid positions to the right, within this arrived block.
+ active = valid[:, offset:].to(dtype=hidden.dtype)
+ right_sum[:, :-offset] = right_sum[:, :-offset] + hidden[:, offset:] * active[:, :, None]
+ right_count[:, :-offset] = right_count[:, :-offset] + active
+ right = right_sum / right_count.clamp_min(1.0)[:, :, None]
+
+ batch, token, head = indices.unbind(dim=1)
+ head_features = self.head_embedding(head)
+ if self.heads_per_order == 1:
+ # Preserve the executed single-head path, including floating-point order.
+ features = torch.cat((hidden[batch, token], right[batch, token], pooled[batch],
+ head_features), dim=-1)
+ evidence = None
+ else:
+ # Both first layers are affine in the shared evidence. Project once
+ # per token rather than copying 3*input_dim features to every head.
+ evidence = torch.cat((hidden, right, pooled[:, None].expand_as(hidden)), dim=-1)
+ features = None
+ event_rows = rows[batch, token, head]
+ return NativeWriterFeatures(event_rows, indices, head_features, evidence, features)
|