Implicit Bias of Optimizers Guides Length Generalization in Forgetting Attention

Autoregressive language models such as transformers process variable-length inputs, yet whether a model trained on bounded-length inputs stays accurate on much longer ones (or length generalize), is known to depend heavily on design choices. We approach it from the optimization perspective, through a task in which a sequence of key-value pairs may reassign a key and a query asks for its key's latest value. Repeated keys force the model to discount older records by their distance from the query, yet a recency bias alone cannot tell a stale record from a correct value that merely sits far back behind newer records of other keys. We analyze this tension in a two-layer attention model with fixed attention masks and value projections, which learns key-value binding, query-key matching, a forget gate, and a linear decoder from binary answer labels. Our main result separates gradient descent (GD) from Adam on this model under suitable per-parameter-group learning rates. From a shared orthogonal initialization of the query and key embeddings with zero initial content logits, GD with a single $1/t$ schedule succeeds whenever the target sits at most two records further back than any training target, and at each greater distance it fails on a fixed sequence. Adam, with its stabilizer $ϵ$ set to zero on the scalar parameters and positive on the embedding matrices, instead guarantees a retrieval distance proportional to its cumulative learning rate, which diverges. Both guarantees hold for an arbitrary prefix of earlier records. Adam's guarantee extends to independent Gaussian initialization. SGD inherits GD's cutoff with increasing batch sizes and exact averages over all key pairs. Controlled experiments compare finite retrieval ranges and isolate the role of the scalar stabilizer $ϵ$.