Subjective Depth and Timescale Transformers: Learning Where and When to Compute

TL;DR

This paper tackles the inefficiency of uniform computation in decoder-only Transformers by introducing two surprise-based conditional compute architectures, SDT and STT, which route computation across depth and time using Bayesian surprise signals. SDT alternates Decision and Dynamic layers to gate tokens via a fixed-capacity Top-K router, while STT uses a temporal Transition Network to predict residuals and gate tokens per step, both trained with a differentiable surprise-based objective. Across transfer-learning experiments, the models demonstrate stable training and meaningful compute savings (self-attention up to $62.5\%$ of the dense baseline and KV-cache reductions around $25\%$), with STT often achieving strong task performance despite the overall accuracy gap relative to the dense model due to reduced compute budgets. The results align with predictive-coding theory, revealing a shift from novelty-driven to prediction-driven gating during training and establishing a flexible framework for efficiency in Transformer-like models, while outlining clear directions for scaling, improved routing, and robust evaluation. $\text{This work advances conditional computation by grounding routing in Bayesian surprise and demonstrating practical compute-accuracy trade-offs with SDT and STT.}$

Abstract

The rigid, uniform allocation of computation in standard Transformer (TF) architectures can limit their efficiency and scalability, particularly for large-scale models and long sequences. Addressing this, we introduce Subjective Depth Transformers (SDT) and Subjective Timescale Transformers (STT), two distinct architectures that leverage Bayesian surprise signals to dynamically route computation, learning where and when to compute within decoder-only TFs. SDT augments a decoder-only stack with alternating Decision and Dynamic layers: a Decision layer computes a full block 'posterior' and a lightweight 'prior,' while a Dynamic layer employs fixed-capacity Top-K routing based on Bayesian surprise (Expected and Unexpected Change), maintaining a static compute graph. STT extends this conditional computation to the temporal domain: a transition network predicts residual updates, forming a temporal 'change hypothesis' that informs a router to dynamically execute or bypass TF blocks for each token, managing KV-cache contributions. Both architectures exhibit the predicted shift from novelty to prediction driven gating over training, suggesting alignment with surprise based principles. While operating at reduced capacity, they offer preliminary insights into the compute-accuracy trade-offs of conditional computation. The proposed architectures establish a flexible framework for efficiency, reducing self-attention computation by 75% and KV-cache requirements by 50% within each compute skipping layer, setting a pathway for more efficient models.

Figures (2)

• Figure 1: Architectural comparison and routing criteria visualisation. (A) A standard decoder-only TF processes all tokens through every block. (B) SDT augments the standard stack with alternating Decision and Dynamic layers. At Decision layers, a lightweight PriorFFN generates a prediction of the main block's output, and both are passed to a Router. (C) STT uses a Transition Network to form a temporal prediction for the current token $t$ based on the previous token's state $t-1$. (D) Plane of decision criteria, providing a geometric intuition for the routing logic. From a starting point or static prior ($x_n^t$), a transition network provides a change prior. The distance from each of these priors to the true posterior ($x_n^{t+1}$) is measured as static surprise ($D_{st}$) and change surprise ($D_{ch}$), respectively. A token is processed if the change prior is a better explanation for the posterior (CE) or if the static surprise is unusually large (CU).
• Figure 2: (A) Training loss for the fixed-capacity SDT, STT, and MoD models at $\gamma=0.5$. Curves are smoothed with an exponential moving average over the last 15 steps. (B) Comparison of the predictive loss ($\mathcal{L}_{pred}$) for the STT's TPN and the SDT's PriorFFN with varying intermediate size factors. The temporal prior (STT) is consistently more accurate. (C) The average proportion of tokens selected per layer by the STT (Dynamic Capacity) model during validation. The model learns to reduce its computational capacity in deeper layers.