| |
| |
| @@ -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) |
|
|
