File size: 5,176 Bytes
7c5e40e | 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 | """Deterministic evaluation-time interventions for predicate memory.
These interventions are not training features. They are causal probes for asking
whether the learned predicate-memory path is example-specific, span-anchored, and
predicate-gated, rather than just extra residual capacity.
"""
from __future__ import annotations
from typing import Literal
import torch
PredicateMemoryIntervention = Literal[
"none",
"typed_graph_memory",
"zero_memory",
"memory_zero",
"uniform_gate",
"untyped_same_topology",
"span_roll",
"span_corrupted",
"span_reverse",
"batch_shuffle",
"batch_shuffled_graph",
"random_same_degree",
]
PREDICATE_MEMORY_INTERVENTIONS: tuple[str, ...] = (
"none",
"typed_graph_memory",
"zero_memory",
"memory_zero",
"uniform_gate",
"untyped_same_topology",
"span_roll",
"span_corrupted",
"span_reverse",
"batch_shuffle",
"batch_shuffled_graph",
"random_same_degree",
)
def _valid_lengths(attention_mask: torch.Tensor | None, batch: int, seq: int) -> list[int]:
if attention_mask is None:
return [seq] * batch
return [int(x) for x in attention_mask.to(torch.long).sum(dim=1).tolist()]
def _roll_valid(x: torch.Tensor, attention_mask: torch.Tensor | None, *, shifts: int = 1) -> torch.Tensor:
out = x.clone()
b, s = x.shape[:2]
for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
if n > 1:
out[i, :n] = torch.roll(x[i, :n], shifts=shifts, dims=0)
return out
def _reverse_valid(x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
out = x.clone()
b, s = x.shape[:2]
for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
if n > 1:
out[i, :n] = torch.flip(x[i, :n], dims=(0,))
return out
def _random_permute_valid(x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
"""Deterministically permute valid memory slots while preserving the budget."""
out = x.clone()
b, s = x.shape[:2]
for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
if n > 1:
generator = torch.Generator(device="cpu")
generator.manual_seed(1729 + i * 1_000_003 + n * 97)
perm = torch.randperm(n, generator=generator).to(x.device)
out[i, :n] = x[i, perm]
return out
def apply_predicate_memory_intervention(
*,
candidate_key: torch.Tensor,
candidate_value: torch.Tensor,
predicate_gate: torch.Tensor,
attention_mask: torch.Tensor | None,
intervention: str = "none",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return intervened ``(key, value, gate)`` tensors.
``zero_memory`` makes the predicate-memory update exactly zero by zeroing the
values read by attention. The other interventions preserve tensor shapes and
most marginal statistics while corrupting a specific part of the learned
predicate-memory computation.
"""
if intervention in {"none", "typed_graph_memory"}:
return candidate_key, candidate_value, predicate_gate
if intervention not in PREDICATE_MEMORY_INTERVENTIONS:
raise ValueError(f"unknown predicate memory intervention {intervention!r}")
if intervention in {"zero_memory", "memory_zero"}:
return candidate_key, torch.zeros_like(candidate_value), predicate_gate
if intervention in {"uniform_gate", "untyped_same_topology"}:
gate = torch.ones_like(predicate_gate)
if attention_mask is not None:
gate = gate * attention_mask.to(gate.dtype).unsqueeze(-1)
return candidate_key, candidate_value, gate
if intervention in {"span_roll", "span_corrupted"}:
return (
_roll_valid(candidate_key, attention_mask),
_roll_valid(candidate_value, attention_mask),
_roll_valid(predicate_gate, attention_mask),
)
if intervention == "span_reverse":
return (
_reverse_valid(candidate_key, attention_mask),
_reverse_valid(candidate_value, attention_mask),
_reverse_valid(predicate_gate, attention_mask),
)
if intervention in {"batch_shuffle", "batch_shuffled_graph"}:
if candidate_key.shape[0] > 1:
return (
torch.roll(candidate_key, shifts=1, dims=0),
torch.roll(candidate_value, shifts=1, dims=0),
torch.roll(predicate_gate, shifts=1, dims=0),
)
return (
_roll_valid(candidate_key, attention_mask),
_roll_valid(candidate_value, attention_mask),
_roll_valid(predicate_gate, attention_mask),
)
if intervention == "random_same_degree":
return (
_random_permute_valid(candidate_key, attention_mask),
_random_permute_valid(candidate_value, attention_mask),
_random_permute_valid(predicate_gate, attention_mask),
)
raise AssertionError(f"unhandled intervention {intervention!r}")
__all__ = [
"PREDICATE_MEMORY_INTERVENTIONS",
"PredicateMemoryIntervention",
"apply_predicate_memory_intervention",
]
|