KV-store memory + new task families

Add sequence-parallel KV and delta memories (#42)

* Add sequence-parallel KV and delta memories

Add configurable causal KV and gated delta-rule memory modules for split-graph controllers, and wire them through training, checkpoint loading, and stochastic write heads. Cover causal reads, stateful execution, gradients, and backbone-compatible initialization with tests.

* Add cross-memory usage diagnostics

Extend memory dependency checks to DNC, KV, and Delta memories, including shuffled-read comparisons. Log memory usage profiles and per-cell or per-head diagnostics during evaluation.

* Add inter-layer sequence memory ablation

Add an opt-in KV/delta memory path between Mamba-1 backbone blocks, with configuration validation and inference support. Add paired graph-traversal diagnostics for memory shuffling, read lag, and layer placement on fixed ID/OOD episodes.

* Enable inter-layer memory with workspace

Allow KV/delta inter-layer memory reads alongside workspace mode. Apply the memory read after the first backbone block and its MoE residual, before the next block, and include workspace runs in memory-path probes.

* Add repeated inter-layer memory reads

Support KV/delta memory reads after every non-final Mamba-1 block, including workspace/MoE paths, with paired evaluation and checkpoint configuration. Keep workspace specialist outputs sparse by default and batch independent projections across timesteps.

* Add causal input prediction and compiled workspace

Generalize prediction objectives and error signals to non-categorical inputs, add an optional compiled workspace recurrence, and extend stochastic write heads to KV values with bounded prior statistics. Also preserve checkpoint compatibility and shorten oversized run IDs for filesystem safety.

* Add workspace capability benchmark

Introduce five synthetic task tracks with oracle-backed generation, curricula, and training/evaluation metrics. Add workspace readback bypass and shuffled-readback comparisons, plus CLI and dataset-registry integration.

* Add semantic workspace binding diagnostics

Add entity-keyed workspace writes, semantic addressing, permutation-state supervision, and readback diagnostics. Separate permutation label and sequence-length generalization evaluations, and add configurable scaling for inter-layer memory reads.

* Add transactional permutation workspace mode

Add atomic paired slot swaps and workspace-based query decoding for permutation episodes, with commit and state-accuracy metrics. Bound semantic workspace content and readback, and document the new mode.

* Canonicalize projected broadcast tensors

Make projected write values and read queries contiguous before recurrent processing so per-step batch strides remain stable for the compiler.

* Add semantic workspace writes and KV bridge

Extend semantic entity addressing to all workspace-capability task tracks with causal event-boundary writes, and add an optional two-pass workspace-to-KV side channel. Add validation, diagnostics, evaluation microbatching, documentation, and coverage for mixed-track behavior.

* Add exact Nth-farthest workspace readout

Add an optional workspace path that decodes candidate vectors from entity-addressed slots and selects the requested rank by squared distance, with entity-ID tie-breaking. Train the vector decoder with auxiliary coordinate loss and report vector and readout metrics. Document the controls and add unit and integration coverage.

* Preserve permutation diagnostics in nth-farthest mode

Include nth-farthest metadata in entity diagnostics while excluding vector targets, keeping permutation-state evaluation compatible. Add an integration test for the diagnostic path.

* Add retrieval readout and memory optimizations

Add exact workspace retrieval supervision and readout metrics, plus compact token-ID inputs and optional CPU activation offload. Reduce diagnostic overhead and semantic address-state memory, with documentation for the new options.

* linted and formatted

* test fixes

* typecheck fixes

prediction error

Add Surprise-Gated Recall Phase 1 (#40)

* Add Surprise-Gated Recall Phase 1

Add a seven-lesson SGR dataset, generator validity checks, training and evaluation diagnostics, and an optional next-token prediction head using split-graph backbone features. Document Phase 0 and Phase 1 usage.

* Trace split-graph forward instability

Add targeted and on-error numeric tracing for split-graph forward boundaries, with CSV output and failed-batch replay. Prevent duplicate CfC residuals and accumulate the graph-level residual in fp32 to reduce overflow risk.

* Add console summaries for SGR diagnostics

Print run-tagged training, phase-one, and evaluation metrics alongside existing logs, including a compact episode-grouped trace of sampled token positions.

* Normalize Mamba outputs and ease lesson gates

Add non-affine output normalization to the parallel Mamba backbone, applying it after workspace deltas to preserve the standard forward path. Lower graph traversal and surprise recall lesson advancement thresholds from 85% to 80%.

* Add SGR error-gated memory and workspace paths

Add isolated Phase 2 backbone-only SGR arms and Phase 4 prediction-error consumers for DNC writes, workspace writes, and reconsideration. Include signal diagnostics, evaluation controls, and lower the surprise-recall lesson advancement threshold.

* Add split-graph backbone-only mode

Allow split-graph runs without DNC memory or its sequential combiner, while preserving backbone MoE, workspace, training, and checkpoint support. Also forward elapsed-time spans to CfC stages and keep Phase 2 Mamba residuals in fp32.

* Prevent NaNs in surprise recall paths

Run the prediction head in float32, compute loss only for valid stream positions, and mask excluded positions with `masked_fill` so NaNs cannot propagate through zero multipliers.

* Stabilize Mamba residuals and configure log intervals

Bound mixer residual RMS across Mamba paths while keeping residual streams in FP32, and apply the configured MoE residual scale to workspace specialist candidates. Add a configurable training summary interval via `--log-every`.

* Linted, formatted + a few bug fixes.

* Harden masked losses and add regression coverage

Select valid positions before cross-entropy to prevent masked non-finite logits from contaminating losses, and return differentiable zeros for empty selections. Add coverage for surprise-recall generation, next-token signals, numerical stability, and workspace diagnostics; update related documentation.

* test and linting fixes

broadcasting

Add MoE workspace and diagnostics (#38)

* Add MoE workspace and diagnostics

Adds an optional shared workspace for split-graph MoE models, including competitive slot writes and broadcast readback to specialists. This also adds per-bank MoE diagnostics, segment-level activation/gradient logging, paired memory-on/off eval replay, and stricter finite-gradient/NaN safeguards for recurrent training. The change keeps existing behavior unchanged unless the new --workspace and related flags are enabled.

* Improve split-graph MoE routing and diagnostics

This patch keeps MoE additive around the original controller signal instead of replacing it at initialization, zero-initializes expert output projections, and fixes top-1 routing so the router still receives useful task gradients. Capacity is now assigned by route confidence rather than flattened token order, which avoids first-source bias during multi-source CfC routing and adds accepted/dropped source diagnostics for training logs and inference reports.

It also adds split-graph MoE placement controls (`all`, `backbone`, `combiner`) so backbone and combiner MoE effects can be isolated cleanly, and updates the README with the placement behavior and a short diagnostic smoke-test matrix for comparing MoE effects on memory-learning transitions. The change keeps the baseline CfC fusion path intact while letting source-aware experts provide learned residual corrections.

* Add backbone MoE isolation diagnostics

Adds optional split-graph backbone MoE isolation per zero-based block and wires it through the training/run, loader, and model configs. This enables backbone-only MoE experiments, validates incompatible options (workspace/bad ranges), and emits valid-vs-padded routing diagnostics for graph padding without masking the actual execution path.

* Add backbone MoE residual scaling control

This change adds a configurable backbone MoE residual scale for split-graph runs, defaulting to 1.0 to preserve existing behavior. The scale is validated as finite and non-negative, only allowed with backbone/all placement, rejected with the workspace broadcast path, and included in run IDs and checkpoint metadata.

The residual is applied before the backbone residual addition as x + scale * moe_delta, while leaving routing and auxiliary load-balancing behavior active. The new option is wired through the MoE layer, Mamba and CfC parallel backbone builders, and inference config loading. The README also documents the experimental usage and follow-up commands for conservative backbone-only scale testing.

* Enable fp32 DNC tracing and diagnostics

This commit forces DNC interface, addressing, and recurrent memory math onto fp32 even under CUDA autocast, and records the precision choice in startup logs/checkpoint metadata. It adds --dnc-backward-trace for split-graph runs to watch DNC read/write addressing, temporal links, usage/allocation, prod/cumprod numerics, and the first non-finite backward boundary with batch-row context; collected diagnostics are emitted during training and failure windows and include product-input stats, read-weight mass checks, and memory/read/link magnitudes.

* Enable residual-scaled workspace MoE broadcasts

This change allows the workspace broadcast path to run with the split-graph backbone MoE residual scale instead of rejecting it. The final-bank broadcast delta is scaled with the matched backbone setting, and active workspace attention/slot updates run in fp32 outside CUDA autocast while the zero-slot/closed-gate bypass remains unchanged. Startup reporting and the README were updated to describe the new phase-1 configuration, fp32 precision behavior, and the relevant MoE/DNC validation commands.

* Add frozen workspace warm-start and MoE diagnostics

This commit adds a warm-start path for frozen workspace experiments: a compatible workspace-disabled checkpoint can be resumed into a new frozen-workspace run without reinitializing backbone, MoE, combiner, DNC, stochastic-write, or task-head weights. It validates model metadata and shared tensor shapes, keeps newly introduced workspace parameters freshly initialized, and resets optimizer, scaler, RNG, curriculum, and step counters to start the workspace from lesson 1. Existing checkpoints with workspace state still follow the normal exact-resume behavior.

It also adds optional MoE routing diagnostics control. Detailed routing/activation diagnostics can be disabled to reduce overhead while keeping the auxiliary load-balancing loss and routing logic active. The code resets cached diagnostics cleanly when toggled, and the training loop now conditionally emits the more expensive MoE diagnostics. The README is updated to document the new warm-start workflow and expected behavior.

* Add optional segment/workspace diagnostics

Adds a segment-diagnostics flag to the training CLI and passes it into SegmentDiagnostics, which now skips hook registration when disabled. WorkspaceBroadcast also gains a diagnostics toggle to avoid expensive per-timestep reductions, and the training entrypoint reports whether MoE and segment diagnostics are enabled. This keeps long recurrent runs lighter while preserving the optional trace output when needed.

* Add eval-only checkpoint evaluation mode

Add evaluation-only resume flows that load checkpoint weights without restoring optimizer state or training updates. The new flags support deterministic eval RNGs and workspace counterfactual modes (checkpoint/reset-each-step/bypass), while preventing lesson advancement during evaluation-only passes so checkpoint scoring stays comparable.

* Add experimental bistable workspace gate

Introduce an opt-in recurrent stochastic write gate with checkpointed noise, gate-only warm starts, stability checks, and training/evaluation diagnostics. Document Phase 4 usage and add unit tests; preserve the model's prior training mode during traversal evaluation.

* Optimize DNC memory allocation

Add an opt-in exact allocation path that scatters weights through the existing usage permutation, eliminating a redundant sort while preserving state-dict compatibility. Wire the option through training, inference, and memory resizing, and add equivalence tests.

* Stabilize bistable gate gradients and diagnostics

Add configurable truncated or unrolled gradients for bistable settling, defaulting to truncated backpropagation while preserving the full forward trajectory. Extend gradient and activation diagnostics, route CfC steps through module hooks, and validate frozen gate-only training parameters.

* Add configurable workspace BPTT truncation

Allow gradients through persistent workspace slots to be truncated at a configurable horizon without changing the forward trajectory. Wire the setting through training, checkpoints, and inference; reset optimizer state when the horizon changes; and extend recurrent diagnostics with state-path metrics and gradient clipping counts.

* Add DNC execution and addressing options

Add compiled memory forwarding, fast and top-k allocation, temporal-link sharpening, and configurable batch sizes. Skip dense traversal products when links are ablated, preserve these settings when resizing memory, and document the tradeoffs.

* Add workspace reconsideration and diagnostics

Allow high-innovation proposals to reopen writes to occupied workspace slots, with configurable thresholds and a rate-matched random trigger. Add bistable, CFC, and DNC gradient tracing and training diagnostics to help investigate numerical issues.

* formatted and linted

* Preserve exact bistable gate outputs

Avoid floating-point round-off in straight-through hard writes and tighten type and nullability handling across MoE, diagnostics, and controller code.

* versioning

* format

* typecheck fix

* linter fix

* Fix typing for segment diagnostic hooks

Mark runtime callback assignments on PyTorch modules as intentional typing exceptions and reformat tensor-gradient hook lambdas without changing their behavior.

* typecheck fix

New experimental memory write and read options

Add split-graph DNC mechanism variants (#37)

* Add split-graph DNC mechanism variants

This change adds configurable split-graph DNC read/write mechanisms: relational-mlp read scoring and a Kanerva-style closed-form write posterior, wired through controller config, CLI arguments, and model loading. It also adds graph traversal probes for relational and Q99 order-sensitivity checks, updates experiment CSV logging for Track A/B diagnostics, and preserves mechanism-aware memory resizing and link-matrix ablations.

* Relax minimum relational probe cell size

Lower the minimum required `cell_size` for relational probes from 12 to 6 so tiny smoke-test configurations can run without aborting before training. The probe only needs three role channels, and the previous 4-dimension-per-channel requirement was unnecessarily restrictive for small fixtures.

* Add relational residual DNC reader

Add a new `relational-residual` read variant that keeps the DNC cosine baseline at initialization by zero-initializing the learned residual output, and extend split-graph diagnostics with read-content entropy and max-weight metrics. Update the config and CLI/docs to expose the variant and document the Track A/B experiment logs and reader choices.

* Add configurable residual read scaling

This change introduces a configurable `--dnc-read-residual-scale` for the relational-residual DNC read variant, preserving the cosine baseline at initialization while allowing the learned correction to be strengthened or weakened. The setting is threaded through config, training, checkpoint loading, and memory resizing, and diagnostics now report the residual RMS and residual-to-cosine RMS ratio. The README is updated with the new flag and experiment guidance.

* Lower DNC residual read scale to 1.0

This updates the default relational-residual read correction from 2.0 to 1.0 across config and runtime defaults, and clarifies the docs that the residual remains cosine-anchored and exact at initialization. It also notes that lower scales such as 0.5 are available for follow-up comparisons while keeping the baseline and Track A setup consistent.

* Tune DNC residual read gating and Q99 probes

Adds a configurable cap for relational DNC residual read scores, exposed through config and CLI, and threads it through model loading and memory resizing. The Q99 order-sensitivity check is also refactored into separate fact-order and query-order probes with clearer semantic metrics and post-training gating. This keeps the residual correction bounded while making the order-sensitivity diagnostics more interpretable and robust. Includes a small .gitignore update for tmp files.

* Stabilize DNC residual normalization

Avoid zero-valued RMS denominators in residual score normalization during the first forward pass. This change adds a small epsilon before taking the square root for both residual and cosine RMS values, preventing NaNs and infinite gradients when the scorer is initialized to zero.

Note that these are purely experimental and do not replace the defaults and are not recommended to use for production yet

usable architecture modules

Add reusable prepared datasets (#34)

* Add reusable prepared datasets

This adds save/load support for prepared datasets so training and inference can reuse processed data without re-tokenizing raw sources. It introduces --save, --load-prepared, and --prepare-only flows, persists the resolved dataset state, and wires prepared datasets through task construction while validating type compatibility and path safety.

* Fix source import path for training entrypoint

Add the src root to sys.path so the training module works when launched either as `python -m initium.core_training` or directly from a source checkout. This keeps imports consistent regardless of execution context.

* Add classic runtime config overrides

Classic datasets now accept runtime overrides for window size, probe distances, gamma, and OOD eval episodes even when loaded from prepared data. Training also adds a --batch-size option with a classic default of 1 to keep full-window DNC activations within memory, and the generated run ID reflects the selected classic batch size. The classic inference path applies the same runtime settings to prepared streams, and the output projection path no longer requires a contiguous tensor.

* Reduce CUDA syncs in training and compat paths

This change avoids expensive per-step host/device syncs during training by keeping scalar loss terms on-device until logging intervals and computing averaged values from accumulated tensors. Finite-value checks are now limited to the logging cadence, reducing noisy overhead in classic-track runs.

It also tightens the classic task metrics path by keeping loss terms as tensors and computing diversity without per-row Python object conversion, and updates the NVRTC compatibility shim to avoid repeated CUDA-to-CPU syncs on non-negative DNC inputs while adding a strict debug mode for negative-input detection.

* Compact classic token encodings

Classic datasets now encode each token with the existing 3-digit codec instead of a 1000-wide one-hot slot, while preserving a separate per-modality probe marker. This updates dataset sizing, window construction, generated-token encoding, and the classic track diagnostic output, and clarifies the CUDA no-sync compatibility patch message.

* formatted

* Add decode_digits and use it for target decoding

Add DigitCodec.decode_digits to convert a sequence of digit IDs (tensor or list) into an integer label with input validation. Replace calls to decode_field on target digit arrays in ClassicDataset and ClassicInferenceTask with decode_digits so already-decoded ID sequences are handled correctly (avoids treating them as logits/floats). Improves error messages for wrong lengths or out-of-range digit IDs.

* Expose stable architecture API

This change turns the package root into a lazy, stable public API for architecture components and controller wrappers. It centralizes exports via __getattr__ so optional training and inference dependencies are not imported eagerly, keeps internal data encoders and task implementations out of the public contract, and adds checkpoint/model loading helpers for easier integration.

* Remove classic long-window task track

This removes the additive classic long-window dataset/task family and its CLI flags, including the classic configs, dataset registry entries, inference generation paths, and training logs tied to memory probe evaluation. It also simplifies OGS scoring to the core offset/memory/perfect components and updates attribution reporting to be architecture-agnostic rather than classic-specific.

* Annotate final_ood_field_log dict type

Add an explicit type annotation for final_ood_field_log in src/initium/core_training.py: dict[tuple[int,int], list[int]]. This improves type clarity and helps static type checkers; no functional behavior changed.

initial release

Py pi publication and automation (#33)

* chore: rename src/model to src/initium

* Adopt src/initium package layout

This commit migrates the project from top-level sibling imports to the conventional `initium.*` package namespace and updates the packaging metadata, entry points, and documentation to match the `src/initium` layout. It also adds the packaged `initium.setup_wheels` module and updates CI to invoke the wheel installer via `python -m initium.setup_wheels`.

* Fix test imports for package layout

Update test fixtures and checkpoint reload tests to import the training module from the package namespace (`initium.core_training`) instead of the top-level module. This matches the project layout and ensures the tests exercise the correct module during synthetic training and checkpoint restore flows.

* Fix dnc direct_url detection

Guard against `None` values when checking whether `dnc` was installed from a VCS checkout. This avoids false negatives when `direct_url.json` is missing and keeps the wheel setup logic aligned with git-based installs.

* Add GitHub workflow to publish to PyPI

Add .github/workflows/publish-pypi.yml which builds and validates source/wheel distributions and publishes them to PyPI on published releases. The workflow checks the release commit is reachable from main and that the release tag matches __version__ in src/__init__.py, uses Python 3.12, runs build and twine check, and publishes via pypa/gh-action-pypi-publish with OIDC (no stored PyPI token). Update README.md with publishing instructions and one-time Trusted Publisher setup for the pypi environment.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* formatted

* Rename PyPI package to initium-deeplearning

Update the package metadata and publishing workflow to publish the distribution under `initium-deeplearning` while keeping the Python import namespace as `initium`. The README now documents the new PyPI project name and Trusted Publisher setup so release instructions match the actual package configuration.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>