Looped Transformer: Adaptive Early Exit
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Implement per-sequence early stopping for a looped Transformer using a simple hidden-state change rule. hidden has shape (B, S, D) with positive dimensions. The supplied block maps the full batch to the same shape, operates independently on each sequence, and does not modify its input. Reuse that same callable throughout.
Start every sequence as active. On each round, call block once on the full current batch. For each active sequence, accept its new state, count one step, and stop it when the largest absolute change across all its tokens and features is at most tol. Keep stopped states frozen while other sequences continue. Stop calling the block once every sequence has stopped, or after max_loops rounds. This rule measures state change, not answer correctness.
Topics: transformer, looped-transformer, pytorch, numpy
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Looped Transformer: Adaptive Early Exit in the browser and run it against the tests.