Phase 1: Implement Eq. 23 DMGD with MLP-parameterized momentum

[georgepullen]

Purpose

Implement DMGD with MLP-parameterized momentum per Eq. 23 intent (not tanh-on-EMA heuristic).

Mandatory Reading (blocking)

First comment must summarize:

• reports/NL_IMPLEMENTATION_ORACLE.md section 6.1.1
• reports/paper/NL-print.extracted.clean.txt Eq. (23)
• src/nested_learning/optim/deep.py current dmgd behavior

Required Code Anchors

• src/nested_learning/optim/deep.py
• src/nested_learning/optim/factory.py
• src/nested_learning/optim/manager.py

Scope

• Add dmgd_mlp variant with small momentum network.
• Add explicit inner-objective update on momentum-network parameters.
• Preserve legacy dmgd behavior under legacy name.

Test Requirements

• Verify momentum-network params update each step.
• Verify deterministic behavior under fixed seed.
• Keep existing tests green.

Deliverables

• Implementation + docs + ablation config.

Acceptance Criteria

• dmgd_mlp tested with non-zero parameter update evidence.
• 5k pilot run stable.
• First issue comment contains mandatory reading summary.

[georgepullen]

Mandatory read summary before implementation (oracle + paper + code):

• Oracle §6.1.1 flags current dmgd as a tanh-on-EMA heuristic and explicitly calls for an MLP-parameterized momentum path (dmgd_mlp) with an inner objective update.
• Paper Eq. (23) defines DMGD as m_{i+1} = alpha_i m_i - eta_t ∇L(2)(m_i; u_i, I) and update W_{i+1} = W_i + m_{i+1}(u_i), i.e., momentum is a learned module, not just a scalar nonlinearity.
• The existing implementation in src/nested_learning/optim/deep.py had no momentum-network params and therefore no inner-objective parameter update step.
• LevelOptimizerManager applied optimizer calls under torch.no_grad(), which prevents any internal learnable optimizer module (e.g., DMGD-MLP) from receiving gradient-based inner updates.
• Context7/PyTorch autograd docs confirm expected pattern: compute internal gradients in grad mode, then mutate parameters under torch.no_grad() to avoid recording update ops in the outer graph.

Files modified for this ticket:

• src/nested_learning/optim/deep.py
• src/nested_learning/optim/manager.py
• tests/test_optim.py
• configs/ablations/selfmod_dmgd_mlp.yaml
• tests/test_paper_faithful_configs.py

Progress implemented:

• Added new inner optimizer variant: variant=dmgd_mlp.
• Implemented a learnable momentum MLP with explicit inner-objective step (mse/cosine) and per-step parameter update.
• Updated manager apply paths so optimizer internals can train while keeping model parameter writes in no_grad.
• Added tests proving DMGD-MLP network parameters update and deterministic behavior under fixed seed.
• Added ablation config for 5k pilot run (configs/ablations/selfmod_dmgd_mlp.yaml) and config composition coverage.

Risks and non-goals:

• This change does not yet claim full equation-level equivalence to every possible L(2) instantiation in Eq. (23); it implements a concrete, explicit inner-objective variant consistent with the paper’s DMGD framing.
• No changes here to Eq. (24) Newton-Schulz output path (muon_ns) or CMS projection equations (those are separate tickets).

[georgepullen]

Implementation evidence update:

• Added dmgd_mlp variant in src/nested_learning/optim/deep.py.
• Added explicit inner-objective parameter update for momentum network (supports dmgd_inner_objective=mse|cosine).
• Updated optimizer manager apply paths (src/nested_learning/optim/manager.py) so optimizer internals can update in grad mode while model parameter writes remain in torch.no_grad().
• Added tests:
  • tests/test_optim.py::test_deep_momentum_dmgd_mlp_updates_network_params
  • tests/test_optim.py::test_deep_momentum_dmgd_mlp_is_deterministic_for_fixed_seed
• Added ablation config:
  • configs/ablations/selfmod_dmgd_mlp.yaml
  • composition coverage in tests/test_paper_faithful_configs.py

Commands run:

• pytest -q tests/test_optim.py tests/test_paper_faithful_configs.py tests/test_checkpoint_resume.py tests/test_validate_fidelity_telemetry.py tests/test_model.py
• uv run python train.py --config-name pilot_smoke train.device=cpu train.steps=3 train.checkpoint.enable=true train.checkpoint.dir=artifacts/checkpoints/dmgd_mlp_smoke train.checkpoint.save_interval=3 train.checkpoint.save_last=true +model.optimizers.titan_opt.params.variant=dmgd_mlp +model.optimizers.titan_opt.params.dmgd_hidden_dim=8 +model.optimizers.titan_opt.params.dmgd_inner_lr=1e-3 +model.optimizers.cms_opt.params.variant=dmgd_mlp +model.optimizers.cms_opt.params.dmgd_hidden_dim=8 +model.optimizers.cms_opt.params.dmgd_inner_lr=1e-3 logging.enabled=true logging.backend=json logging.path=logs/dmgd_mlp_smoke_metrics.json
• uv run python scripts/checkpoint/verify.py --checkpoint artifacts/checkpoints/dmgd_mlp_smoke/step_000003.pt
• uv run python scripts/checks/validate_fidelity_telemetry.py --log logs/dmgd_mlp_smoke_metrics.json

Result: all commands passed.