Hitting-Time Adaptive Transformer Depth

Wrap a stack of optional Transformer blocks with a token-level routing controller. Choose a base layer r and M=2--4 extra blocks. At layer r, retain its attention matrix A^(r); after every optional block, obtain A^(r+m) and update R^(m)=A^(r+m)@R^(m-1). For a language-modeling MVP, let T_i be all non-padding tokens; for retrieval, let T_i be the evidence span or a set of designated memory tokens. Compute q_i^(m)=sum over j in T_i of R^(m)_ij. During training, apply each next block to all tokens but multiply its residual update for token i by a soft gate based on whether q_i^(m) has crossed C. Add the budget penalty shown above to target an active-token fraction b. During inference, replace soft gates with hard decisions: exit token i once q_i^(m)>=C, and cache only active-token states. Exact dense matrix products are acceptable for sequences up to 256; for longer sequences, retain top-k attention entries or propagate only K target probe vectors. Stop gradients through hard routing or use a straight-through estimator. The mathematical quantities are R, q, and the first threshold crossing; C, b, M, and the approximation method must be tuned on validation data. First test a 6-layer decoder Transformer on 512--2048-token synthetic copy and retrieval tasks, comparing full-depth inference with standard token pruning. Record block-token FLOPs, latency, perplexity, exact-match retrieval accuracy, and the correlation between h_i and errors. A useful target is 25--40% fewer block-token evaluations with no more than 1% perplexity degradation.