Fixed-Point Reasoning: Stable and Adaptive Deep Looped Models

Sajad Movahedi, Shlomo Libo Feigin, Vera Milovanović, Alexander Theus, Thomas Hofmann, Valentina Boeva, T. Konstantin Rusch, Antonio Orvieto

Keywords: Deep Learning, Compositional Reasoning, State Tracking, Looped Transformers, Optimization, Fixed-point Models

TL;DR: We solve the depth-induced signal propagation problem in fixed-point reasoning models.

Abstract: Looped-in-depth architectures provide an inductive bias toward learning step-by-step procedures for tasks that require compositional reasoning. The effective depth reached by looping determines the quality of the solution. Similar to deep architectures, looped architectures are prone to signal propagation issues as the halting decision is postponed. In this paper, we address these signal propagation issues by using pre-norm layers and residual scaling. Furthermore, we propose FPRM: a Fixed-Point Reasoning Model that uses fixed-point convergence as an end-to-end halting mechanism in a looped architecture. We show that fixed-point halting allows FPRM to adapt its compute to the difficulty of the task. FPRM proves effective on common reasoning benchmarks, namely Sudoku, Maze, and state tracking.

Submission Number: 118