Source-linked AI summary
Recurrent Independent Mechanisms
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, Bernhard Schölkopf
TL;DR
Many dynamical systems contain mostly independent processes, but conventional recurrent models encourage all processes to interact, challenging generalization under changing task distributions. RIMs address this with recurrent modules that have separate default dynamics, selective activation, and sparse attention-based communication. The paper reports specialization and improved generalization across variations, while showing that sparse activation, input attention, and inter-RIM communication are important design components.
Problem
The paper asks whether models can learn independent, sparsely interacting recurrent mechanisms from finite data in changing environments, rather than relying on fully connected dynamics.
Method
RIMs divide a recurrent model into modules with distinct dynamics, selectively activate relevant modules through top-down input attention, and allow sparse communication between modules.
Results
RIMs specialize over temporal patterns and objects and generalize better than recurrent baselines when subsets of patterns or distractors change between training and evaluation.
Takeaways & Limitations
The results support modular recurrent architectures as a way to model independent mechanisms and improve robustness to systematic changes in task factors.
Takeaways & Limitations
The paper assumes high-level variables have structured dependencies rather than being marginally independent, limiting direct applicability to disentangling settings built on the latter assumption.
Abstract
from arXiv · showhide
Learning modular structures which reflect the dynamics of the environment can lead to better generalization and robustness to changes which only affect a few of the underlying causes. We propose Recurrent Independent Mechanisms (RIMs), a new recurrent architecture in which multiple groups of recurrent cells operate with nearly independent transition dynamics, communicate only sparingly through the bottleneck of attention, and are only updated at time steps where they are most relevant. We show that this leads to specialization amongst the RIMs, which in turn allows for dramatically improved generalization on tasks where some factors of variation differ systematically between training and evaluation.
1 INDEPENDENT MECHANISMS
The paper motivates modeling dynamic systems as compositions of mostly independent mechanisms that interact sparsely. Such modular inductive biases may improve learning from finite data and generalization when only some mechanisms change.
- Physical processes often comprise simpler subsystems that evolve independently and interact sparsely.
- Independent mechanisms are central to causal modeling because localized interventions can change one module while others remain robust or invariant.
- Sparse interactions can reduce unnecessary interference by requiring systems to jointly consider only subsystems that meaningfully interact.
- The paper asks whether gradient-based models can learn independent, sparsely interacting recurrent mechanisms from finite data in changing environments.
- For k independent mechanisms, a fully connected recurrent model would need a fraction k−1 of connections set exactly to zero, approaching 100% as k grows.
- The authors argue that modularity and dynamic recombination provide an advantageous inductive bias for sample efficiency and out-of-distribution generalization.
2 RIMS WITH SPARSE INTERACTIONS
RIMs divide a recurrent model into modules with separate default dynamics, selectively activate modules through input attention, and permit sparse communication among them. This design aims to preserve modularity while modeling relevant interactions.
- RIMs divide the model into k recurrent subsystems with distinct learned functions and parameters shared across time.
- Each RIM has independent default dynamics, while the architecture allows occasional interaction with relevant RIMs and selected input elements.
- 2.1 KEY-VALUE ATTENTION TO PROCESS SETS OF NAMED INTERCHANGEABLE VARIABLES: The input is represented as rows of elements, transformed into keys and values, while RIM states are transformed into queries for attention-based processing.
- 2.2 SELECTIVE ACTIVATION OF RIMS AS A FORM OF TOP-DOWN MODULATION: Input attention lets RIMs compete for activation, selecting the top kA of K RIMs with the least attention to a null input.
- 2.2 SELECTIVE ACTIVATION OF RIMS AS A FORM OF TOP-DOWN MODULATION: Queries depend on each RIM’s state, enabling top-down selective attention to input information relevant to that RIM.
- 2.3 INDEPENDENT RIM DYNAMICS: Activated RIMs update through separate GRU or LSTM transition dynamics, whereas inactive RIM states remain unchanged.
- 2.4 COMMUNICATION BETWEEN RIMS: Activated RIMs can read from all other RIMs through attention, allowing contextual information to be shared without updating inactive states.
3 RELATED WORK
Related architectures also use independent memory elements, modular experts, or dormant recurrent state. RIMs combine these ideas with separate recurrent dynamics and attention-controlled input and communication.
- Neural Turing Machines update independent memory cells through attention-based targeted reads and writes, a principle RIMs also use for sparse input impact.
- IndRNN gives each recurrent unit independent dynamics, whereas RIMs use separate recurrent models that can communicate through attention.
- Modular-network approaches route information through gated experts, but commonly activate only one expert per time step; RIMs can activate and connect multiple modules.
- Computation-on-demand architectures keep parts of recurrent state dormant, while RIMs additionally modularize recurrent dynamics and control each module’s inputs with attention.
4 EXPERIMENTS
Across temporal, object-based, and reinforcement-learning settings, RIMs are evaluated for specialization and generalization under changed environments or factors of variation. The experiments report improved robustness over recurrent baselines and identify attention, communication, and sparse activation as important design components.
- Temporal patterns: RIMs specialize on distinct temporal patterns and generalize better when selected patterns change, especially patterns unrelated to the class label.This section evaluates generalization based on temporal patterns before moving to object-based and combined settings.
- Temporal patterns: On the copying task, RIMs retain perfect performance when the dormant phase increases from 50 training steps to 200 testing steps.The task separates sequence-reception and dormant temporal patterns.
- Temporal patterns: Sequential MNIST results show that RIMs are more robust than LSTMs to test resolutions different from the 14x14 training resolution and outperform several strong baselines.The cited results name Transformers, EntNet, RMC, and DNC as baselines that RIMs outperform.
- Objects: In bouncing-balls experiments, RIMs substantially improve over LSTMs when object counts or visual regions differ between training and evaluation.The evaluations include training on four balls and testing with six to eight or four balls, plus tracking through an occluding curtain.
- Objects: For object picking with distractors, RIMs outperform LSTMs, including with novel unseen distractors, where top-down attention suppresses irrelevant representations.With known distractors, RIMs achieve perfect performance while LSTMs struggle; with novel distractors, RIMs strongly outperform LSTMs without achieving perfect performance.
- Discussion and ablations: Ablations show that sparse activation, input attention, and communication between RIMs materially affect performance, while moderate sparsity from 30% to 70% performs optimally on copying.Removing input attention substantially degrades Atari results, and removing communication is reported to harm copying and sequential MNIST performance.
5 CONCLUSION
RIMs model dynamical systems as modules that operate independently by default and interact sparsely, supporting modular computation and potential robustness to changing environments. The conclusion also positions the architecture alongside related modular-memory and attention-based recurrent models.
- Conclusion: RIMs divide a dynamical system into recurrent modules with distinct functions learned automatically from data.
- Conclusion: Most RIMs follow their own default dynamics and share information only rarely through sparse interactions.
- Related work: RIMs relate to separate recurrent models, selective-attention theories, capsules, and graph-based methods through their emphasis on modular or sparse computation.
- Desiderata: Sparse information flow is intended to keep irrelevant or distracting inputs from affecting dormant RIMs.
- Desiderata: Modular parameterization associates most parameters with individual RIMs, allowing specialization on simpler subproblems and preserving functionality when other modules change.
- Related work: RIMs share sparse memory access with Neural Turing Machines, while differing from Relational RNNs by maintaining separate transition dynamics and limiting information flow.
C.1 RIMS IMPLEMENTATION
The implementation combines input attention, selective activation, independent per-RIM recurrence, and inter-RIM communication. The appendix also records hyperparameter, optimization, decoder, and input-processing choices, including several untested or preliminary alternatives.
- Architecture: The RIMs model has three main components: input attention, activated-RIM selection, and communication between RIMs.
- Architecture: Top-k attention selects activated RIMs, while separate LSTMs provide independent dynamics and multihead attention enables communication between modules.
- Hyperparameters: The architecture introduces two hyperparameters beyond LSTM/GRU layers: the total number of RIMs and the number activated per step.
- Limitations: The authors did not systematically optimize the proposed architecture and suggest principled tuning could improve performance.
- Explored changes: An additional output key-value attention layer seemed beneficial empirically but was not included in the paper’s implementation.
- Explored changes: The experiments used a shared decoder for all RIMs, while a more structured decoder was identified as a possible future improvement.
D.1 EFFECT OF VARYING NUMBER AND ACTIVE RIMS ON COPYING TASK
The copying-task experiments vary model capacity and evaluate robustness to longer dormant intervals. RIMs retain generalization beyond the training length, whereas the LSTM baseline fails catastrophically.
- Setup: 600 recurrent units were used for both models: RIMs split them across 6 blocks, while the LSTM used 600 units in total.
- Task: The copying sequences contain 10 random digits, a zero span, and an indicator that starts reproduction of the original sequence.
- Evaluation: Training used zero spans of length 50, while evaluation used length 200 to test out-of-distribution generalization.
- Results: RIMs generalize to sequences longer than those seen during training, whereas the LSTM fails catastrophically.
D.2 EFFECT OF VARYING NUMBER AND ACTIVE RIMS ON ADDING TASK
The adding task tests whether models can select designated numbers from a stream while ignoring irrelevant inputs. The evaluation extends sequence length and includes configurations varying the number of RIMs and active RIMs.
- Task: The task presents real-valued number streams and identifies two randomly varying input streams whose values should be added.
- Evaluation: Testing uses longer input sequences than training to assess whether the model ignores numbers it is not tasked with adding.
- Results: The reported results show better generalization for longer testing sequences and for adding multiple numbers.
D.3 SEQUENTIAL MNIST RESOLUTION TASK
The supplied passages describe sequential MNIST resolution evaluation and bouncing-ball analyses, including RIM specialization, rollout comparisons, and occlusion tracking.
- Training uses 14 x 14 MNIST sequences, with evaluation on higher resolutions of 16 x 16, 19 x 19, and 24 x 24.
- RIM activations initially correlate strongly with the locations of four balls, but this correlation deteriorates later as predictable dynamics require less constant attention.
- RIMs track balls through occlusion without difficulty, while the LSTM baseline also tracks the ball through the invisible curtain.
- Across four experiments, RIMs capture ball trajectories more accurately than two LSTM baselines during rollout, despite hard collisions making modeling difficult.
D.5 VIDEO PREDICTION FROM CROPS
The passages cover video prediction under distribution shifts and a MiniGrid reinforcement-learning setting, with RIMs outperforming non-oracle baselines on OOD one-step prediction.
- Models are trained on 3-ball videos and evaluated on five OOD test sets containing 1 to 6 bouncing balls.
- The evaluation reports balanced accuracy and F1-scores to account for class imbalance.
- On the OOD one-step forward prediction task, the proposed method outperforms all non-oracle baselines while balancing in-distribution and OOD performance.
- The MiniGrid experiments use goal-conditioned A2C across maze multi-room tasks, with LSTM state encoding for partial observability.
D.7.6 ADDITIONAL ABLATION
The additional ablation compares full RIMs against variants lacking input attention or using one attention head, alongside Atari comparisons with PPO and LSTM-PPO.
- The ablation removes input attention while forcing continual communication, or replaces multi-headed key-value attention with one head.
- The one-head and no-input-attention models perform worse than normal RIMs during rollout.
- Using kA = 5 performs slightly worse than kA = 4 but still outperforms PPO across most games.
- Atari comparisons use identical PPO settings while changing only the recurrent architecture to RIMs with kA = 5.
E.1 IMITATION LEARNING: ROBUSTNESS TO NOISE IN STATE DISTRIBUTION
The passages report RIM evaluations across noisy imitation learning, Atari transfer, masking ablations, and machine translation transfer, emphasizing robustness and modular behavior.
- Under joint perturbations unseen during training, RIMs suffer a performance drop but not as severe as the LSTM baseline.
- RIMs transfer positively to 9 of 12 target Atari games, compared with 8 of 12 for progressive networks and 3 of 12 for LSTM.
- Masking experiments find that sparsity is needed to prevent RIMs from co-adapting, while several alternative models fail to remain meaningful after masking.
- In WMT-to-IWSLT translation, multitask training hurts LSTM performance but helps RIMs, while transformer performance remains about the same.