File size: 6,486 Bytes
e29f88c | 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 | """Explicit text routing and learned interfaces for isolated pointer execution."""
from functools import lru_cache
import re
from typing import Any
import torch
from torch import Tensor, nn
from .number_reader import SharedNumberReader
class PromptRouter:
"""Remove the Steps line without parsing its value or changing public inputs.
Preserve the original answer index for existing diagnostic callers. Padding
inserted before that position is masked and RoPE positions are compacted, so
neither count digits nor their token length reach the executor. Cached CPU
tokenization is bounded; no prompt or activation survives in a checkpoint.
"""
def __init__(self, tokenizer: Any) -> None:
self.tokenizer = tokenizer
self.executor_ids = lru_cache(maxsize=4096)(self._executor_ids)
def _executor_ids(self, ids: tuple[int, ...]) -> tuple[int, ...]:
prompt = self.tokenizer.decode(list(ids), skip_special_tokens=False,
clean_up_tokenization_spaces=False)
if not re.fullmatch(r"Rules: [^\n]+\nStart: [A-Z]\nSteps: [1-9][0-9]*\nAnswer:", prompt):
raise ValueError("Isolated execution requires the raw Rules/Start/Steps/Answer prompt")
routed = re.sub(r"\nSteps: [1-9][0-9]*", "", prompt)
return tuple(self.tokenizer.encode(routed, add_special_tokens=False))
def __call__(self, ids: Tensor, mask: Tensor) -> tuple[Tensor, Tensor, Tensor]:
rows = ids.detach().cpu().tolist()
masks = mask.detach().cpu().tolist()
routed = torch.full_like(ids, self.tokenizer.pad_token_id)
attention = torch.zeros_like(mask)
for row, (tokens, valid) in enumerate(zip(rows, masks, strict=True)):
positions = [i for i, active in enumerate(valid) if active]
executor = self.executor_ids(tuple(tokens[i] for i in positions))
end = positions[-1]
if not executor or len(executor) > end + 1:
raise ValueError("Routed executor prompt does not fit original input positions")
routed[row, :len(executor) - 1] = torch.tensor(executor[:-1], device=ids.device)
routed[row, end] = executor[-1]
attention[row, :len(executor) - 1] = 1
attention[row, end] = 1
positions = (attention.cumsum(-1) - 1).clamp_min(0)
return routed, attention, positions
class ReentryBridge(nn.Module):
"""Learned map from R output back to R input, applied only on re-entry.
Normalize direction, then restore the *initial* working-state RMS scale.
The scale contains rules/start information only and stays differentiable.
Identity initialization of the projection does not claim T>1 equivalence.
"""
def __init__(self, width: int) -> None:
super().__init__()
self.projection = nn.Linear(width, width, bias=False)
nn.init.eye_(self.projection.weight)
self.gain = nn.Parameter(torch.ones(width))
def forward(self, state: Tensor, initial: Tensor) -> Tensor:
rms = state.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
scale = initial.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
normalized = (state.float() / rms * scale).to(state.dtype)
return self.projection(normalized * self.gain)
class RecurrentController(nn.Module):
"""Own learned memory; detached prompt/executor features, no numeric clock."""
kind = 'gru'
def __init__(self, width: int, intermediate: int = 128) -> None:
super().__init__()
self.intermediate = intermediate
self.context = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
self.observation = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
self.cell = nn.GRUCell(intermediate, intermediate)
self.readout = nn.Linear(intermediate, 1)
def initialize(self, prompt_state: Tensor) -> Tensor:
return self.context(prompt_state.detach())
def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
memory = self.cell(self.observation(state.detach()), memory)
return self.readout(memory).squeeze(-1), memory
class AffineSuffixController(RecurrentController):
"""Learned prompt reading and unbounded scalar recurrence; no parsed clock.
Initialization consumes flattened final-eight frozen token embeddings.
The trainable cell implements a*m+b, initialized to identity (b=0), never
a programmed decrement. Numerical labels train the cell through free-running
recurrence; stop BCE trains its readout without distorting counting dynamics.
Executor observations are intentionally ignored for this requested-count
controller: timing is independent of the pointer's identity or correctness.
"""
kind = 'affine_suffix'
suffix_tokens = 8
def __init__(self, width: int) -> None:
nn.Module.__init__(self)
self.width = width
self.intermediate = 1
self.context = nn.Linear(width * self.suffix_tokens, 1)
self.observation = nn.Identity()
self.cell = nn.Linear(1, 1)
self.readout = nn.Linear(1, 1)
nn.init.ones_(self.cell.weight)
nn.init.zeros_(self.cell.bias)
nn.init.zeros_(self.readout.weight)
nn.init.zeros_(self.readout.bias)
def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
self.observation(state.detach()) # Preserve observation hooks for cache/native checks.
memory = self.cell(memory)
return self.readout(memory.detach()).squeeze(-1), memory
class SharedNumberController(AffineSuffixController):
"""Same learned countdown; share the number reader across positions/lengths."""
kind = 'shared_number'
def __init__(self, width: int) -> None:
super().__init__(width)
self.context = SharedNumberReader(width)
def controller_head(width: int, intermediate: int, kind: str = 'gru') -> RecurrentController:
"""Explicit architecture factory; historical checkpoints default to GRU."""
if kind == 'gru':
return RecurrentController(width, intermediate)
if kind == 'affine_suffix' and intermediate == 1:
return AffineSuffixController(width)
if kind == 'shared_number' and intermediate == 1:
return SharedNumberController(width)
raise ValueError('Unknown controller kind or incompatible memory width')
|