SPB Git

spb/wp31_uqo

Public
6.1 KB · 98 lines latex
1\section{The \prisme{} architecture}\label{sec:architecture}23\prisme{} (Predictive Routing with Intervention-Supervised Memory Editing) adds to a standard decoder a read-only4\emph{source bank}, an ephemeral $M$-slot \emph{workspace}, a shared \emph{sparse editor}, a \emph{value controller}5and a \emph{differential readout}. Figure~\ref{fig:architecture} gives the data flow; Table~\ref{tab:configs}6the three dimensioned configurations. All notation is collected in Appendix~\ref{app:notation}.78\begin{figure}[t]\centering9\input{figures/architecture.tex}10\caption{\prisme{} data flow for one position $t$. Solid arrows are read paths, the dashed arrow is the only write11path (into 1--2 of the $M$ slots per action). The backbone residual stream $H$, the source bank $(K,V)$ and the12summary $z_t$ are inputs to the editor and are never modified by it; the workspace is discarded after the token.13With zero executed edits the output is bitwise the backbone output (explicit bypass).}14\label{fig:architecture}15\end{figure}1617\subsection{Backbone and source bank}\label{sec:backbone}18The backbone $B_\phi$ is a pre-norm causal decoder with RMSNorm \citep{zhang2019rmsnorm}, rotary position19embeddings \citep{su2021roformer}, grouped-query attention \citep{ainslie2023gqa}, SwiGLU feed-forward blocks20\citep{shazeer2020glu}, tied input/output embeddings and one final normalisation. It returns the final residual21stream $H_t=B_\phi(x_{1:t})\in\R^{t\times d}$, with $h_t=H_t[t]$, and its own head is22$\ell^{\mathrm{bb}}_t = E\,\mathrm{RMSNorm}(h_t)$. The \emph{source bank} is a projection of the history to the23workspace width $w\ll d$:24\begin{equation}25\hat H_t=\mathrm{RMSNorm}_{\mathrm{bank}}(H_t),\qquad K_t=\hat H_t W_K,\qquad V_t=\hat H_t W_V,\qquad26z_t = P\,\hat h_t,\qquad W_K,W_V,P^{\!\top}\in\R^{d\times w}.27\label{eq:bank}28\end{equation}29The bank is distinct from the Transformer KV cache: it stores projected \emph{final} representations, it is30append-only, and a query at position $t$ may read rows $1,\dots,t$ only.3132\subsection{Ephemeral workspace}33For every position a fresh workspace $S^{(0)}\in\R^{M\times w}$ is formed from $M$ trainable address embeddings34$e_j$ and the current summary,35\begin{equation}36S^{(0)}_j = e_j + z_t,\qquad j=0,\dots,M-1 .37\label{eq:init}38\end{equation}39$e_j$ carries no imposed semantic role; it only makes slots distinguishable. No state of $S$ survives to position40$t+1$; history is carried exclusively by the autoregressive prefix and the two caches.4142\subsection{Action catalogue}\label{sec:actions}43The nonterminal actions are $\textsc{single}_j$ ($j<M$, edit one slot) and $\textsc{pair}_{2i,2i+1}$44($i<M/2$, edit two fixed slots), plus \Stop. For $M=16$ this gives $|\Act|=24$ nonterminal actions and 25 in45total; the $\binom{16}{2}=120$ arbitrary pairs are \emph{not} available. We write $I(a)\subset\{0,\dots,M-1\}$ for46the written index set of action $a$ ($|I(a)|=k(a)\in\{1,2\}$, $I(\Stop)=\varnothing$). \Stop{} is absorbing and a47hard cap $R_{\max}=8$ bounds the number of actions per position.4849\subsection{Shared sparse editor}\label{sec:editor}50Given the current workspace $S$, an action $a$ with active set $I=I(a)$, the step index $r$ and the remaining budget51$\rho=R_{\max}-r-1$, the editor forms the active queries52$X = S_{I} + \mathrm{ctx}([r/R_{\max},\;\rho/R_{\max},\;k/2])\in\R^{k\times w}$ and applies $L_w=2$ \emph{shared-weight}53blocks. Each block computes, with pre-normalisation and residual connections,54\begin{align}55X &\leftarrow X + \mathrm{MHA}\big(q=X,\;k=K_t,\;v=V_t;\;\text{mask}_{\le t}\big)56   && \text{(read the causal source bank)}\label{eq:blk1}\\57X &\leftarrow X + \mathrm{MHA}\big(q=X,\;k=\bar S,\;v=\bar S\big)58   && \text{(read all $M$ slots of the pre-write snapshot $\bar S=S$)}\label{eq:blk2}\\59X &\leftarrow X + \mathrm{SwiGLU}(X). && \label{eq:blk3}60\end{align}61The bank attention uses query and output projections only (the bank is already in width $w$ and is cached), so the62bank is never re-projected per action. The slot attention re-projects the $M$ slot keys/values at every block63(reference implementation; an optimised kernel could cache the keys of unchanged slots). The two lanes of a pair64read the \emph{same} snapshot $\bar S$, so their writes are simultaneous by construction.6566\paragraph{Clipped, gated write.}67Let $\tilde S_j$ be the lane output for $j\in I$. The write is68\begin{equation}69\delta_j=\tilde S_j - S_j,\qquad70\delta^{c}_j=\delta_j\min\!\Big(1,\frac{c}{\max(\rms(\delta_j),\epsilon)}\Big),\qquad71\eta_j=\eta_{\max}\,\sigma\!\big(g([S_j;\tilde S_j])\big),\qquad72S'_j = S_j+\eta_j\delta^{c}_j,73\label{eq:write}74\end{equation}75and $S'_m=S_m$ \emph{exactly} for $m\notin I$ (the implementation materialises the new rows through a one-hot76scatter from the snapshot, so the result is independent of lane order; Proposition~\ref{prop:locality}). We use77$\eta_{\max}=0.25$, $c=1$: these are untested research hyper-parameters, not tuned values. The gate $g$ and the78context map are zero-initialised, so $\eta_j=\eta_{\max}/2$ at initialisation.7980\subsection{Differential readout and exact bypass}\label{sec:readout}81A single attention read with a query that depends on $z_t$ only,82$r(S,z_t)=\sum_j \alpha_j S_j$, $\alpha=\mathrm{softmax}_j\big((W_q z_t)\cdot (W_k S_j)/\sqrt{w}\big)$, is83evaluated at the final and at the initial workspace, and their difference is mapped back to the residual width:84\begin{equation}85\Delta_t = s\,U\big(r(S^{(n_t)},z_t)-r(S^{(0)},z_t)\big),\qquad86h^{\mathrm{corr}}_t=h_t+\Delta_t,\qquad87\ell_t = E\,\mathrm{RMSNorm}(h^{\mathrm{corr}}_t).88\label{eq:readout}89\end{equation}90$U\in\R^{d\times w}$ is initialised at zero and $s>0$ is a fixed scalar set once at the start of editor training to91the measured root-mean-square of the backbone residual stream (Section~\ref{sec:learning}); because the final92RMSNorm is scale invariant, $s$ only reparametrises the learning rate of $U$. When $n_t=0$, the term93\eqref{eq:readout} is algebraically zero; the implementation additionally branches out before allocating any slot94when the budget is zero, so the fast path performs no workspace computation at all (Proposition~\ref{prop:identity}).9596\subsection{Configurations and parameter accounting}97\input{tables/configs.tex}98
Files · 20e88d79efe28fb13610284e7be1ef0ef40c265d
Menu

Theme