Invertible Memory Flow Networks

TL;DR

Invertible Memory Flow Networks (IMFN) address long-context memory by factorizing compression into a binary-tree of locally invertible 2→1 merges, enabling $O(log N)$-depth memory with an explicit inverse path. The approach allows training of a high-capacity teacher and distillation into a constant-time online student, achieving sublinear error growth on long MNIST sequences and high-dimensional video data like UCF-101. Empirical results show IMFN outperforms larger end-to-end baselines (Transformer and Mamba) on reconstruction fidelity while demanding fewer parameters, supported by a Merkle-style update mechanism for efficient trajectory updates. The work highlights the practical viability of invertible, structured memory for long-horizon tasks and outlines future extensions to real-time video, audio, and agentic planning contexts with bounded memory.

Abstract

Long sequence neural memory remains a challenging problem. RNNs and their variants suffer from vanishing gradients, and Transformers suffer from quadratic scaling. Furthermore, compressing long sequences into a finite fixed representation remains an intractable problem due to the difficult optimization landscape. Invertible Memory Flow Networks (IMFN) make long sequence compression tractable through factorization: instead of learning end-to-end compression, we decompose the problem into pairwise merges using a binary tree of "sweeper" modules. Rather than learning to compress long sequences, each sweeper learns a much simpler 2-to-1 compression task, achieving O(log N) depth with sublinear error accumulation in sequence length. For online inference, we distilled into a constant-cost recurrent student achieving O(1) sequential steps. Empirical results validate IMFN on long MNIST sequences and UCF-101 videos, demonstrating compression of high-dimensional data over long sequences.

Figures (6)

• Figure 1: Overview of Invertible Memory Flow Networks (IMFN). The teacher builds a binary-tree memory by repeatedly merging adjacent states with learned "sweeper" modules ($2\!\rightarrow\!1$), yielding a single root memory in $O(\log T)$ depth. An inverse pathway applies the corresponding learned decoders ($1\!\rightarrow\!2$) to reconstruct leaf states (and inputs), providing local reconstruction losses that make each merge approximately invertible. For online use, the tree computation is distilled into a recurrent student that updates a fixed-size memory with constant per-step cost. Due to Merkle-style distillation, trajectory simulation using the teacher becomes much more efficient.
• Figure 2: IMFN teacher reconstruction error (MNIST). Average MSE vs. memory dimension $d$, averaged over 5 seeds.
• Figure 3: IMFN teacher reconstruction error scaling with horizon. Average MSE vs. sequence length $T$, averaged over 5 seeds.
• Figure 4: Student online prefix reconstruction (MNIST, $d=1024$). After $t$ online updates, we decode the first $t$ frames and report MSE averaged over the prefix. Curves correspond to different horizons $T$ (separate student per $T$).
• Figure 5: End-of-sequence reconstruction scaling with horizon (MNIST, $d=1024$). Pixel MSE of teacher and student reconstructions vs. the original sequence at the final step. Error bars show variance across seeds.