Do Transformers Use their Depth Adaptively? Evidence from a Relational Reasoning Task

TL;DR

This paper examines whether transformer models adaptively utilize their depth for complex relational reasoning tasks, finding evidence of adaptive depth use especially in finetuned models.

Contribution

It provides empirical evidence that transformers, particularly when finetuned, use their depth adaptively based on task difficulty and complexity.

Findings

Larger models need fewer layers for easier tasks.

Models use more layers as chain length increases.

Finetuned models show clearer adaptive depth use.

Abstract

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.