Looped Transformers as Programmable Computers

ICML 2023 · Angeliki Giannou et al.

核心思想：让一个参数固定的浅层 Transformer 反复执行，通过循环次数获得计算深度，从而把 Transformer 变成一个可编程计算机

Image

Image

Image

1. Motivation：为什么需要 Looped Transformer？

传统 Transformer：

Input
 ↓
Layer 1
 ↓
Layer 2
 ↓
...
 ↓
Layer N
 ↓
Output

如果想执行一个需要 N 次迭代 的算法，通常需要增加 Transformer 的层数。

问题是：

模型深度 ≈ 计算步数

因此模型参数量和计算能力会同时增长。

论文希望解决：

能不能只用一个固定的小 Transformer，通过不断重复执行它，完成很长的计算？

于是：

        ┌─────────────┐
        │ Transformer │
        └──────┬──────┘
               ↓
        ┌─────────────┐
        │ Transformer │
        └──────┬──────┘
               ↓
        ┌─────────────┐
        │ Transformer │
        └──────┬──────┘
               ↓
              ...

核心变化：

Parameter depth ≠ Computation depth

参数只需要一套，但是计算可以执行很多步。