Do Transformers Use their Depth Adaptively? Evidence from a Relational Reasoning Task
For researchers studying transformer interpretability and adaptive computation, this work provides evidence that models can adjust depth usage based on task complexity, though the effect is task-specific and not universal.
The paper investigates whether transformers use depth adaptively based on task difficulty in a multi-hop relational reasoning task. They find limited evidence in pretrained models but clearer adaptive depth use in finetuned models, especially under less constrained finetuning.
We investigate whether transformers use their depth adaptively across tasks of increasing difficulty. Using a controlled multi-hop relational reasoning task based on family stories, where difficulty is determined by the number of relationship hops that must be composed, we monitor (i) how predictions evolve across layers via early readouts (the logit lens) and (ii) how task-relevant information is integrated across tokens via causal patching. For pretrained models, we find some limited evidence for adaptive depth use: some larger models need fewer layers to arrive at plausible answers for easier tasks, and models generally use more layers to integrate information across tokens as chain length increases. For models finetuned on the task, we find clearer and more consistent evidence of adaptive depth use, with the effect being stronger for less constrained finetuning regimes that do not preserve general language modeling abilities.