Initial implementation: all models, methods, toy and CIFAR experiments

Debug phase. Toy LQ experiments (3 seeds) complete with terminal gradient matching.
Credit bridge matches state bridge on linear system (~0.94 cosine).
CIFAR experiments in progress.

1 files changed, 335 insertions, 0 deletions

diff --git a/results/toy_lq/toy_lq_seed42.json b/results/toy_lq/toy_lq_seed42.json
new file mode 100644
index 0000000..0a821be
--- /dev/null
+++ b/results/toy_lq/toy_lq_seed42.json
@@ -0,0 +1,335 @@
+{
+  "config": {
+    "d_hidden": 64,
+    "output_dim": 10,
+    "num_layers": 12,
+    "sigma": 0.03,
+    "batch_size": 256,
+    "num_steps": 5000,
+    "lr_fb": 0.001,
+    "lam": 0.1,
+    "K": 8,
+    "ema_momentum": 0.995,
+    "sigma_bridge": 0.03,
+    "eval_every": 500,
+    "seed": 42,
+    "gpu": 0,
+    "output_dir": "results/toy_lq"
+  },
+  "log": {
+    "steps": [
+      1,
+      500,
+      1000,
+      1500,
+      2000,
+      2500,
+      3000,
+      3500,
+      4000,
+      4500,
+      5000
+    ],
+    "state_bridge_loss": [
+      66.01814270019531,
+      2.058396816253662,
+      2.191567897796631,
+      2.1112077236175537,
+      2.2748100757598877,
+      2.1154427528381348,
+      2.0553503036499023,
+      1.966332197189331,
+      2.1147494316101074,
+      2.007577896118164,
+      2.0389387607574463
+    ],
+    "credit_bridge_loss": [
+      112.22884368896484,
+      0.5559965372085571,
+      0.36617255210876465,
+      0.22275042533874512,
+      0.13448813557624817,
+      0.0859837457537651,
+      0.06448937207460403,
+      0.03930438682436943,
+      0.033317387104034424,
+      0.03704077750444412,
+      0.02309737727046013
+    ],
+    "dfa_costate_cos": [
+      0.0011988391992277824,
+      0.007705975646296373,
+      -0.0006242827870524847,
+      0.0037568132393062115,
+      0.004209253507164808,
+      -0.0012397096635630499,
+      0.0033803660816584644,
+      0.003754911944270134,
+      0.002183924035097544,
+      0.0057398807973815845,
+      0.0028941635615550554
+    ],
+    "state_costate_cos": [
+      0.009303608121207127,
+      0.9451924241506137,
+      0.9427064611361577,
+      0.9456643370481638,
+      0.947175892499777,
+      0.9486432488148029,
+      0.9405259856810937,
+      0.9429925313362708,
+      0.9463646090947665,
+      0.9434385620630704,
+      0.9468063895518963
+    ],
+    "credit_costate_cos": [
+      0.03412332414434506,
+      0.3051199523302225,
+      0.2641584941974053,
+      0.17990898627501267,
+      0.11904546274588658,
+      0.03094297769264533,
+      0.022915477577883463,
+      0.008693795901938127,
+      0.0027211602920523058,
+      0.020937439484091904,
+      0.033342053540624104
+    ],
+    "dfa_rho": [
+      0.005011526596111556,
+      0.0020135376447190842,
+      -0.011304221504057447,
+      0.003935044708972176,
+      0.0159942601264144,
+      -0.011545649264007807,
+      0.01096861291443929,
+      0.0007782066240906715,
+      -0.015019190264865756,
+      0.007689292387415965,
+      0.006176682732378443
+    ],
+    "state_rho": [
+      0.011923154350370169,
+      0.9337877084811529,
+      0.9292550335327784,
+      0.9341330577929815,
+      0.9364036619663239,
+      0.9322425921758016,
+      0.9240961174170176,
+      0.9329939633607864,
+      0.9331585764884949,
+      0.9324665367603302,
+      0.9281178514162699
+    ],
+    "credit_rho": [
+      0.031667908265565835,
+      0.31275976697603863,
+      0.2433429310719172,
+      0.1816188059747219,
+      0.11667139704028766,
+      0.013198353117331862,
+      0.022044080891646445,
+      -0.007547003333456814,
+      -0.011566813724736372,
+      -0.003948230994865298,
+      0.016950203105807304
+    ],
+    "dfa_nudge": [
+      -0.0003799900102118651,
+      -0.0025626374408602715,
+      0.0017628272374471028,
+      -0.001205168974896272,
+      -0.0011821148606638114,
+      0.0014717701512078445,
+      -8.787959814071655e-05,
+      -0.0006076549955954155,
+      0.0005303900688886642,
+      -0.0014991160326947768,
+      -0.0005284918782611688
+    ],
+    "state_nudge": [
+      -0.002327537008871635,
+      -0.34574924657742184,
+      -0.3358767156799634,
+      -0.33698513607184094,
+      -0.3561149264375369,
+      -0.33514803399642307,
+      -0.3557068184018135,
+      -0.3218521823485692,
+      -0.33325668424367905,
+      -0.34358637283245724,
+      -0.34090926001469296
+    ],
+    "credit_nudge": [
+      -0.014598140881086389,
+      -0.11370646270612876,
+      -0.09322128010292847,
+      -0.06459770910441875,
+      -0.04436610918492079,
+      -0.008244700108965239,
+      -0.00636714743450284,
+      0.0006991980286935965,
+      0.0036097665627797446,
+      -0.0035287897723416486,
+      -0.007594603579491377
+    ],
+    "bridge_residual": [
+      0.06566914729773998,
+      0.37236853316426277,
+      0.3114783614873886,
+      0.26258066420753795,
+      0.20779911428689957,
+      0.13781529137243828,
+      0.10999641008675098,
+      0.07969006771842639,
+      0.06178054213523865,
+      0.07378745886186759,
+      0.0832565538585186
+    ]
+  },
+  "final_per_layer": {
+    "dfa_costate_cos": [
+      -0.039167389273643494,
+      -0.0378018394112587,
+      0.005690325051546097,
+      -0.023073989897966385,
+      -0.0005057593807578087,
+      -0.014485953375697136,
+      0.03301015496253967,
+      0.04401148855686188,
+      0.054177843034267426,
+      0.03981431573629379,
+      -0.04246171563863754,
+      -0.0012151142582297325,
+      0.01963176019489765
+    ],
+    "state_costate_cos": [
+      0.9444395303726196,
+      0.9453534483909607,
+      0.9460644721984863,
+      0.9466040134429932,
+      0.9469730257987976,
+      0.9472954273223877,
+      0.9476633667945862,
+      0.9478192925453186,
+      0.947782576084137,
+      0.9473406672477722,
+      0.947346568107605,
+      0.9471475481987,
+      0.9466531276702881
+    ],
+    "credit_costate_cos": [
+      0.04752141237258911,
+      0.04377397149801254,
+      0.04051002860069275,
+      0.03716123104095459,
+      0.0350164994597435,
+      0.03194836527109146,
+      0.02999947965145111,
+      0.02913709171116352,
+      0.027683185413479805,
+      0.0277146864682436,
+      0.027703404426574707,
+      0.027434173971414566,
+      0.027843166142702103
+    ],
+    "dfa_rho": [
+      -0.04803554713726044,
+      0.001050771214067936,
+      0.008967258036136627,
+      -0.0271889790892601,
+      0.02336559258401394,
+      -0.018210411071777344,
+      0.05891512706875801,
+      0.040720634162425995,
+      0.07478035986423492,
+      0.04802168905735016,
+      -0.06280035525560379,
+      -0.0254659466445446
+    ],
+    "state_rho": [
+      0.9311305284500122,
+      0.9222633838653564,
+      0.9287852644920349,
+      0.9287664890289307,
+      0.9245603084564209,
+      0.928197979927063,
+      0.9275168180465698,
+      0.9290561676025391,
+      0.9267844557762146,
+      0.927483081817627,
+      0.9315245151519775,
+      0.9313452243804932
+    ],
+    "credit_rho": [
+      0.05740518122911453,
+      0.035541512072086334,
+      0.002091987058520317,
+      0.024556485936045647,
+      -0.006993812508881092,
+      0.03284040838479996,
+      0.012268777936697006,
+      -0.004999782890081406,
+      0.014774687588214874,
+      -0.010628825053572655,
+      0.05940534919500351,
+      -0.012859531678259373
+    ],
+    "dfa_nudge": [
+      0.01363457553088665,
+      0.013758538290858269,
+      -0.0032786596566438675,
+      0.010209780186414719,
+      -0.0013850200921297073,
+      0.004463233053684235,
+      -0.012735363095998764,
+      -0.01801125332713127,
+      -0.019914839416742325,
+      -0.012350432574748993,
+      0.017156170681118965,
+      0.002111367881298065
+    ],
+    "state_nudge": [
+      -0.34207087755203247,
+      -0.3417494297027588,
+      -0.3413676619529724,
+      -0.3414176106452942,
+      -0.3412688076496124,
+      -0.34025871753692627,
+      -0.340381920337677,
+      -0.34054237604141235,
+      -0.3402785062789917,
+      -0.34061968326568604,
+      -0.34013575315475464,
+      -0.340819776058197
+    ],
+    "credit_nudge": [
+      -0.012597911059856415,
+      -0.011231275275349617,
+      -0.009987404569983482,
+      -0.008699672296643257,
+      -0.007977090775966644,
+      -0.006919408217072487,
+      -0.006156830117106438,
+      -0.005891043692827225,
+      -0.00546320341527462,
+      -0.005353953689336777,
+      -0.005423387512564659,
+      -0.005434062331914902
+    ],
+    "bridge_residual": [
+      0.10028190910816193,
+      0.09995798766613007,
+      0.09832010418176651,
+      0.09528136253356934,
+      0.09090456366539001,
+      0.0843343734741211,
+      0.07823348790407181,
+      0.0723348930478096,
+      0.06821676343679428,
+      0.06773027032613754,
+      0.06943590939044952,
+      0.0740470215678215
+    ]
+  }
+}
\ No newline at end of file