Implements the core Neural Memory mechanism from the Titans paper with advanced features including associative scan for efficient updates, first and higher-order momentum, per-parameter learning rate modulation, spectral norming of surprises, and weight residual connections. Highly modular with low coupling - depends only on external assoc-scan and injectable memory model, making it easily reusable in other projects.

Contains multiple self-contained memory model implementations: basic MemoryMLP, GatedResidualMemoryMLP, FactorizedMemoryMLP for efficient factorized weights, MemorySwiGluMLP based on popular SwiGLU architecture, and MemoryAttention treating attention as memory. Each implementation is isolated with minimal dependencies (primarily torch and einops), demonstrating various architectural patterns for neural memory that could be extracted and used independently.

Implements Implicit MLP Attention where an MLP serves as implicit weights for attention computation, replacing traditional key-value matrices. Features talking heads mechanism and rotary embeddings. The module is self-contained with clear separation of concerns, implementing a novel attention mechanism that could be dropped into other transformer architectures with minimal modification.

Implements Nested Attention that computes attention in multiple stages where queries, keys, and values are derived from previous attention outputs. Creates a chain of attention operations that can capture more complex dependencies. Like the implicit attention module, it's highly self-contained with minimal external dependencies, representing a reusable attention pattern.