The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization

Despite progress across a broad range of applications, Transformers have\nlimited success in systematic generalization. The situation is especially\nfrustrating in the case of algorithmic tasks, where they often fail to find\nintuitive solutions that route relevant information to the right node/operation\nat the right time in the grid represented by Transformer columns. To facilitate\nthe learning of useful control flow, we propose two modifications to the\nTransformer architecture, copy gate and geometric attention. Our novel Neural\nData Router (NDR) achieves 100% length generalization accuracy on the classic\ncompositional table lookup task, as well as near-perfect accuracy on the simple\narithmetic task and a new variant of ListOps testing for generalization across\ncomputational depths. NDR's attention and gating patterns tend to be\ninterpretable as an intuitive form of neural routing. Our code is public.\n