stochastic gd

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.9%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20return%20mo%2C%20np%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Stochastic%20Gradient%20Descent%0A%0A%20%20%20%20Instead%20of%20computing%20the%20gradient%20over%20all%20%24%5Cell%24%20samples%2C%20SGD%20uses%20a%20random%0A%20%20%20%20**mini-batch**%20of%20size%20%24b%24%20at%20each%20step%3A%0A%0A%20%20%20%20%24%24w_%7Bk%2B1%7D%20%3D%20w_k%20-%20%5Ceta_k%5C%2C%20%5Cnabla_w%20q_%7Bi_k%7D(w_k)%24%24%0A%0A%20%20%20%20where%20%24i_k%24%20is%20a%20randomly%20drawn%20batch%20of%20indices%20and%3A%0A%0A%20%20%20%20%24%24%5Cnabla_w%20q_%7Bi_k%7D(w)%20%3D%20%5Cfrac%7B2%7D%7Bb%7D%5C%2C%20X_%7Bi_k%7D%5E%5Ctop%20%5Cbigl(X_%7Bi_k%7D%20w%20-%20y_%7Bi_k%7D%5Cbigr)%24%24%0A%0A%20%20%20%20This%20is%20noisier%20than%20full%20GD%20but%20much%20cheaper%20per%20step.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20%23%20---%20student%3A%20begin%20---%0A%0A%20%20%20%20def%20sgd_gradient(X%3A%20np.ndarray%2C%20y%3A%20np.ndarray%2C%20w%3A%20np.ndarray%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_size%3A%20int%2C%20rng%3A%20np.random.Generator)%20-%3E%20np.ndarray%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Sample%20a%20random%20mini-batch%20of%20%60batch_size%60%20rows%20from%20X%2C%20y%2C%0A%20%20%20%20%20%20%20%20and%20return%20the%20gradient%20of%20MSE%20on%20that%20mini-batch.%0A%20%20%20%20%20%20%20%20Use%20%60rng.choice(len(X)%2C%20size%3Dbatch_size%2C%20replace%3DFalse)%60%20to%20pick%20indices.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20sgd_gradient%22)%0A%0A%20%20%20%20%23%20---%20student%3A%20end%20---%0A%20%20%20%20return%20(sgd_gradient%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo%2C%20np%2C%20sgd_gradient)%3A%0A%20%20%20%20rng%20%3D%20np.random.default_rng(7)%0A%20%20%20%20X%20%3D%20rng.standard_normal((80%2C%203))%0A%20%20%20%20w_true%20%3D%20np.array(%5B1.0%2C%20-1.5%2C%202.0%5D)%0A%20%20%20%20y%20%3D%20X%20%40%20w_true%20%2B%200.1%20*%20rng.standard_normal(80)%0A%20%20%20%20w_sgd%20%3D%20np.zeros(3)%0A%0A%20%20%20%20_solved%20%3D%20False%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20g%20%3D%20sgd_gradient(X%2C%20y%2C%20w_sgd%2C%20batch_size%3D16%2C%20rng%3Dnp.random.default_rng(0))%0A%20%20%20%20%20%20%20%20assert%20g.shape%20%3D%3D%20(3%2C)%2C%20f%22gradient%20shape%3A%20expected%20(3%2C)%2C%20got%20%7Bg.shape%7D%22%0A%20%20%20%20%20%20%20%20assert%20not%20np.allclose(g%2C%200)%2C%20%22gradient%20is%20all%20zeros%20%E2%80%94%20check%20batch%20sampling%22%0A%0A%20%20%20%20%20%20%20%20%23%20Verify%20it%20matches%20manual%20mini-batch%20computation%0A%20%20%20%20%20%20%20%20rng_check%20%3D%20np.random.default_rng(99)%0A%20%20%20%20%20%20%20%20idx%20%3D%20rng_check.choice(len(X)%2C%20size%3D8%2C%20replace%3DFalse)%0A%20%20%20%20%20%20%20%20g_manual%20%3D%202%20%2F%208%20*%20X%5Bidx%5D.T%20%40%20(X%5Bidx%5D%20%40%20w_sgd%20-%20y%5Bidx%5D)%0A%20%20%20%20%20%20%20%20g_auto%20%20%20%3D%20sgd_gradient(X%2C%20y%2C%20w_sgd%2C%20batch_size%3D8%2C%20rng%3Dnp.random.default_rng(99))%0A%20%20%20%20%20%20%20%20assert%20np.allclose(g_auto%2C%20g_manual%2C%20rtol%3D1e-5)%2C%20%5C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22gradient%20does%20not%20match%20manual%20mini-batch%20computation%22%0A%0A%20%20%20%20%20%20%20%20%23%20Run%20SGD%20and%20check%20convergence%0A%20%20%20%20%20%20%20%20w_cur%20%3D%20np.zeros(3)%0A%20%20%20%20%20%20%20%20for%20k%20in%20range(2000)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20eta%20%3D%200.3%20*%20(1%20%2F%20(1%20%2B%20k))%20**%200.5%0A%20%20%20%20%20%20%20%20%20%20%20%20g_k%20%3D%20sgd_gradient(X%2C%20y%2C%20w_cur%2C%20batch_size%3D16%2C%20rng%3Drng)%0A%20%20%20%20%20%20%20%20%20%20%20%20w_cur%20-%3D%20eta%20*%20g_k%0A%20%20%20%20%20%20%20%20loss%20%3D%20float(np.mean((X%20%40%20w_cur%20-%20y)%20**%202))%0A%20%20%20%20%20%20%20%20assert%20loss%20%3C%200.05%2C%20f%22SGD%20did%20not%20converge%3A%20MSE%3D%7Bloss%3A.4f%7D%22%0A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%E2%9C%85%20SGD%20correct!%20After%202000%20steps%3A%20MSE%20%3D%20%7Bloss%3A.5f%7D%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22w%20%E2%89%88%20%7Bnp.round(w_cur%2C%203)%7D%22%0A%20%20%20%20%20%20%20%20)%2C%20kind%3D%22success%22)%0A%20%20%20%20%20%20%20%20_solved%20%3D%20True%0A%20%20%20%20except%20NotImplementedError%20as%20e%3A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(f%22%E2%9C%8F%EF%B8%8F%20%7Be%7D%22)%2C%20kind%3D%22neutral%22)%0A%20%20%20%20except%20Exception%20as%20e%3A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(f%22%E2%9D%8C%20%7Be%7D%22)%2C%20kind%3D%22danger%22)%0A%20%20%20%20%23%20---%20capture%20%26%20report%20(runs%20every%20time%2C%20pass%20or%20fail)%20---%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20import%20inspect%20as%20_inspect%2C%20json%20as%20_json%2C%20js%20as%20_js%0A%20%20%20%20%20%20%20%20from%20pyodide.ffi%20import%20to_js%20as%20_to_js%0A%20%20%20%20%20%20%20%20_srcs%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20_fn%20in%20(sgd_gradient%2C)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_srcs.append(_inspect.getsource(_fn))%0A%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pass%0A%20%20%20%20%20%20%20%20_payload%20%3D%20_json.dumps(%7B%22notebookId%22%3A%20%22gradient_descent%2Fstochastic_gd%22%2C%20%22code%22%3A%20%22%5Cn%5Cn%22.join(_srcs)%2C%20%22solved%22%3A%20bool(_solved)%7D)%0A%20%20%20%20%20%20%20%20if%20_solved%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch%20%3D%20_js.BroadcastChannel.new(%22mlp-notebooks%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch.postMessage(_js.JSON.parse(_json.dumps(%7B%22type%22%3A%20%22mlp%3Anotebook-solved%22%2C%20%22notebookId%22%3A%20%22gradient_descent%2Fstochastic_gd%22%7D)))%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch.close()%0A%20%20%20%20%20%20%20%20_js.fetch(%22%2Fapi%2Fnotebook-progress%22%2C%20_to_js(%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%22method%22%3A%20%22POST%22%2C%20%22headers%22%3A%20%7B%22Content-Type%22%3A%20%22application%2Fjson%22%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%22credentials%22%3A%20%22include%22%2C%20%22body%22%3A%20_payload%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20dict_converter%3D_js.Object.fromEntries))%0A%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20pass%20%20%23%20not%20running%20in%20Pyodide%20WASM%0A%20%20%20%20mo.output.replace(_result)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A