Initium is an experimental research framework for combining a Differentiable Neural Computer (DNC)/augmented memory with interchangeable sequence controllers and task/data pipelines. Its central idea is to keep an explicit, addressable external memory while exploring controllers such as LSTM, Mamba-family state-space models, and closed-form continuous-time (CfC) networks. The code also includes stochastic memory-write regularization, sparse expert layers, split-graph models, memory-structure experiments, and checkpoint-driven inference.
Research status: Alpha and actively changing. Results below are experiment-specific observations, not guarantees that a configuration will reproduce on another machine, code revision, or dataset. This README is deliberately maintained as a working guide: update it when the code, default configuration, or evidence changes.
The project studies how a neural controller can interact with a differentiable external memory on sequential tasks, especially graph traversal. It asks questions at several levels:
- Memory and controller: Can a controller process a sequence and use DNC read/write operations to store and retrieve information needed later?
- Generalization: How does performance change between the training curriculum and a held-out or out-of-distribution (OOD) evaluation distribution?
- Write regularization: Does placing a probabilistic prior on each memory write improve generalization, and what are the costs or seed-to-seed variations?
- Compute structure: Can the controller's sequence computation be parallelized or sparsified while preserving the sequential memory interface?
- Task breadth and deployment: Can trained checkpoints be evaluated through a common inference interface across graph, text, audio, video, and multimodal tasks?
These are research questions, not all established project claims. In particular, supported dataset interfaces do not imply that every modality has been trained and validated to the same standard as graph traversal.
You can either fork/clone the repository and customize everything you need or you can use pip/pdm/uv for ready to use installation.
Guide for pip usage:
To install the whole framework: pip install initium-deeplearning
For ANY mamba related operation you need: python -m initium.setup_wheels
Optionally you can try: pip install "initium-deeplearning[text]" pip install "initium-deeplearning[video]" pip install "initium-deeplearning[real-data]"
To see all CLI optionality: initium-infer --help python initium.core_training --help
Example usage:
python -m initium.core_training 0.0 --controller lstm python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 ...
For any inference you need a checkpoint first, the built in inferencing tool will analyze the checkpoint and if any of the parameters are not supported you will NOT be able to launch the checkpoint (in case you made some custom changes to the training code).
Configuration and random seed
│
▼
Dataset registry ──► task episodes / curriculum ──► batch and sequence tensors
│
▼
Controller + DNC external memory
├─ controller processes sequence
├─ write head updates memory
└─ read heads return memory content
│
▼
task loss + optional KL / MoE auxiliary losses
│
▼
optimizer, evaluation, logs, checkpoints, curriculum
│
▼
checkpoint loader ──► task inference ──► metrics and run logs
The training core delegates task construction and task-specific behavior to a dataset interface. The model combines a selected controller with DNC memory operations. During training, the loop computes task and optional auxiliary objectives, updates model parameters, periodically evaluates ID and OOD performance, advances the curriculum when its criteria are met, writes diagnostics, and saves resumable checkpoints. Inference reconstructs a model from a checkpoint, checks that the requested task and dimensions are supported, creates the task runner, and reports episode metrics. Optional inference caches can reuse eligible results or shared prefixes.
The DNC supplies a memory matrix and differentiable read/write interface. The memory can retain information beyond a controller's immediate hidden state. Content-based addressing, read heads, and temporal-link structures are part of the underlying DNC implementation. The temporal link matrix has quadratic cost in the number of memory cells, which motivates the project's link-matrix ablations and memory-size investigations.
The training core exposes controller choices through its command-line interface and configuration:
- LSTM: baseline recurrent controller.
- Mamba: selective state-space controller.
- Mamba-2 and Mamba-3: additional state-space variants; current experiment notes report that these configurations underperformed in earlier tests and need further investigation.
- CfC: closed-form continuous-time controller.
- Hybrid chains: compositions of supported controller stages, such as a Mamba/CfC sequence.
- Split-graph: a parallel sequence backbone plus a separate combiner for the sequential memory-addressing work. It is an experimental way to expose more of the controller computation to parallel execution while retaining the memory interface and currently performs the best out of all options.
Some options require optional or hardware-sensitive dependencies, particularly mamba-ssm and causal-conv1d. A controller being selectable in code should not be read as evidence that it is equally mature or portable on every platform.
The write-head extension models a write vector as a diagonal Gaussian. It can sample a write using the reparameterization trick and add a KL penalty to a prior distribution. The implementation supports a fixed standard-normal prior and a learned prior updated from recent write statistics on a periodic, detached snapshot schedule. The loss can include KL annealing and free bits; the learned-prior path logs its snapshot state separately.
The intended locality distinction is important: the per-write KL is computed from the current write distribution and a prior snapshot. In the learned-prior variant, the prior itself is fitted from recent writes periodically, outside the per-step gradient path. That makes the overall mechanism periodically data-conditioned even though the per-write calculation uses the current write and frozen prior snapshot.
Also KL regularization has specific, uninvestigated behaviour: it tends to stabilize grad_norms allowing the model to not collapse (detected once and needs further investigations) - it was mildly visible in graph-traversal experiments but its full strength occured in multimodal tasks where a model managed to collapse in lesson 14 with beta=0.0, meanwhile with beta=0.001(KL term active) the same model kept going with no issues.
- Mixture of Experts (MoE): sparse expert feed-forward capacity with routing and load-balancing diagnostics. The current project notes caution that this overhead may be unnecessary or harmful at smaller model sizes; its motivation is future scaling, not a demonstrated benefit for the current compact experiments.
- Dynamic memory size (Dynamic-N): usage-triggered growth of the memory cell array within configured bounds.
- Link-matrix options: ablations or sparsification of temporal-link computation.
- Dynamic beta: an optional bounded controller for the KL weight, distinct from a fixed-beta sweep.
- Multimodal and real-data loaders: interfaces for synthetic and selected real text, audio, video, graph, and multimodal inputs. Their practical coverage and validation are still evolving.
| Path | Role |
|---|---|
src/initium/core_training.py |
Dataset-agnostic training loop, CLI, evaluation, logging, checkpointing, and orchestration of model options. |
src/initium/config/ |
Training, regularization, controller, and architecture defaults. |
src/initium/data/ |
Dataset registry, task episodes, curriculum/data handling, codecs, and real-data utilities. |
src/initium/mamba_controller/ |
DNC controller integration, Mamba variants, hybrid controllers, and split-graph implementation. |
src/initium/LNN_controller/ |
CfC and hybrid continuous-time controller implementations. |
src/initium/memory_manipulation/ |
Stochastic write heads, prior snapshots, memory resizing, link-matrix ablations, and compatibility helpers. |
src/initium/MoE/ |
Sparse expert layer and its local usage notes. |
src/initium/inference/ |
Checkpoint loading, task registry, inference engine, metrics, logging, and cache support. |
setup_wheels.py |
Environment/dependency setup entry point for specialized wheels. |
pyproject.toml |
Package metadata, base and optional dependencies, and development tooling groups. |
The PyPI distribution is named initium-deeplearning; its Python import package is initium. The project targets Python 3.10 or later. Core dependencies are declared in pyproject.toml; accelerated state-space controllers may need platform-specific PyTorch/CUDA wheels and additional packages. Install the project and optional task dependencies in an environment appropriate to your hardware. For specialized dependency setup, inspect the repository's setup_wheels.py and its help output before running it; GPU package compatibility depends on the installed PyTorch, CUDA, and Python versions.
Typical editable installation for development:
python -m pip install -e .Optional dependencies are grouped by modality, for example text, video, and real-data in the project metadata. Audio/video processing may also rely on system tools such as FFmpeg. Mamba-family components can require packages beyond the base install. The default training path is intended to be usable without those optional controller dependencies.
The source tree is packaged as initium; internal modules use the initium.* namespace.
The main entry point is src/initium/core_training.py. From the repository root, inspect the available arguments with:
python -m initium.core_training --helpThe script accepts a positional KL weight (beta); omitting it runs the configured beta sweep. It also accepts a random seed and controller/task/architecture options. A representative graph-task run with the Mamba controller and the dynamic-memory/split-graph experiments enabled is:
python -m initium.core_training 0.0 --seed 0 --controller mamba --dynamic-n-mode --split-graphThis is an example of the CLI shape, not a universal best configuration. Check src/initium/config/train_config.py and src/initium/config/controller_config.py for the defaults in the checkout you're running. In particular, training duration, model size, beta sweep, dataset, evaluation cadence, and memory options may have changed since older experiment records were produced.
Training runs typically produce:
- periodic and final checkpoints, including model, optimizer, random-generator, curriculum, and architecture/prior state needed to resume or evaluate;
- step-level or periodic logs for task loss, KL, accuracy, OOD evaluation, and training health;
- additional CSV diagnostics for learned-prior snapshots and optional architecture features;
- a run identifier that encodes relevant settings to reduce accidental collisions.
Keep seeds, source revision, full command, configuration, device/dependency versions, and evaluation protocol with any result you intend to compare. A single final metric is not enough to characterize these experiments: curriculum position, checkpoint step, sampled versus deterministic writes, and evaluation sample count can change the interpretation.
The checkpoint-based inference CLI lives at src/initium/inference/run_inference.py. Its help output is the current reference for arguments:
python -m initium.inference.run_inference --helpThe command accepts a checkpoint, task type, optional data link, episode count, seed, device, and task-specific settings. It can inspect checkpoint capabilities, compare reset and persistent experience modes, report memory-ablation behavior, and optionally configure inference caches. A basic invocation has this form:
python -m initium.inference.run_inference path/to/checkpoint.pt --dataset-type graphThe loader checks checkpoint/task compatibility and refuses unsupported task types or dimensions. Graph, text, audio, and video task runners are registered; multimodal execution is also represented in the code but its input dimensions depend on the selected modality combination. Verify the checkpoint's recorded dataset configuration before comparing results across task types.
The supplied experiment records describe a staged investigation of a KL penalty on DNC memory writes. They are summarized here as results rather than by the names of the individual working files, so this overview can remain useful if those files are renamed or reorganized. -> If anyone wishes to see the original files, you can create an issue and will be contacted afterwards.
The first study compared a no-KL anchor with a small fixed-prior KL weight over multiple random seeds with --controller lstm. It found meaningful heterogeneity: several seeds showed non-collapsed OOD gains, some showed little benefit, and two seeds showed genuine reversals. The strongest observed seed improved OOD triple accuracy by roughly 11–12 percentage points over its own anchor. The study therefore supports a conditional positive result for the tested configuration, not a claim that the effect is universal or statistically settled. The measured OOD perfect-traversal fraction also did not track triple accuracy uniformly, so the metric definition matters.
The follow-up replaced the fixed prior with a prior periodically estimated from recent write statistics. Four deliberately selected seeds were tested: two previous reversal cases and two stronger fixed-prior cases. Outcomes were mixed: the two former reversals improved, one strong case was nearly flat, and the strongest fixed-prior result regressed by about five to six OOD points. Across the four seeds, the mean OOD change was small and not statistically distinguishable from zero; ID accuracy was lower by about 1.1–2.7 points in all four comparisons. Diagnostics did not indicate KL collapse. A regression-to-the-mean explanation is plausible (the reported correlation was large at n=4), but the sample is too small to distinguish that from a true repair of the earlier reversal cases. The structural audit found the learned prior's snapshot update was detached and periodic, with no second online posterior; that is a design characterization, not proof of a generalization benefit.
A separate study compared beta 0 and beta 0.001 in a Mamba, dynamic-memory, split-graph configuration across five seeds. The logs showed no statistically clear OOD or ID-to-OOD offset improvement from the KL term; beta 0 had a small nominal OOD edge in four of five seeds. The beta 0.001 runs consistently left the first curriculum lesson later. The comparison is provisional: only one seed reached the final curriculum lesson, four runs were unfinished, and there was no wide-evaluation sweep for that batch. This configuration should not be described as a completed replication of the earlier study because its controller, memory size, training budget, and evaluation coverage differ.
- Results above apply to the exact runs and evaluation procedures recorded; they do not establish broad performance across modalities, hardware, or model scales.
- Phase 1's seed variation and Phase 2's four-seed selection limit aggregate claims. Report per-seed values and uncertainty alongside averages.
- The learned prior's apparent variance-dampening/mean-reversion effect needs more seeds and controls to separate it from failure-mode-specific repair.
- The alternate Mamba batch needs completion of unfinished seeds and late-curriculum evaluation before a gate-level conclusion.
- Mamba-2/Mamba-3 performance, MoE's useful scaling regime, memory-link trade-offs, dynamic memory growth, and real-data/multimodal behavior remain active investigation areas.
- Experiment status can change. Update this section when new runs are verified; retain the method, sample size, seed set, incompleteness, and uncertainty so summaries do not outrun evidence.
The project draws on several complementary lines of work. These references motivate components and questions; their inclusion does not imply that Initium reproduces each paper's results.
- Graves et al., “Hybrid computing using a neural network with dynamic external memory.” DNC memory, differentiable addressing, and the graph-traversal motivation.
- Hasani et al., “Closed-form Continuous-time Neural Networks.” Continuous-time controller design.
- Gu and Dao, “Mamba-3: Improved Sequence Modeling using State Space Principles.” State-space sequence-modeling direction explored by controller variants.
- Dao and Gu, “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.” Structured state-space/sequence computation background.
- “Gated Delta Networks: Improving Mamba 2 with Delta Rule.” Gated delta-rule sequence processing.
- “Parallelizing Linear Transformers with the Delta Rule over Sequence Length.” Parallelization ideas related to sequence-length computation.
- Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” Sparse expert routing foundations.
- Fedus et al., “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” Simplified sparse expert scaling.
- “Multi-Rate Mixture of Experts for Accelerating Liquid Neural Network Training.” Motivation for combining expert computation with liquid-network training.
- “MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts.” MoE and selective state-space combination.
- Higgins et al., “beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework.” Gaussian KL regularization and reparameterized latent sampling adapted here to memory writes.
- Graves et al., “Hybrid computing using a neural network with dynamic external memory.” Baseline memory architecture and task family used by the Q21 study.
- Additional work on local credit assignment, priors, and continual-learning regularization informed the experiment design and interpretation. Consult the experiment notes for the specific comparisons and caveats.
Because the repository is new and research directions are still changing:
- Prefer describing mechanisms and measured outcomes over calling a configuration “best” without a scoped comparison.
- Distinguish implemented, exercised, and validated features. A registry entry is implementation coverage, not empirical validation.
- Preserve negative and mixed results; record the seed, comparison anchor, metric, evaluation window, and incomplete runs.
- Keep CLI help, configuration defaults, this README, and experiment records aligned as the project evolves.
- Treat research papers as motivation and technical background, not as evidence that this implementation matches their methods or results unless a reproduction has been explicitly checked.
Initium is distributed under the GNU General Public License v3.0. See LICENSE.
Package versions come from __version__ in src/__init__.py. After merging a version update into main, create a GitHub Release whose tag matches that version (for example, v0.1.0). The Publish to PyPI workflow builds and validates the wheel and source distribution, then publishes them automatically. It refuses releases whose tag points to a commit outside main or whose tag does not match the package version.
The PyPI distribution name is initium-deeplearning; the source directory and import namespace remain initium. The configured PyPI Trusted Publisher should target project initium-deeplearning with owner samueleternity, repository initium, workflow publish-pypi.yml, and GitHub environment pypi. Publishing uses short-lived OIDC credentials, so no PyPI API token is stored in GitHub.