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