Object Representations as Fixed Points
Training Iterative Inference Algorithms with Implicit Differentiation
We show that conceptualizing object representations as stable points of a fixed-point procedure enables us to take advantage of implicit differentiation techniques to stabilize the training and improve the learning performance of object-centric models, specifically slot attention.
Our proposed method is simple to implement, requiring adding only one line of code to the original slot attention implementation and also enjoys lower space and time complexity in the backward pass.
Training with implicit differentiation (in blue) stabilizes the Jacobian norm of the slot attention cell and prevents gradients from exploding.
The vanilla version sometimes misses objects...
... changes their size...
... or changes their color.
The implicit version generally matches the ground truth much more closely. In terms of mean squared error, the implicit version has almost a 7x improvement over its vanilla counterpart.
