{
  "title": "dimuon-lite vs AdamW, LoRA SFT on Qwen3.5-2B, no_robots held-out NLL",
  "hypothesis": "With per-optimizer learning rates picked on a validation split, LoRA SFT with dimuon-lite cuts held-out no_robots token_mean NLL by at least 0.02 nats versus AdamW. Both arms use the same 100 steps, batches and seeds.",
  "method": "Train a rank-32 LoRA on Qwen/Qwen3.5-2B over HuggingFaceH4/no_robots, using one tinker SDK script for both arms so the only difference is the optimizer. The cookbook chat_sl recipe is AdamW-only, so running AdamW through the same script removes any recipe difference. DiMuon is selected at creation with optimizer=tinker.types.DimuonOptimizerConfig() and stepped with DimuonParams(learning_rate, beta1=0.9, grad_clip_norm=0.0). AdamW uses tinker.AdamParams(learning_rate, beta1=0.9, beta2=0.95, eps=1e-8). dimuon-lite is a recently landed reverse-engineered update rule: Polar Express orthogonalization with about 5 iterations, rank_lr_correction_exponent 0.175, Muon on LoRA B, and no beta2. Its update is checked against the trusted DiMuon reconstruction before the comparison runs. Phase 1, LR selection: each optimizer runs 40 steps at seed 0 on a 3-point LR grid fixed now (AdamW {5e-5, 1e-4, 2e-4}; DiMuon {3e-4, 1e-3, 3e-3}). Each run is scored by token_mean NLL on 500 validation items held out from the train split, and the best LR per optimizer is frozen. Phase 2, confirmation: each optimizer trains 100 steps at its frozen LR for seeds {0,1,2}. Batches are identical across arms within a seed: the same shuffled order of the remaining 9,000 train rows, batch_size 32, max_seq_len 2048. Data order and LoRA init seed are shared per seed. Primary eval is token_mean NLL on the full 500-example no_robots test split, computed from per-example summed logprobs times loss weights via forward(). Secondary, descriptive only: IFEval prompt-level strict accuracy (541 prompts, greedy, max_tokens 1024) on each final checkpoint. IFEval is dropped first if feasibility does not fit.",
  "base_model": "Qwen/Qwen3.5-2B",
  "lora_rank": 32,
  "training_type": "sft",
  "dataset": {
    "name": "no_robots",
    "hf_path": "HuggingFaceH4/no_robots",
    "split": "train (9000 for training, 500 held out for LR selection); test (500) for confirmation",
    "n_examples": 9500
  },
  "optimizer": {
    "name": "dimuon-lite",
    "lr": 0.001
  },
  "steps": 100,
  "batch_size": 32,
  "max_seq_len": 2048,
  "eval": {
    "benchmark": "no_robots test split (HuggingFaceH4/no_robots, split=test)",
    "metric": "held-out token_mean NLL (nats per assistant loss token)",
    "n_examples": 500,
    "aggregation": "token_mean"
  },
  "primary_endpoint": "no_robots_test_token_mean_nll_at_step_100",
  "arms": [
    {
      "name": "lrsearch_adamw_5e-5_s0",
      "phase": "selection",
      "optimizer": "adam",
      "lr": 5e-05,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "lrsearch_adamw_1e-4_s0",
      "phase": "selection",
      "optimizer": "adam",
      "lr": 0.0001,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "lrsearch_adamw_2e-4_s0",
      "phase": "selection",
      "optimizer": "adam",
      "lr": 0.0002,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "lrsearch_dimuon_3e-4_s0",
      "phase": "selection",
      "optimizer": "dimuon-lite",
      "lr": 0.0003,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "lrsearch_dimuon_1e-3_s0",
      "phase": "selection",
      "optimizer": "dimuon-lite",
      "lr": 0.001,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "lrsearch_dimuon_3e-3_s0",
      "phase": "selection",
      "optimizer": "dimuon-lite",
      "lr": 0.003,
      "steps": 40,
      "seed": 0
    },
    {
      "name": "baseline_adamw_s0",
      "phase": "confirmation",
      "optimizer": "adam",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 0
    },
    {
      "name": "baseline_adamw_s1",
      "phase": "confirmation",
      "optimizer": "adam",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 1
    },
    {
      "name": "baseline_adamw_s2",
      "phase": "confirmation",
      "optimizer": "adam",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 2
    },
    {
      "name": "treatment_dimuon_s0",
      "phase": "confirmation",
      "optimizer": "dimuon-lite",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 0
    },
    {
      "name": "treatment_dimuon_s1",
      "phase": "confirmation",
      "optimizer": "dimuon-lite",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 1
    },
    {
      "name": "treatment_dimuon_s2",
      "phase": "confirmation",
      "optimizer": "dimuon-lite",
      "lr": "frozen from selection",
      "steps": 100,
      "seed": 2
    }
  ],
  "baseline": "AdamW (tinker.AdamParams, beta1 0.9, beta2 0.95, eps 1e-8) through the same SDK script, with the same model, LoRA rank, init seed, batches and order, 100 steps, and eval. Its LR is chosen by the same 3-point validation search.",
  "experiment_features": {
    "custom_loss": false,
    "custom_optimizer": true,
    "paired_comparison": true,
    "paired_training_arms": true,
    "checkpoint_restore": false,
    "warmup_adaptation": false,
    "generated_training_data": false,
    "learned_generator": false,
    "heldout_loss": true,
    "adaptive_selection": true
  },
  "correctness_preflight": {
    "required": true,
    "status": "unresolved: spindle and miles commits for #85/#17 and the commit of Aaron's DiMuon reconstruction must be pinned before any full sweep",
    "paths": [
      {
        "id": "spindle_dimuon_lite_optimizer",
        "kind": "optimizer",
        "baseline_revision": "UNRESOLVED: pin the modal-projects/spindle commit deployed at the shared Spindle URL (post-#85). AdamW updates on this deployment must match the pre-#85 AdamW update on identical grads/state.",
        "reference_revision": "UNRESOLVED: pin the commit of Aaron's DiMuon reconstruction (0.09% relative update error vs Tinker) and the modal-projects/miles commit of #17",
        "tolerances": {
          "baseline_update": {
            "atol": 1e-06,
            "rtol": 0.001,
            "min_cosine": 0.99999
          },
          "reference_update": {
            "atol": 1e-05,
            "rtol": 0.01,
            "min_cosine": 0.999
          }
        },
        "lr_scaling": {
          "learning_rates": [
            0.0001,
            0.001,
            0.01
          ],
          "rtol": 0.01
        }
      }
    ],
    "notes": "Tolerances are fixed for bf16 LoRA weights with fp32 optimizer math before measuring. reference_update rtol 1e-2 leaves headroom over the reported 0.09% error. Feed the same LoRA state and gradient to the Spindle DiMuon step and to the reference for one step. With grad_clip_norm=0 and beta1 momentum seeded from the first gradient, the update should scale linearly in LR across a 100x span."
  },
  "activation_gate": {
    "applicable": true,
    "reason": "The paired comparison only tests DiMuon if the executed optimizer really orthogonalizes LoRA B updates and both arms ran the declared configuration. Polar Express makes the first-step B update near-orthogonal (flat singular spectrum); AdamW's is not. This is measured from saved weight diffs, not SDK introspection.",
    "checks": [
      {
        "metric": "executed_base_model",
        "measures": "model name passed to create_lora_training_client in every arm",
        "one_of": [
          "Qwen/Qwen3.5-2B"
        ]
      },
      {
        "metric": "executed_optimizer_step_type",
        "measures": "type of params object passed to optim_step in each arm (logged client-side)",
        "one_of": [
          "DimuonParams",
          "AdamParams"
        ]
      },
      {
        "metric": "dimuon_B_update_normalized_stable_rank_step1",
        "measures": "||dB||_F^2 / (||dB||_2^2 * r) for the first DiMuon update of LoRA B, averaged over layers, from downloaded checkpoints at step 0 and 1",
        "min": 0.6,
        "max": 1.05
      },
      {
        "metric": "batch_hash_match_across_arms",
        "measures": "fraction of the 100 training batches whose token-id hash is identical between paired arms within a seed",
        "min": 1.0,
        "max": 1.0
      },
      {
        "metric": "pilot_train_nll_drop_dimuon_10steps",
        "measures": "train token_mean NLL at step 0 minus step 10 in the DiMuon pilot at the middle LR; must be finite and positive",
        "min": 0.0
      }
    ]
  },
  "restoration_preflight": {
    "applicable": true,
    "reason": "AdamW and DiMuon arms are trained from a shared fresh-LoRA start within each seed. Before the first update, each arm calls save_weights and downloads the adapter. The SHA-256 of the LoRA tensors must match across the paired arms; B=0 at init, and A comes from the seeded init. Base-model logprobs on a fixed tokenized batch must also match. A hash mismatch makes that seed inconclusive. Optimizer type is verified by the activation-gate step-1 spectrum check, not by invented introspection.",
    "fixed_batch_required": true,
    "groups": [
      {
        "arms": [
          {
            "name": "baseline_adamw_s0",
            "optimizer": "adam",
            "optimizer_config": {
              "beta1": 0.9,
              "beta2": 0.95,
              "eps": 1e-08,
              "lr": "frozen"
            }
          },
          {
            "name": "treatment_dimuon_s0",
            "optimizer": "dimuon-lite",
            "optimizer_config": {
              "beta1": 0.9,
              "grad_clip_norm": 0.0,
              "lr": "frozen"
            }
          }
        ]
      },
      {
        "arms": [
          {
            "name": "baseline_adamw_s1",
            "optimizer": "adam",
            "optimizer_config": {
              "beta1": 0.9,
              "beta2": 0.95,
              "eps": 1e-08,
              "lr": "frozen"
            }
          },
          {
            "name": "treatment_dimuon_s1",
            "optimizer": "dimuon-lite",
            "optimizer_config": {
              "beta1": 0.9,
              "grad_clip_norm": 0.0,
              "lr": "frozen"
            }
          }
        ]
      },
      {
        "arms": [
          {
            "name": "baseline_adamw_s2",
            "optimizer": "adam",
            "optimizer_config": {
              "beta1": 0.9,
              "beta2": 0.95,
              "eps": 1e-08,
              "lr": "frozen"
            }
          },
          {
            "name": "treatment_dimuon_s2",
            "optimizer": "dimuon-lite",
            "optimizer_config": {
              "beta1": 0.9,
              "grad_clip_norm": 0.0,
              "lr": "frozen"
            }
          }
        ]
      }
    ]
  },
  "success_criterion": "Validated if the simultaneous 95% paired_t interval, across 3 training seeds, for (AdamW NLL \u2212 DiMuon NLL) on no_robots test token_mean NLL at step 100 lies entirely above 0.02 nats. Failed if it lies entirely below 0.02. Otherwise inconclusive. IFEval strict accuracy is reported descriptively and does not decide.",
  "decision_rule": {
    "version": 1,
    "estimand": "Reduction in held-out no_robots test token_mean NLL at step 100 from replacing AdamW with dimuon-lite (each at its validation-selected LR) for rank-32 LoRA SFT on Qwen3.5-2B; larger is better for dimuon-lite.",
    "unit": "training_seed",
    "seeds": [
      0,
      1,
      2
    ],
    "alpha": 0.05,
    "method": "paired_t",
    "assumptions": {
      "independent_seed_effects": true,
      "identically_distributed_seed_effects": true,
      "normal_seed_effects": true
    },
    "criteria": [
      {
        "id": "primary",
        "metric": "no_robots_test_token_mean_nll_at_step_100",
        "weights": {
          "baseline": 1,
          "treatment": -1
        },
        "margin": 0.02
      }
    ],
    "adaptive_selection": {
      "kind": "hyperparameter",
      "rationale": "Muon-family and AdamW LoRA optimizers have very different good LRs. Each gets a prespecified 3-point grid, scored only on validation items held out from train, and the choice is frozen before the single confirmation eval on the test split.",
      "search_item_ids": "no_robots train rows held out by seeded shuffle (seed 1234), first 500 after shuffle; never used for training",
      "confirmation_item_ids": "all 500 rows of no_robots split=test"
    }
  },
  "budget": {
    "gpu_hours": 2.0,
    "max_wall_minutes": 150
  },
  "generation_audit": {
    "required": false,
    "answer_comparison": null,
    "min_sample_size": 0,
    "max_label_errors": 0,
    "max_eval_overlaps": 0,
    "reason": "fixed public dataset; no generated training data"
  },
  "feasibility_required": true,
  "feasibility": {
    "config": null,
    "effective_config": {
      "base_model": "Qwen/Qwen3.5-2B",
      "arms": 12,
      "selection_steps": 40,
      "confirmation_steps": 100,
      "batch_size": 32,
      "max_seq_len": 2048,
      "n_eval_validation": 500,
      "n_eval_test": 500,
      "n_ifeval_secondary": 541,
      "planned_arm_steps": 840
    },
    "primary_endpoint": "no_robots_test_token_mean_nll_at_step_100",
    "source": null,
    "measurements": null,
    "generation": null,
    "recovery": null,
    "downsizing_order": "drop IFEval secondary, then shrink LR grid to 2 points per optimizer, then shorten selection runs to 25 steps; keep 3 seeds and the 100-step test NLL. If that still does not fit, report inconclusive."
  },
  "command": "# Custom part: one SDK SFT loop shared by both optimizers (cookbook chat_sl is AdamW-only).\nimport os, sys, tinker, datasets, random\nfrom tinker_cookbook import renderers, tokenizer_utils\nfrom tinker_cookbook.supervised.data import conversation_to_datum\nOPT, LR, SEED, STEPS = sys.argv[1], float(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4])\nMODEL='Qwen/Qwen3.5-2B'; BS=32; MAXLEN=2048\ntok=tokenizer_utils.get_tokenizer(MODEL); rend=renderers.get_renderer(renderers.get_recommended_renderer_name(MODEL) if hasattr(renderers,'get_recommended_renderer_name') else 'qwen3', tok)\nds=datasets.load_dataset('HuggingFaceH4/no_robots'); tr=list(ds['train']); random.Random(1234).shuffle(tr)\nval, train = tr[:500], tr[500:]; random.Random(SEED).shuffle(train)\nmk=lambda r: conversation_to_datum(r['messages'], rend, MAXLEN)\nsc=tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\nkw={'optimizer': tinker.types.DimuonOptimizerConfig()} if OPT=='dimuon' else {}\ntc=sc.create_lora_training_client(base_model=MODEL, rank=32, seed=SEED, **kw)\ntc.save_weights_for_sampler(name=f'init-{OPT}-s{SEED}').result()  # hash for restoration preflight\nstep_params = (lambda: tinker.types.DimuonParams(learning_rate=LR, beta1=0.9, grad_clip_norm=0.0)) if OPT=='dimuon' else (lambda: tinker.AdamParams(learning_rate=LR, beta1=0.9, beta2=0.95, eps=1e-8))\nfor s in range(STEPS):\n    batch=[mk(r) for r in train[s*BS:(s+1)*BS]]\n    tc.forward_backward(batch, loss_fn='cross_entropy').result(); tc.optim_step(step_params()).result()\ndef nll(rows):  # token_mean from per-example totals, never batch loss:mean\n    tot=0.0; n=0.0\n    for i in range(0,len(rows),32):\n        d=[mk(r) for r in rows[i:i+32]]; out=tc.forward(d, loss_fn='cross_entropy').result()\n        for di,o in zip(d,out.loss_fn_outputs):\n            w=di.loss_fn_inputs['weights'].to_numpy(); lp=o['logprobs'].to_numpy(); tot+=-(lp*w).sum(); n+=w.sum()\n    return tot/n\nsplit = val if STEPS==40 else list(ds['test'])\nprint({'opt':OPT,'lr':LR,'seed':SEED,'steps':STEPS,'token_mean_nll':nll(split)})\ntc.save_weights_for_sampler(name=f'final-{OPT}-lr{LR}-s{SEED}').result()\n# run: python sft_opt.py {adam|dimuon} <lr> <seed> {40|100}",
  "feasibility_check": {
    "status": "unknown",
    "problems": [
      "Effective configuration not recorded",
      "Pilot evidence source not recorded",
      "Representative training timing not recorded",
      "Representative setup timing not recorded",
      "Representative evaluation timing not recorded",
      "Positive recovery reserve not recorded"
    ],
    "wall_seconds": null,
    "gpu_hours": null,
    "known_gpu_hours": 0,
    "phases": {},
    "generation_counts": null,
    "generation_applicability": {
      "check": "generation_audit",
      "status": "not_applicable",
      "enforce": false,
      "reasons": [
        "spec.generation_audit opts out and no feature it covers is true"
      ],
      "features": {
        "generated_training_data": false
      }
    }
  }
}