When Does Depth Matter For In-Context Learning? Adaptive Inference in Deep Transformers
Abstract
Transformers perform computations through many successive attention and feedforward blocks, allowing them to learn complex correlations between a large collection of coupled variables. When does stacking successive attention-feedforward computations over many layers provide a computational advantage over a single transformer block? We address this question by examining in-context learning in generalized linear attention transformers. We first introduce a general theory of distributed inference in such transformers, subject to constraints on communication and depth. We show that such systems can exploit internal representations (`function vectors') to infer a latent context variable at increasingly finer scales over its layers. For an in-context linear regression task, the theory predicts that while one-layer transformers without feedforward blocks are optimal for Gaussian priors over the context variable, multi-layer transformers are superior for non-Gaussian, tree-like priors. Trained linear attention transformers reproduce quantitative predictions from the theory. Using causal key-patching experiments, we verify that function vectors in intermediate layers mediate adaptive routing of information. Our results suggest that depth and feedforward blocks enable transformers to implement adaptive inference, and this is advantageous when the distribution over latent variables has hierarchical structure.
Related
- PaECTER: Patent-level representation learning using citation-informed transformers
- PHACTn enables training-free, context-independent inference of nucleotide variant tolerance across the genome
- LoCoFlow: Local-Context Flow Matching for Single-Cell Dynamics and Fate Inference
- SpaFactor: Lightweight Spatial Context-Aware Gene Program Modeling for Histology-to-Transcriptomics Inference
- flodl-bio: First Rust data, transformers and model layer for entire deep learning on biological sequences
BibTeX
@misc{raj2026when,
title = {{When Does Depth Matter For In-Context Learning? Adaptive Inference in Deep Transformers}},
author = {Ravin Raj and Gautam Reddy},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2606.16694},
eprint = {2606.16694},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2606.16694},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - When Does Depth Matter For In-Context Learning? Adaptive Inference in Deep Transformers
AU - Ravin Raj
AU - Gautam Reddy
PY - 2026
DA - 2026/09/27
PB - arXiv
DO - 10.48550/arxiv.2606.16694
UR - https://arxiv.org/abs/2606.16694
AB - Transformers perform computations through many successive attention and feedforward blocks, allowing them to learn complex correlations between a large collection of coupled variables. When does stacking successive attention-feedforward computations over many layers provide a computational advantage over a single transformer block? We address this question by examining in-context learning in generalized linear attention transformers. We first introduce a general theory of distributed inference in such transformers, subject to constraints on communication and depth. We show that such systems can exploit internal representations (`function vectors') to infer a latent context variable at increasingly finer scales over its layers. For an in-context linear regression task, the theory predicts that while one-layer transformers without feedforward blocks are optimal for Gaussian priors over the context variable, multi-layer transformers are superior for non-Gaussian, tree-like priors. Trained linear attention transformers reproduce quantitative predictions from the theory. Using causal key-patching experiments, we verify that function vectors in intermediate layers mediate adaptive routing of information. Our results suggest that depth and feedforward blocks enable transformers to implement adaptive inference, and this is advantageous when the distribution over latent variables has hierarchical structure.
ER -