SPB Git

spb/wp31_uqo

Public
5.0 KB · 62 lines latex
1\section{Introduction}\label{sec:intro}23Small language models spend the same computation on every token. A growing body of work relaxes this by4re-applying shared layers a variable number of times---recurrent depth, looped models, mixtures of recursions---and5by learning when to stop \citep{dehghani2019universal,geiping2025recurrentdepth,zhu2025ouro,bae2025mixtureofrecursions,wang2026recurtrace}.6In all of these the extra computation is \emph{global}: every refinement step rewrites the entire hidden state of7the position. This paper asks a narrower question. If the extra computation is instead confined to a small8latent workspace of $M$ slots, and each step may rewrite only one or two of them, can a learned policy that9chooses \emph{where} to write---supervised by executed counterfactual branches that measure what each write would10have done to the next-token loss---improve the quality--compute frontier beyond uniform or random slot selection,11dense updates of every slot, and global latent recurrence at matched cost?1213We call the resulting reference architecture \prisme{} (Predictive Routing with Intervention-Supervised Memory14Editing). A causal decoder is left unchanged; a read-only \emph{source bank} caches its past representations15projected to a small width; at every position a fresh 16-slot workspace is initialised from the current16representation; a shared two-block \emph{sparse editor} rewrites one or two slots per action, chosen among1724 nonterminal actions and \Stop{} by a \emph{value controller}; a \emph{differential readout} adds the change of18an attention read over the workspace to the backbone's representation, so that with no edit the model returns19the backbone output bitwise. The controller is trained on \emph{intervention gains}20$g(s,a;y)=\Loss(D(s),y)-\Loss(D(T_a(s)),y)$ obtained by actually executing the paired edit / no-edit branches at21training time, with logged propensities.2223\paragraph{What this paper is and is not.}24It is an executed, pre-registered research programme (Appendix~\ref{app:prereg}) with a complete25implementation, invariant tests, controlled baselines and an honest accounting of parameters, operations,26latency and memory. It is not a claim of architectural priority: an execution-time literature audit27(Section~\ref{sec:related}) shows that each ingredient---sparse slot writes, loss-gain supervision of28continue/stop decisions, zero-initialised identity insertion---has verified predecessors, and that only their29combination lacks one. Nor is it a claim of superiority: the hypotheses of Section~\ref{sec:hypotheses} are30stated so that the experiments can refute them, and several are refuted.3132\paragraph{Hypotheses.}\label{sec:hypotheses}33H1 (editor usefulness): with backbone and data fixed, sparse workspace edits improve held-out loss and34compositional accuracy over the frozen backbone. H2 (spatial selection): learned selection of one or two slots35beats uniform or random selection with the same editor at fixed cost. H3 (intervention supervision): predicting36paired gains routes better than uncertainty-, entropy-, norm- or fixed-budget rules, controller overhead37included. H4 (adaptivity): \prisme{} moves the quality-versus-cost or quality-versus-latency frontier relative to38the best matched baseline. H5 (identity): at budget zero the output equals the backbone's. H639(generalisation): gains survive unseen templates, deeper compositions and seeds. H7 (scaling, conditional).40The pre-registered null possibilities are equally informative: an oracle may find beneficial edits that no41learned controller can predict; token loss may improve while complete answers do not; a parameter-matched42plain Transformer or a global recurrence may win at equal compute; fixed two-slot editing may match learned43routing; dispatch overhead may dominate the theoretical sparse savings.4445\paragraph{Summary of findings.}46\input{sections/findings}4748\paragraph{Contributions.}49(i) A complete, tested reference implementation of \prisme{} with exact parameter and operation accounting for50103M, 459M and 1.18B configurations, an incremental inference engine with separate backbone and bank caches,51and a 15-invariant test suite (Sections~\ref{sec:architecture}--\ref{sec:compute}).52(ii) A counterfactual labelling and controller-training pipeline with logged propensities, exhaustive53one-step and two-step oracle audits, and a variance decomposition of intervention gains into state and action54components (Sections~\ref{sec:learning}, \ref{sec:results-oracle}).55(iii) Controlled, seeded comparisons against a parameter-matched plain decoder, a dense-slot editor, a56global-recurrence adapter trained with the identical recipe, and five reference routers, with paired57uncertainty (Sections~\ref{sec:design}--\ref{sec:results}).58(iv) Measured latency and memory profiles on Apple Silicon showing where sparse computation does and does not59pay (Section~\ref{sec:results-latency}).60(v) A results-to-claim audit, pre-registration, deviations log and failure log that make every number in this61paper traceable to a run identifier (Section~\ref{sec:repro}).62
Files · d6746e03e0d6828cde427e384b9ef1d2b92c62d6
Menu

Theme