1"""Differential readout: delta_out = U (r(S_final, z) - r(S0, z)); U initialised at zero."""2from __future__ import annotations34import math56import torch7import torch.nn as nn89from .config import PrismeConfig101112class DifferentialReadout(nn.Module):13 def __init__(self, cfg: PrismeConfig):14 super().__init__()15 self.cfg = cfg16 self.q_proj = nn.Linear(cfg.w, cfg.w, bias=False)17 self.k_proj = nn.Linear(cfg.w, cfg.w, bias=False)18 self.U = nn.Linear(cfg.w, cfg.d_model, bias=False)19 if cfg.readout_zero_init:20 nn.init.zeros_(self.U.weight)21 else:22 nn.init.normal_(self.U.weight, std=cfg.readout_init_std)2324 def read(self, S: torch.Tensor, z: torch.Tensor) -> torch.Tensor:25 """Attention read over slots with a query that depends only on z. S: [...,M,w]; z: [...,w] -> [...,w]."""26 q = self.q_proj(z).unsqueeze(-2) # [...,1,w]27 k = self.k_proj(S) # [...,M,w]28 att = torch.softmax((q * k).sum(-1) / math.sqrt(self.cfg.w), dim=-1) # [...,M]29 return (att.unsqueeze(-1) * S).sum(-2)3031 def forward(self, S_final: torch.Tensor, S0: torch.Tensor, z: torch.Tensor) -> torch.Tensor:32 return self.U(self.read(S_final, z) - self.read(S0, z)) * self.cfg.readout_scale33