about
Memory-augmented Transformers can implement Linear first-Order Optimization (arxiv.org)
1 point by PaulHoule on Oct 28, 2024 | hide | past | pdf | discuss on HN

In plain words: By giving a Transformer a memory of past gradients, it can learn any rule that linearly combines those gradients, including momentum and other advanced variants, instead of plain gradient descent. Tests showed it picked up these rules and adapted to unfamiliar data at test time.

Abstract · Toward generalizable learning of all (linear) first-order methods via memory augmented Transformers

We show that memory-augmented Transformers can implement the entire class of linear first-order methods (LFOMs), a class that contains gradient descent (GD) and more advanced methods such as conjugate gradient descent (CGD), momentum methods and all other variants that linearly combine past gradients. Building on prior work that studies how Transformers simulate GD, we provide theoretical and empirical evidence that memory-augmented Transformers can learn more advanced algorithms. We then take a first step toward turning the learned algorithms into actually usable methods by developing a mixture-of-experts (MoE) approach for test-time adaptation to out-of-distribution (OOD) samples. Lastly, we show that LFOMs can themselves be treated as learnable algorithms, whose parameters can be learned from data to attain strong performance.

Sanchayan Dutta, Suvrit Sra
arXiv:2410.07263 · cs.LG, math.OC · submitted Oct 8, 2024 · updated Feb 1, 2025
abstract · pdf · html

add comment on HN