3c:null 3d:["$","title",null,{"children":"Understanding diffusion models in 2-D | Gabe Hope | Observable"}] 43:T5b5,dataset = { class dataset { static defaults = { batchsize: 16, fill: 'pink', stroke: 'black', r: 3, strokeWidth: 1, }; constructor(data, options) { const overwrites = options || {}; this.options = {...dataset.defaults, ...overwrites}; this.data = data; this.element = document.createElement("div"); this.element.value = this; } batch(batchsize) { const bs = batchsize || this.options.batchsize; return tf.tidy(() => { const inds = tf.randomUniformInt([bs], 0, this.data.shape[0]); return tf.gather(this.data, inds); }); } add(point) { if (this.data.shape[0] < 250){ const oldData = this.data; const newData = tf.tensor(point); this.data = tf.tidy(() => tf.concat([oldData, newData.reshape([1, oldData.shape[1]])])); oldData.dispose(); newData.dispose(); this.element.dispatchEvent(new CustomEvent("input")); } } marks(options) { return [Plot.dot(this.data.arraySync(), {r: this.options.r, fill: this.options.fill, stroke: this.options.stroke, strokeWidth: this.options.strokeWidth, ...options})] } tensor() { return this.data; } array() { return this.data.arraySync(); } watch() { return this.element; } } function create(data, options) { return new dataset(data, options); } return create; }44:T44a,estGradPointer = { if (!showCheckboxes.includes("Estimated gradient")) { return; } networkObs; const svg = d3.select(mainPlot.svg()); svg.selectAll('#estgradpointer').remove(); if (isFinite(coords[0]) && isFinite(coords[1])){ const coordst = tf.tensor([coords], undefined, 'float32'); const [dev, scales, weights] = tf.tidy(() => scheduler.schedule(sigma)); const gradtf = tf.tidy(() => scheduler.predictMean(coordst, network, dev, dev.mul(0.).add(sigma))); const grad = gradtf.arraySync()[0]; coordst.dispose(); gradtf.dispose(); scales.dispose(); weights.dispose(); dev.dispose(); const scalex = mainPlot.scalex; const scaley = mainPlot.scaley; const normGrad = Math.sqrt(grad[0] * grad[0] + grad[1] * grad[1]); svg.append('line').attr('id', 'estgradpointer') .attr("x1", scalex.apply(coords[0])) .attr("y1", scaley.apply(coords[1])) .attr("x2", scalex.apply(grad[0])) .attr("y2", scaley.apply(grad[1])) .attr('stroke', uiOptions.estgrad) .attr('stroke-width', 2); return grad; } }45:T55ee,makeSampler = { class BaseSampler { static defaults = { num_steps: 400, epsilon_s: 1e-3, M: 1000, seed: -1, latents: 'none', batchsize: 500, trace_color: "#000000", }; constructor(net, scheduler, options) { const overwrites = options || {}; this.options = {...BaseSampler.defaults, ...overwrites}; this.seed = this.options.seed; this.net = net; this.solver = this.options.solver; var epsilon_s = this.options.epsilon_s this.finished = false; this.scheduler = scheduler; this.num_steps = this.scheduler.options.num_steps; this.t_steps = tf.tidy(() => tf.linspace(1., epsilon_s, this.num_steps).arraySync()); this.t_steps.push(0); // Main sampling loop. this.reset(false); this.element = document.createElement("div"); this.element.value = this; } sigma(t) { return this.scheduler.sigma(t); } sigma_deriv(t) { return this.scheduler.sigma_deriv(t); } sigma_inv(t){ return this.scheduler.sigma_inv(t); } s(t) { return this.scheduler.s(t); } s_deriv(t) { return this.scheduler.s_deriv; } reset(do_dispose=true) { if (do_dispose){ this.dispose(); } const latentSeed = (this.options.seed >= 0 ? this.options.seed : undefined); const shape = [this.options.batchsize, 2]; this.t_next = this.t_steps[0]; this.x_next = tf.tidy(() => tf.randomNormal(shape, 0, 1, 'float32', latentSeed).mul(this.sigma(this.t_next))); this.i = 0; this.trace = [this.x_next.arraySync()[0]]; this.finished = false; } step(obs) { const output = tf.tidy(() => this.step_internal(obs)); this.update(); return output; } scatterMark(options){ return Plot.dot(this.x_next.arraySync(), {x: "0", y: "1", stroke: "white", r: 1, ...options}); } getTrace(){ return {color: this.options.trace_color, trace: this.trace}; } marks(){ return [Plot.dot(this.x_next.arraySync(), {x: "0", y: "1", stroke: "white", r: 1})]; } inpaintMarks() { return this.marks(); } dispose() { this.x_next.dispose(); } update() { this.element.dispatchEvent(new CustomEvent("input")); } watch() { return this.element; } } class AncestralSampler extends BaseSampler { step_internal(obs) { if (this.finished){ return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.i]; let t_next = this.t_steps[this.i+1]; this.t_next = t_next; this.x_next = this.scheduler.p_reverse_sample(this.net, x_cur, t_cur, t_next); this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); return this.x_next; } } class BlendedSampler extends BaseSampler { step_internal(obs) { if (this.finished){ return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.i]; let t_next = this.t_steps[this.i+1]; this.t_next = t_next; this.x_next = this.scheduler.p_reverse_sample(this.net, x_cur, t_cur, t_next); if (obs) { const tfobs = tf.tensor([obs]); const obs_sample = this.scheduler.q_forward_sample(tfobs.mul(tf.onesLike(this.x_next)), 0., t_next); this.x_next = tf.where(tf.isFinite(tfobs), obs_sample, this.x_next); } this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); return this.x_next; } } class EulerODESampler extends BaseSampler { step_internal(obs) { if (this.finished){ return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.i]; let t_next = this.t_steps[this.i+1]; const h = t_next - t_cur; this.t_next = t_next; this.x_next = x_cur.add(this.scheduler.predictDxDt(x_cur, this.net, this.scheduler.sigma(t_cur), t_cur).mul(h)); this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); x_cur.dispose(); return this.x_next; } } class HuenODESampler extends BaseSampler { step_internal(obs) { if (this.finished){ return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.i]; let t_next = this.t_steps[this.i+1]; const h = t_next - t_cur; this.t_next = t_next; const dxdt = this.scheduler.predictDxDt(x_cur, this.net, this.scheduler.sigma(t_cur), t_cur); const x_inter = x_cur.add(dxdt.mul(h)); const correction = this.scheduler.predictDxDt(x_inter, this.net, this.scheduler.sigma(t_next), t_next); this.x_next = x_cur.add(dxdt.add(correction).mul(h * 0.5)); this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); x_cur.dispose(); return this.x_next; } } class ForwardSampler extends BaseSampler { reset(do_dispose=true) { if (do_dispose){ this.dispose(); } const latentSeed = (this.options.seed >= 0 ? this.options.seed : undefined); const shape = [this.options.batchsize, 2]; this.t_next = this.t_steps[this.t_steps.length - 1]; this.x_next = this.options.data.batch(this.options.batchsize); this.i = 0; this.trace = [this.x_next.arraySync()[0]]; this.finished = false; } step_internal(obs) { if (this.finished){ return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.t_steps.length - 1 - this.i]; let t_next = this.t_steps[this.t_steps.length - 2 - this.i]; this.t_next = t_next; this.x_next = this.scheduler.q_forward_sample(x_cur, t_cur, t_next); this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); x_cur.dispose(); return this.x_next; } } class ForwardODESampler extends BaseSampler { reset(do_dispose=true) { if (do_dispose){ this.dispose(); } const latentSeed = (this.options.seed >= 0 ? this.options.seed : undefined); const shape = [this.options.batchsize, 2]; this.t_next = this.t_steps[this.t_steps.length - 1]; if (typeof this.options.data.batch === 'function') { this.x_next = this.options.data.batch(this.options.batchsize); } else { this.x_next = tf.clone(this.options.data); } this.lptrace = tf.tidy(() => tf.zerosLike(this.x_next.sum(1))); this.lp = tf.tidy(() => tf.zerosLike(this.x_next.sum(1))); this.i = 0; this.trace = [this.x_next.arraySync()[0]]; this.finished = false; } dispose() { this.lptrace.dispose(); this.lp.dispose(); this.x_next.dispose(); } step_internal(obs) { if (this.finished){ console.log(this.lp.arraySync()) return this.x_next; } let x_cur = this.x_next; let t_cur = this.t_steps[this.t_steps.length - 1 - this.i]; let t_next = this.t_steps[this.t_steps.length - 2 - this.i]; this.t_next = t_next; const h = t_next - t_cur; const x1 = tf.tensor([[1., 0.]]).mul(tf.onesLike(x_cur)); const x2 = tf.tensor([[0., 1.]]).mul(tf.onesLike(x_cur)); const pgrad = tf.grad((x_in) => this.scheduler.predictDxDt(x_in, this.net, this.scheduler.sigma(t_cur), t_cur)); const oldtrace = this.lptrace; const update = pgrad(x_cur, x1).add(pgrad(x_cur, x2)).sum(1).mul(h); this.lptrace = tf.add(this.lptrace, update); oldtrace.dispose(); const oldLp = this.lp; this.lp = norm.logpdf(x_cur, tf.zerosLike(x_cur), tf.onesLike(x_cur).mul(this.scheduler.sigma(1.))).sum(1).add(this.lptrace); oldLp.dispose(); this.x_next = x_cur.add(this.scheduler.predictDxDt(x_cur, this.net, this.scheduler.sigma(t_cur), t_cur).mul(h)); this.i += 1; this.trace.push(this.x_next.arraySync()[0]); this.finished = this.i >= (this.t_steps.length - 1); x_cur.dispose(); return [this.x_next, this.lp, this.lptrace]; } } class QAVISampler extends BaseSampler { constructor(net, scheduler, options) { const new_options = {...options, epsilon_s: 0.1, num_steps: Math.min(options.num_steps, 20), batchsize: Math.min(options.batchsize, 50)}; super(net, scheduler, new_options); this.num_steps = this.options.num_steps; this.reset(false, true); } reset(do_dispose=true, reset_parameters=false) { if (do_dispose) { this.dispose(reset_parameters); } const latentSeed = (this.options.seed >= 0 ? this.options.seed : undefined); const shape = [this.options.batchsize, 2]; this.t_next = this.t_steps[0]; this.i = 0; // EDM paper multiplies by s(t), not clear to me why? this.x_next = tf.tidy(() => tf.randomNormal(shape, 0, 1, 'float32', latentSeed).mul(this.sigma(this.t_next) * this.s(this.t_next))); this.trace = [this.x_next.arraySync()[0]]; this.finished = false; if (!reset_parameters) { return; } this.parameters = []; this.all_parameters = []; for (let i = 0; i < this.num_steps; i++) { let lambda = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01).sub(1.)); let m = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01)); let q_sigma = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01)); this.parameters.push([lambda, m, q_sigma]); this.all_parameters.push(lambda); this.all_parameters.push(m); this.all_parameters.push(q_sigma); } this.optimizer = tf.train.adam(0.01); this.qavi_trace = [[this.x_next.arraySync()[0]]]; this.loss_trace = []; this.current_step = 0; this.updatePlot(); this.update(); } step_internal(obs_in) { const trace = []; var optimizing = Boolean(obs_in); if (!optimizing && this.finished){ return this.x_next; } const computeLoss = () => { var loss = 0; var i = 0; var obs = obs_in; const t_0 = this.t_steps[i]; var x_cur = tf.randomNormal(this.x_next.shape, 0, 1, 'float32').mul(this.sigma(t_0) * this.s(t_0)); var p_mu = tf.zerosLike(x_cur); var p_sigma = tf.onesLike(x_cur).mul(this.sigma(t_0) * this.s(t_0)); if (!optimizing) { i = this.i; obs = tf.tensor([[0, 0]]); x_cur = this.x_next; } else { obs = tf.tensor([obs]); } trace.push(x_cur.arraySync()); while (i < (this.num_steps)) { const t_cur = this.t_steps[i]; const t_next = this.t_steps[i+1]; i++; if (i == this.num_steps){ var [p_mu, p_sigma] = this.scheduler.p_reverse(this.net, x_cur, t_cur, t_next); const lik = tf.where(obs.isFinite(), norm.logpdf(tf.where(obs.isFinite(), obs, tf.zerosLike(obs)) , p_mu, p_sigma.add(0.01)), tf.zerosLike(obs)); loss = lik.mean(0).sum().neg().add(loss); x_cur = tf.randomNormal(x_cur.shape, 0, 1, 'float32').mul(p_sigma).add(p_mu); x_cur = tf.where(tf.isFinite(obs), obs, x_cur); } else { var pred_x0 = this.scheduler.predictX0(x_cur, this.net, null, t_cur); const [lambda_logit, m, q_sigma_logit] = this.parameters[i]; const lambda = tf.sigmoid(lambda_logit); const q_x0 = pred_x0.mul(lambda.neg().add(1.)).add(m.mul(lambda)); const q_mu = this.scheduler.q_reverse(x_cur, q_x0, t_cur, t_next)[0]; const q_sigma = tf.softplus(q_sigma_logit); var [p_mu, p_sigma] = this.scheduler.q_reverse(x_cur, pred_x0, t_cur, t_next); //const q_mu = p_mu.mul(lambda.neg().add(1.)).add(m.mul(lambda)); x_cur = tf.randomNormal(x_cur.shape, 0, 1, 'float32').mul(q_sigma).add(q_mu); loss = norm.kl_divergence(q_mu, q_sigma, p_mu, p_sigma).mean(0).sum().add(loss); } trace.push(x_cur.arraySync()); if (!optimizing){ break; } } return loss; }; if (optimizing){ const loss = this.optimizer.minimize(computeLoss, true, this.all_parameters); this.loss_trace.push([this.current_step, loss.arraySync()]); if (this.loss_trace.length > 250){ this.loss_trace.shift(); } this.current_step = this.current_step + 1; this.updatePlot(); this.qavi_trace = trace; } else { computeLoss(); this.i += 1; this.finished = this.i >= (this.t_steps.length - 1); } this.trace.push(trace[trace.length - 1]); this.x_next.dispose(); this.x_next = tf.tensor(trace[trace.length - 1]); return this.x_next; } updatePlot() { const lossplot = Plot.plot({ height: 150, y: {label: "Loss"}, x: {label: "Step"}, marks: [ Plot.ruleY([0]), Plot.lineY(this.loss_trace, Plot.windowY({x: "0", y: "1", k: 25, reduce: "mean", stroke: "firebrick"})), Plot.lineY(this.loss_trace, {x: "0", y: "1", stroke: "steelblue"}), ] }); this.element = lossplot; } inpaintMarks() { var marks = this.marks(); const points = []; for (let i=0; i < this.parameters.length; i++) { const pxy = this.parameters[i][1].arraySync()[0]; const w = tf.tidy(() => tf.sigmoid(this.parameters[i][0]).mean()).arraySync(); points.push([...pxy, 5 * w, d3.interpolatePlasma(i / this.parameters.length), i / this.parameters.length]); } marks.push(Plot.lineY(points, {x: "0", y: "1", z: null, strokeWidth: 2, stroke: ((x)=> x[3])})); marks.push(Plot.dot(points, {x: "0", y: "1", r: "2", fill: "3", stroke: "3"})); return marks; } step(obs) { tf.tidy(() => this.step_internal(obs)); this.update(); return this.x_next; } dispose(dispose_parameters=true) { this.x_next.dispose(); if (dispose_parameters) { this.all_parameters.forEach((pi) => { pi.dispose(); }); this.optimizer.dispose(); } } } class QAVISamplerCont extends BaseSampler { constructor(net, scheduler, options) { const new_options = {...options, epsilon_s: 0.01, num_steps: Math.min(options.num_steps, 20), batchsize: Math.min(options.batchsize, 50)}; super(net, scheduler, new_options); this.num_steps = this.options.num_steps; this.reset(false, true); } reset(do_dispose=true, reset_parameters=false) { if (do_dispose) { this.dispose(reset_parameters); } const latentSeed = (this.options.seed >= 0 ? this.options.seed : undefined); const shape = [this.options.batchsize, 2]; this.t_next = this.t_steps[0]; this.i = 0; // EDM paper multiplies by s(t), not clear to me why? this.x_next = tf.tidy(() => tf.randomNormal(shape, 0, 1, 'float32', latentSeed).mul(this.sigma(this.t_next) * this.s(this.t_next))); this.trace = [this.x_next.arraySync()[0]]; this.finished = false; if (!reset_parameters) { return; } this.parameters = []; this.all_parameters = []; for (let i = 0; i < this.num_steps; i++) { let lambda = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01).sub(1.)); let m = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01)); let q_sigma = tf.variable(tf.randomNormal([1, this.x_next.shape[1]], 0, 0.01)); this.parameters.push([lambda, m, q_sigma]); this.all_parameters.push(lambda); this.all_parameters.push(m); this.all_parameters.push(q_sigma); } this.optimizer = tf.train.adam(0.01); this.qavi_trace = [[this.x_next.arraySync()[0]]]; this.loss_trace = []; this.current_step = 0; this.updatePlot(); this.update(); } step_internal(obs_in) { const trace = []; var optimizing = Boolean(obs_in); if (!optimizing && this.finished){ return this.x_next; } const computeLoss = () => { var loss = 0; var i = 0; var obs = obs_in; const t_0 = this.t_steps[i]; var x_cur = tf.randomNormal(this.x_next.shape, 0, 1, 'float32').mul(this.sigma(t_0) * this.s(t_0)); var p_mu = tf.zerosLike(x_cur); var p_sigma = tf.onesLike(x_cur).mul(this.sigma(t_0) * this.s(t_0)); if (!optimizing) { i = this.i; obs = tf.tensor([[0, 0]]); x_cur = this.x_next; } else { obs = tf.tensor([obs]); } trace.push(x_cur.arraySync()); while (i < this.num_steps) { const t_cur = this.t_steps[i]; const t_next = this.t_steps[i+1]; i++; if (i == this.num_steps){ var [p_mu, p_sigma] = this.scheduler.p_reverse(this.net, x_cur, t_cur, t_next) const lik = tf.where(obs.isFinite(), norm.logpdf(tf.where(obs.isFinite(), obs, tf.zerosLike(obs)) , p_mu, p_sigma.add(0.01)), tf.zerosLike(obs)); loss = lik.mean(0).sum().neg().add(loss); x_cur = tf.randomNormal(x_cur.shape, 0, 1, 'float32').mul(p_sigma).add(p_mu); x_cur = tf.where(tf.isFinite(obs), obs, x_cur); } else { var pred_x0 = this.scheduler.predictX0(x_cur, this.net, null, t_cur); const [lambda_logit, m, q_sigma_logit] = this.parameters[i]; const lambda = tf.sigmoid(lambda_logit); const q_x0 = pred_x0.mul(lambda.neg().add(1.)).add(m.mul(lambda)); const t_mid = tf.randomUniform([q_x0.shape[0], 1], t_next, t_cur); const x_mid = this.scheduler.q_reverse_sample(x_cur, q_x0, t_cur, t_mid); const p_x0 = this.scheduler.predictX0(x_mid, this.net, null, t_mid); var [sigma, scale, weights] = this.scheduler.schedule(t_mid.squeeze()); loss = p_x0.sub(q_x0).square().mul(weights.reshape([-1, 1])).mean(0).sum().add(loss); //loss = p_x0.sub(q_x0).square().mean(0).sum().add(loss); x_cur = this.scheduler.q_reverse_sample(x_mid, q_x0, t_mid, t_next); } trace.push(x_cur.arraySync()); if (!optimizing){ break; } } return loss; }; if (optimizing){ const loss = this.optimizer.minimize(computeLoss, true, this.all_parameters); this.loss_trace.push([this.current_step, loss.arraySync()]); if (this.loss_trace.length > 250){ this.loss_trace.shift(); } this.current_step = this.current_step + 1; this.updatePlot(); this.qavi_trace = trace; } else { computeLoss(); this.i += 1; this.finished = this.i >= (this.t_steps.length - 1); } this.trace.push(trace[trace.length - 1]); this.x_next.dispose(); this.x_next = tf.tensor(trace[trace.length - 1]); return this.x_next; } updatePlot() { const lossplot = Plot.plot({ height: 150, y: {label: "Loss"}, x: {label: "Step"}, marks: [ Plot.ruleY([0]), Plot.lineY(this.loss_trace, Plot.windowY({x: "0", y: "1", k: 25, reduce: "mean", stroke: "firebrick"})), Plot.lineY(this.loss_trace, {x: "0", y: "1", stroke: "steelblue"}), ] }); this.element = lossplot; } inpaintMarks() { var marks = this.marks(); const points = []; for (let i=0; i < this.parameters.length; i++) { const pxy = this.parameters[i][1].arraySync()[0]; const w = tf.tidy(() => tf.sigmoid(this.parameters[i][0]).mean()).arraySync(); points.push([...pxy, 5 * w, d3.interpolatePlasma(i / this.parameters.length), i / this.parameters.length]); } marks.push(Plot.lineY(points, {x: "0", y: "1", z: null, strokeWidth: 2, stroke: ((x)=> x[3])})); marks.push(Plot.dot(points, {x: "0", y: "1", r: "2", fill: "3", stroke: "3"})); return marks; } step(obs) { tf.tidy(() => this.step_internal(obs)); this.update(); return this.x_next; } dispose(dispose_parameters=true) { this.x_next.dispose(); if (dispose_parameters) { this.all_parameters.forEach((pi) => { pi.dispose(); }); this.optimizer.dispose(); } } } const samplerMap = { euler_ode: (n, s, o) => new EulerODESampler(n, s, {...o}), huen_ode: (n, s, o) => new HuenODESampler(n, s, {...o}), ancestral: (n, s, o) => new AncestralSampler(n, s, {...o}), blended: (n, s, o) => new BlendedSampler(n, s, {...o}), forward: (n, s, o) => new ForwardSampler(n, s, {...o}), forward_ode: (n, s, o) => new ForwardODESampler(n, s, {...o}), qavi: (n, s, o) => new QAVISampler(n, s, {...o}), qavi_cont: (n, s, o) => new QAVISamplerCont(n, s, {...o}), } var currentSampler = false; function create(net, scheduler, options) { if (currentSampler){ currentSampler.dispose(); } currentSampler = samplerMap[options['sampler']](net, scheduler, options); return currentSampler; } return create; }46:T4687,scheduler = { function allPairs(A, B, sigma) { const Asq = tf.square(A).sum(1, true); const Bsq = tf.square(B).sum(1, true).transpose(); const AB = tf.matMul(A, B, false, true).mul(-2.); const sigmasq = tf.square(sigma); const distsq = AB.add(Asq).add(Bsq).div(sigmasq).mul(-0.5).exp(); const constant = tf.pow(sigmasq.mul(2. * Math.PI), -0.5); return distsq.sub(constant).mean(1); } class BaseSchedule { static defaults = { num_steps: 1000, parameterization: "eps", noise_cond: "t", M: 1000, discrete_loss: false, loss: "simple", epsilon_t: 1e-5, }; // Constructor/destructor constructor(options) { const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...overwrites}; } dispose() {} // Scheduling functions, use mp math for generality // Minimally need to override variance schedule in subclass sigma(t) {} // Variance (really std. dev) schedule sigma_inv(s) {} s(t) { // Scaling schedule, default to VE schedule return mp.add(mp.mul(t, 0.), 1.); } weight(t) { // Weighting schedule (for continuous time model) return mp.mul(-0.5, this.snr_deriv(t)); } // Derivatives of scaling and variance schedules s_deriv(t) { return tf.tidy(() => { const tv = tf.variable(tf.clone(t)); const grad = tf.grad((x) => tf.sum(this.s(x)))(tv); tv.dispose(); return grad; }); } sigma_deriv(t) { return tf.tidy(() => { const tv = tf.variable(tf.clone(t)); const grad = tf.grad((x) => tf.sum(this.sigma(x)))(tv); tv.dispose(); return grad; }); } snr(t) { // Signal-to-noise return mp.square(mp.div(this.s(t), this.sigma(t))); } snr_deriv(t) { return tf.tidy(() => { const tv = tf.variable(tf.clone(t)); const grad = tf.grad((x) => tf.sum(this.snr(x)))(tv); tv.dispose(); return grad; }); } logsnr(t) { // Log Signal to noise return mp.mul(2., mp.sub(mp.log(this.s(t)), mp.log(this.sigma(t)))); } variance(t) { return mp.square(this.sigma(t)); } // ODE/SDE coefficients f(t) { return mp.div(this.s_deriv(t), this.s(t)); } g(t) { const t1 = mp.mul(this.sigma_deriv(t), this.sigma(t)); const t2 = mp.mul(mp.square(this.sigma(t)), this.f(t)); return mp.sqrt(mp.mul(2, mp.sub(t1, t2))); } schedule(t) { return [tf.clone(this.sigma(t)), tf.clone(this.s(t)), tf.clone(this.weight(t))]; } call(t) { // Deprecated name return this.schedule(t); } noise_cond(t) { if (this.options.noise_cond == 'log_snr') { const [nlsnr_min, nlsnr_max] = [mp.neg(this.logsnr(this.options.epsilon_t)), mp.neg(this.logsnr(1.))]; const nlsnr = mp.neg(this.logsnr(t)); // Rescale to [0, M] return mp.mul(mp.div(mp.sub(nlsnr, nlsnr_min), nlsnr_max), this.options.M); } else if (this.options.noise_cond == 't') { return mp.mul(t, this.options.M); } } // Prediction functions, use tf predictX0(x, net, sig, t, param) { const noise = tf.clone(this.noise_cond(t)).reshape([-1]); const scale = tf.clone(this.s(t)).reshape([-1, 1]); const sigma = tf.clone(this.sigma(t)).reshape([-1, 1]); const precond = 1.;// tf.sqrt(tf.square(sigma).add(1.)); const output = net.apply(tf.div(x, precond), noise); const parameterization = param || this.options.parameterization; if (parameterization == 'x0') { return output; } else if (parameterization == 'eps') { return tf.div(tf.sub(x, tf.mul(sigma, output)), scale); } else if (parameterization == 'score') { return tf.div(tf.add(x, tf.mul(tf.square(sigma), output)), scale); } } predictEps(x, net, sig, t, param) { const noise = tf.clone(this.noise_cond(t)).reshape([-1]); const scale = tf.clone(this.s(t)).reshape([-1, 1]); const sigma = tf.clone(this.sigma(t)).reshape([-1, 1]); const precond = 1.; //tf.sqrt(tf.square(sigma).add(1.)); const output = net.apply(tf.div(x, precond), noise); const parameterization = param || this.options.parameterization; if (parameterization == 'x0') { return tf.div(tf.sub(x, tf.mul(scale, output)), sigma); } else if (parameterization == 'eps') { return output; } else if (parameterization == 'score') { return tf.mul(tf.neg(output), sigma); } } predictScore(x, net, sig, t, param) { const noise = tf.clone(this.noise_cond(t)).reshape([-1]); const scale = tf.clone(this.s(t)).reshape([-1, 1]); const sigma = tf.clone(this.sigma(t)).reshape([-1, 1]); const precond = 1.;//tf.sqrt(tf.square(sigma).add(1.)); const output = net.apply(tf.div(x, precond), noise); const parameterization = param || this.options.parameterization; if (parameterization == 'x0') { return tf.div(tf.sub(tf.mul(scale, output), x), tf.square(sigma)); } else if (parameterization == 'eps') { return tf.div(tf.neg(output), sigma); } else if (parameterization == 'score') { return output; } } predictDxDt(x, net, sig, t, param) { const noise = tf.clone(this.noise_cond(t)).reshape([-1]); const scale = tf.clone(this.s(t)).reshape([-1, 1]); const sigma = tf.clone(this.sigma(t)).reshape([-1, 1]); const precond = 1.;//tf.sqrt(tf.square(sigma).add(1.)); const output = net.apply(tf.div(x, precond), noise); var h = mp.sub(this.sigma_deriv(t), mp.mul(this.sigma(t), this.f(t))); h = tf.clone(h).reshape([-1, 1]) const f = tf.clone(this.f(t)).reshape([-1, 1]).mul(x); const parameterization = param || this.options.parameterization; if (parameterization == 'x0') { return f.sub(tf.div(tf.sub(tf.mul(scale, output), x), sigma).mul(h)); } else if (parameterization == 'eps') { return f.sub(tf.neg(output).mul(h)); } else if (parameterization == 'score') { return f.sub(output.mul(h).mul(sigma)); } } predictMean(x, net, sig, t, param) { const scale = tf.clone(this.s(t)).reshape([-1, 1]); const sigma = tf.clone(this.sigma(t)).reshape([-1, 1]); const parameterization = param || this.options.parameterization; if (parameterization == 'x0') { return this.predictX0(x, net, sig, t).mul(scale); } else if (parameterization == 'eps') { return tf.sub(x, tf.mul(sigma, this.predictEps(x, net, sig, t))); } else if (parameterization == 'score') { return this.predictX0(x, net, sig, t).mul(scale); } } predict(x, net, sigma, t) { // Deprecated name return this.predictX0(x, net, sigma, t); } // Conditional distributions, use tf q_reverse(x_cur, x0, t_cur, t_prev){ // Reverse step const alpha_t = this.s(t_cur); const alpha_s = this.s(t_prev); const alpha_ts = mp.div(alpha_t, alpha_s); const v_t = mp.square(this.sigma(t_cur)); const v_s = mp.square(this.sigma(t_prev)); const v_ts = mp.sub(v_t, mp.mul(v_s, mp.square(alpha_ts))); const v_Q = mp.mul(v_ts, mp.div(v_s, v_t)); const sigma_Q = mp.sqrt(v_Q); const c_cur = mp.mul(alpha_ts, mp.div(v_s, v_t)); const c_0 = mp.mul(alpha_s, mp.div(v_ts, v_t)); const q_mu = x_cur.mul(c_cur).add(x0.mul(c_0)); const q_sigma = tf.onesLike(x0).mul(sigma_Q); return [q_mu, q_sigma]; } q_reverse_sample(x_cur, x0, t_cur, t_prev, seed){ const [p_mu, p_sigma] = this.q_reverse(x_cur, x0, t_cur, t_prev); const sample = Boolean(seed) ? tf.randomNormal(x_cur.shape, 0., 1., 'float32', seed) : tf.randomNormal(x_cur.shape); return sample.mul(p_sigma).add(p_mu); } p_reverse(net, x_cur, t_cur, t_prev){ // Reverse step var x_hat = this.predictX0(x_cur, net, tf.ones([x_cur.shape[0]]).mul(this.sigma(t_cur)), tf.ones([x_cur.shape[0]]).mul(t_cur)); return this.q_reverse(x_cur, x_hat, t_cur, t_prev); } p_reverse_sample(net, x_cur, t_cur, t_prev, seed){ var x_hat = this.predictX0(x_cur, net, tf.ones([x_cur.shape[0]]).mul(this.sigma(t_cur)), tf.ones([x_cur.shape[0]]).mul(t_cur)); return this.q_reverse_sample(x_cur, x_hat, t_cur, t_prev, seed); } q_forward(x_cur, t_cur, t_next) { const alpha_t = this.s(t_next); const alpha_s = this.s(t_cur); const alpha_ts = mp.div(alpha_t, alpha_s); const v_t = mp.square(this.sigma(t_next)); const v_s = mp.square(this.sigma(t_cur)); const v_ts = mp.sub(v_t, mp.mul(v_s, mp.square(alpha_ts))); const sigma_ts = mp.sqrt(v_ts); const q_mu = x_cur.mul(alpha_ts); const q_sigma = tf.onesLike(x_cur).mul(sigma_ts); return [q_mu, q_sigma]; } q_forward_sample(x_cur, t_cur, t_next, seed){ // forward step const [q_mu, q_sigma] = this.q_forward(x_cur, t_cur, t_next); const sample = Boolean(seed) ? tf.randomNormal(x_cur.shape, 0., 1., 'float32', seed) : tf.randomNormal(x_cur.shape); return sample.mul(q_sigma).add(q_mu); } // Loss functions, use tf lossCont(x, net, t0=0.) { // Standard, continuous time loss const t = tf.randomUniform([x.shape[0]]); var [sigma, scale, weights] = this.schedule(t); const n = tf.randomNormal(x.shape).mul(sigma.reshape([-1, 1])); const x0_pred = this.predictX0(x.mul(scale.reshape([-1, 1])).add(n), net, sigma, t); const loss = x0_pred.sub(x).square().mul(weights.reshape([-1, 1])); const output = loss.sum(1).mean(); return output; } lossSimpleCont(x, net, t0=0.) { // Standard, continuous time loss const t = tf.randomUniform([x.shape[0]]); var [sigma, scale, weights] = this.schedule(t); const n = tf.randomNormal(x.shape); const eps_pred = this.predictEps(x.mul(scale.reshape([-1, 1])).add(n.mul(sigma.reshape([-1, 1]))), net, sigma, t); const loss = eps_pred.sub(n).square(); const output = loss.sum(1).mean(); return output; } lossDiscrete(x, net, t0=0.) { const t = tf.randomUniformInt([x.shape[0]], 1, this.options.num_steps).cast('float32'); const t_cur = t.div(this.num_steps - 1); const t_next = t.sub(1).div(this.num_steps - 1); var [sigma, scale, weights] = this.schedule(t); const n = tf.randomNormal(x.shape).mul(sigma.reshape([-1, 1])); const x0_pred = this.predictX0(x.mul(scale.reshape([-1, 1])).add(n), net, sigma, t_cur); weights = this.snr(t_next) - this.snr(t_cur); const loss = x0_pred.sub(x).square().mul(weights.reshape([-1, 1])); const output = loss.sum(1).mean(); return output; } lossSimpleDiscrete(x, net, t0=0.) { const t = tf.randomUniformInt([x.shape[0]], 1, this.options.num_steps).cast('float32').div(this.num_steps - 1); var [sigma, scale, weights] = this.schedule(t); const n = tf.randomNormal(x.shape); const eps_pred = this.predictEps(x.mul(scale.reshape([-1, 1])).add(n.mul(sigma.reshape([-1, 1]))), net, sigma, t); const loss = eps_pred.sub(n).square(); const output = loss.sum(1).mean(); return output; } loss(x, net) { if (this.options.discrete && this.options.loss == "vlb") { return this.lossDiscrete(x, net); } else if (this.options.discrete && this.options.loss == "simple") { return this.lossSimpleDiscrete(x, net); } else if (!this.options.discrete && this.options.loss == "vlb") { return this.lossCont(x, net); } else if (!this.options.discrete && this.options.loss == "simple") { return this.lossSimpleCont(x, net); } } // "Tidy" functions for visualization heatmapPlot(A, B, t) { return tf.tidy(() => { const [sigma, scale, weight] = this.schedule(t); return allPairs(A, B.mul(scale), tf.clone(sigma)); }); } linePlot(y){ const outputs = tf.tidy(() => { const t = tf.linspace(0., 1., 200); return tf.stack([t, this[y](t)], 1); }); const arr = outputs.arraySync(); outputs.dispose(); return arr; } trueGrad(p, x, t, astf=false) { const [sigma, scale, weights] = tf.tidy(() => this.schedule(t)); const outputs = tf.tidy(() => { const A = p; const B = x.mul(scale); const Asq = tf.square(A).sum(1, true); const Bsq = tf.square(B).sum(1, true).transpose(); const AB = tf.matMul(A, B, false, true).mul(-2.); const sigmasq = tf.square(sigma); const distsq = AB.add(Asq).add(Bsq).div(sigmasq).mul(-0.5).exp(); const z_prob = distsq.div(distsq.sum(1, true)); return z_prob.matMul(x.mul(scale)); }); sigma.dispose(); scale.dispose(); weights.dispose(); if (astf) { return outputs; } const arr = outputs.arraySync(); outputs.dispose(); return arr; } } class VPSchedule extends BaseSchedule { static defaults = { beta_max: 19.9, beta_min: 0.1, epsilon_t: 1e-3, }; constructor(options) { super(options); const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...VPSchedule.defaults, ...overwrites}; } sigma(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1.); const beta_min = this.options.beta_min; const beta_max = this.options.beta_max; const t1 = mp.mul(-0.5 * (beta_max - beta_min), mp.square(t)); const t2 = mp.mul(beta_min, t); return mp.sqrt(mp.neg(mp.expm1(mp.sub(t1, t2)))); } s(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1.); const beta_min = this.options.beta_min; const beta_max = this.options.beta_max; const t1 = mp.mul(-0.25 * (beta_max - beta_min), mp.square(t)); const t2 = mp.mul(0.5 * beta_min, t); return mp.exp(mp.sub(t1, t2)); } f(t_in) { const t = mp.clipByValue(t_in, this.options.epsilon_t, 1.); const beta_min = this.options.beta_min; const beta_max = this.options.beta_max; const out = mp.mul(-0.5, mp.add(beta_min, mp.mul(beta_max - beta_min, t))); return out; } } class SubVPSchedule extends VPSchedule { static defaults = { beta_max: 19.9, beta_min: 0.1, epsilon_t: 1e-3, }; constructor(options) { super(options); const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...SubVPSchedule.defaults, ...overwrites}; } sigma(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1.); const beta_min = this.options.beta_min; const beta_max = this.options.beta_max; const t1 = mp.mul(-0.5 * (beta_max - beta_min), mp.square(t)); const t2 = mp.mul(beta_min, t); return mp.neg(mp.expm1(mp.sub(t1, t2))); } } class VESchedule extends BaseSchedule { static defaults = { sigma_max: 19.9, sigma_min: 0.1, epsilon_t: 1e-3, }; constructor(options) { super(options); const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...VESchedule.defaults, ...overwrites}; } sigma(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1.); const sigma_min = this.options.sigma_min; const sigma_max = this.options.sigma_max; const t1 = mp.pow(mp.div(sigma_max, sigma_min), mp.mul(2., t)); return mp.sqrt(mp.mul(mp.square(sigma_min), t1)); } } class CosineSchedule extends BaseSchedule { static defaults = { epsilon_t: 1e-3, }; constructor(options) { super(options); const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...VPSchedule.defaults, ...overwrites}; } helper(t) { const epsilon_t = this.options.epsilon_t; return mp.cos(mp.mul(mp.div(mp.add(t, epsilon_t), 1 + epsilon_t), Math.PI / 2.)); } sigma(t_in){ const t = mp.mul(t_in, 1. - this.options.epsilon_t); return mp.sqrt(mp.sub(1., mp.square(this.s(t)))); } s(t_in){ const t = mp.mul(t_in, 1. - this.options.epsilon_t); return mp.div(this.helper(t), this.helper(0)); } } class LinearSchedule extends BaseSchedule { static defaults = { epsilon_t: 1e-3, }; constructor(options) { super(options); const overwrites = options || {}; this.options = {...BaseSchedule.defaults, ...LinearSchedule.defaults, ...overwrites}; } sigma(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1. - this.options.epsilon_t); return t; } s(t_in){ const t = mp.clipByValue(t_in, this.options.epsilon_t, 1. - this.options.epsilon_t); return mp.sub(1., t); } } const vp = new VPSchedule({num_steps: samplerOptions.num_steps, discrete_loss: lossType == 'Discrete', loss: lossWeighting, parameterization: net_parameterization}); const subvp = new SubVPSchedule({num_steps: samplerOptions.num_steps, discrete_loss: lossType == 'Discrete', loss: lossWeighting, parameterization: net_parameterization}); const ve = new VESchedule({num_steps: samplerOptions.num_steps, discrete_loss: lossType == 'Discrete', loss: lossWeighting, parameterization: net_parameterization}); const cosine = new CosineSchedule({num_steps: samplerOptions.num_steps, discrete_loss: lossType == 'Discrete', loss: lossWeighting, parameterization: net_parameterization}); const linear = new LinearSchedule({num_steps: samplerOptions.num_steps, discrete_loss: lossType == 'Discrete', loss: lossWeighting, parameterization: net_parameterization}); const schedules = {VP: vp, VE: ve, SubVP: subvp, Cosine: cosine, Linear: linear}; const selectedScheduler = schedules[select]; return selectedScheduler; } 47:T55e,mp = { // Functions to work on either tf Tensors or javascript numbers // Avoid converting to fp32 and creating tensors that need to be disposed function makeOp(numOp, tfOp) { function op(...args) { if (args.reduce((c, n) => c || (n instanceof tf.Tensor), false)){ return tfOp(...args); } return numOp(...args); } return op; } return { add: makeOp((a, b) => a + b, tf.add), sub: makeOp((a, b) => a - b, tf.sub), mul: makeOp((a, b) => a * b, tf.mul), div: makeOp((a, b) => a / b, tf.div), neg: makeOp((a) => -a, tf.neg), maximum: makeOp(Math.max, tf.maximum), minimum: makeOp(Math.min, tf.minimum), exp: makeOp(Math.exp, tf.exp), log: makeOp(Math.log, tf.log), expm1: makeOp(Math.expm1, tf.expm1), log1p: makeOp(Math.log1p, tf.log1p), pow: makeOp(Math.pow, tf.pow), sin: makeOp(Math.sin, tf.sin), cos: makeOp(Math.cos, tf.cos), abs: makeOp(Math.abs, tf.abs), square: makeOp((a) => a * a, tf.square), sqrt: makeOp(Math.sqrt, tf.sqrt), floor: makeOp(Math.floor, tf.floor), ceil: makeOp(Math.ceil, tf.ceil), reciprocal: makeOp((a) => 1. / a, tf.reciprocal), tanh: makeOp(Math.tanh, tf.tanh), sigmoid: makeOp((a) => Math.tanh(a) / 2., tf.tanh), clipByValue: makeOp((a, min, max) => Math.min(Math.max(a, min), max), tf.clipByValue), } }48:T5ae,makeTrainStep = { const optimizers = {SGD: tf.train.sgd, Adam: tf.train.adam, RMSProp: tf.train.rmsprop, Momentum: tf.train.momentum } var optimizer = tf.train.sgd(); var currentstep = 0; var lossTrace = []; function makeTrainStep(network, rate, opt, scheduler) { optimizer.dispose(); optimizer = optimizers[opt](rate); currentstep = 0; lossTrace = []; function trainStepHelper(x) { return optimizer.minimize(() => scheduler.loss(x, network), true); } function trainStep(x, update=true) { var loss = 0; if (update) { const losstf = tf.tidy(() => trainStepHelper(x)); loss = losstf.arraySync(); currentstep += 1; losstf.dispose(); lossTrace.push([currentstep, loss]); if (lossTrace.length > 250){ lossTrace.shift(); } network.update(); } const lossplot = Plot.plot({ height: 150, y: {label: "Loss"}, x: {label: "Step"}, marks: [ Plot.ruleY([0]), Plot.lineY(lossTrace, Plot.windowY({x: "0", y: "1", k: 25, reduce: "mean", stroke: "firebrick"})), Plot.lineY(lossTrace, {x: "0", y: "1", stroke: "steelblue"}), ] }); lossplot.value = loss; return lossplot; } return trainStep; } return makeTrainStep; }49:T8b4,getDataset = { function makeMoons(n){ return tf.tidy(() => { const outer_circ_x = tf.cos(tf.linspace(0., Math.PI, n)); const outer_circ_y = tf.sin(tf.linspace(0., Math.PI, n)); const inner_circ_x = tf.cos(tf.linspace(0., Math.PI, n)).neg().add(1.); const inner_circ_y = tf.sin(tf.linspace(0., Math.PI, n)).neg().add(1.).sub(0.5); const X = tf.stack( [tf.concat([outer_circ_x, inner_circ_x]).sub(0.5), tf.concat([outer_circ_y, inner_circ_y])], 1 ).mul(4.); return X; }); } function makeSpiral(n){ return tf.tidy(() => { const r = tf.linspace(0.5, 6, n); const circ_x = tf.cos(tf.linspace(0.01, 1., n).pow(0.5).mul(6 * Math.PI)).mul(r); const circ_y = tf.sin(tf.linspace(0.01, 1., n).pow(0.5).mul(6 * Math.PI)).mul(r); const X = tf.stack( [circ_x, circ_y], 1 ); return X; }); } function makeCircles(n){ return tf.tidy(() => { const outer_circ_x = tf.cos(tf.linspace(2 * Math.PI / n, 2 * Math.PI, n)).mul(3.); const outer_circ_y = tf.sin(tf.linspace(2 * Math.PI / n, 2 * Math.PI, n)).mul(3.); const inner_circ_x = tf.cos(tf.linspace(2 * Math.PI / n, 2 * Math.PI, n)).mul(5.); const inner_circ_y = tf.sin(tf.linspace(2 * Math.PI / n, 2 * Math.PI, n)).mul(5.); const X = tf.stack( [tf.concat([outer_circ_x, inner_circ_x]), tf.concat([outer_circ_y, inner_circ_y])], 1 ); return X; }); } function makeBlobs(n){ return tf.tidy(() => { const q1 = tf.randomStandardNormal([Math.floor(n / 4) + 1, 2]).add(tf.tensor([[3., 3.]])); const q2 = tf.randomStandardNormal([Math.floor(n / 4) + 1, 2]).add(tf.tensor([[3., -3.]])); const q3 = tf.randomStandardNormal([Math.floor(n / 4) + 1, 2]).add(tf.tensor([[-3., 3.]])); const q4 = tf.randomStandardNormal([Math.floor(n / 4) + 1, 2]).add(tf.tensor([[-3., -3.]])); const X = tf.concat([q1, q2, q3, q4]); return X; }); } const datasets = {Moons: makeMoons, Spiral: makeSpiral, Circles: makeCircles, Blobs: makeBlobs}; function getDataset(dataset="Moons", n=50) { return datasets[dataset](n); } return getDataset; }4a:T1848,makeNetwork = { class layerbase { apply(x) { var output = x; for (let i = 0; i < this.all_layers.length; i++){ output = this.all_layers[i].apply(output); } return output; } dispose() { for (let i = 0; i < this.all_layers.length; i++){ this.all_layers[i].dispose(); } } } class dense extends layerbase { static defaults = {}; constructor(in_dims, out_dims, options) { super(); const overwrites = options || {}; this.options = {...dense.defaults, ...overwrites}; this.all_layers = []; this.in_dims = in_dims; this.out_dims = out_dims; if (this.options.normalization == "batch") { this.normlayer = tf.layers.batchNormalization(); this.all_layers.push(this.normlayer); } else if (this.options.normalization == "layer") { this.normlayer = tf.layers.layerNormalization(); this.all_layers.push(this.normlayer); } if (this.options.dropout) { this.dropoutlayer = tf.layers.gaussianDropout(); this.all_layers.push(this.dropoutlayer); } this.layer = tf.layers.dense({units: out_dims, activation: this.options.activation}); this.all_layers.push(this.layer); } } class residual extends layerbase { static defaults = { activation: 'relu', dropout: false, normalization: false, }; constructor(in_dims, layers, options) { super(); const overwrites = options || {}; this.options = {...dense.defaults, ...overwrites}; this.all_layers = []; this.in_dims = in_dims; this.out_dims = layers[layers.length-1]; var dims = in_dims; for (let i = 0; i < layers.length; i++){ this.all_layers.push(new dense(dims, layers[i], this.options)); dims = layers[i]; } } } class embedding extends layerbase { static defaults = { max_positions: 10000., }; constructor(in_dims, embedding_dim, options) { super(); const overwrites = options || {}; this.options = {...embedding.defaults, ...overwrites}; this.in_dims = in_dims; this.out_dims = embedding_dim; this.embedding_dim = embedding_dim; this.inputLayer = tf.layers.dense({units: embedding_dim}); this.addLayer = tf.layers.add(); this.all_layers = [this.inputLayer, this.addLayer]; } apply(x, t) { var freqs = tf.range(0, this.embedding_dim / 2); freqs = freqs.div(this.embedding_dim / 2); freqs = tf.pow(1. / this.options.max_positions, freqs); var e = tf.outerProduct(t.reshape([-1]), freqs); e = tf.concat([e.cos(), e.sin()], 1); var output = x; output = this.inputLayer.apply(x); output = this.addLayer.apply([output, e]); return output; } } class embedding2d extends layerbase { static defaults = { max_positions: 10000., xmin: -300., xmax: 300. }; constructor(in_dims, embedding_dim, options) { super(); const overwrites = options || {}; this.options = {...embedding2d.defaults, ...overwrites}; this.in_dims = in_dims; this.out_dims = embedding_dim; this.embedding_dim = embedding_dim; this.all_layers = []; } apply(x, t) { var freqs = tf.range(0, this.embedding_dim / 2); freqs = freqs.div(this.embedding_dim / 2); freqs = tf.pow(1. / this.options.max_positions, freqs); var e = tf.outerProduct(t.reshape([-1]), freqs); const [x1, x2] = tf.unstack(x.sub(this.options.xmin).div(this.options.xmax).mul(1000.), 1); const ex1 = tf.outerProduct(x1.reshape([-1]), freqs); const ex2 = tf.outerProduct(x2.reshape([-1]), freqs); e = e.mul(tf.onesLike(ex1)); e = tf.concat([e.cos(), e.sin(), ex1.cos(), ex1.sin(), ex2.cos(), ex2.sin()], 1); return e; } } class network extends layerbase { static defaults = { hidden_layers: [[256, 256], [256, 256]], output_dim: 2, input_dim: 2, activation: 'elu', embedding_dim: 256, embedding_2d: true, dropout: false, normalization: 'none', }; constructor(options) { super(); const overwrites = options || {}; this.options = {...network.defaults, ...overwrites}; this.all_layers = []; var dims = this.options.input_dim; if (this.options.embedding_dim && this.options.embedding_2d) { this.embedding_layer = new embedding2d(this.options.input_dim, this.options.embedding_dim); this.all_layers.push(this.embedding_layer); dims = 3 * this.options.embedding_dim; } else if (this.options.embedding_dim) { this.embedding_layer = new embedding(this.options.input_dim, this.options.embedding_dim); this.all_layers.push(this.embedding_layer); dims = this.options.embedding_dim; } for (let i = 0; i < this.options.hidden_layers.length; i++){ let layer = this.options.hidden_layers[i]; if(Array.isArray(layer)){ layer = new residual(dims, layer, this.options); } else { layer = new dense(dims, layer, this.options); } this.all_layers.push(layer); dims = layer.out_dims; } this.all_layers.push(new dense(dims, this.options.output_dim)) this.element = document.createElement("div"); this.element.value = this; // Make sure layers get initialized tf.tidy(() => this.apply(tf.ones([1, this.options.input_dim]), tf.ones([1]))).dispose(); } apply(x, t) { var output = x; for (let i = 0; i < this.all_layers.length; i++){ output = this.all_layers[i].apply(output, t); } return output; } watch() { return this.element; } update() { this.element.dispatchEvent(new CustomEvent("input")); } } var currentNetwork = new network(); function create(options) { currentNetwork.dispose(); currentNetwork = new network(options); return currentNetwork; } return create; }4b:T97f,marks = { var output = []; if (showCheckboxes.includes("Data")) { output = [...(data.marks())]; } if (sampleView == "Samples") { const sample_options = {r: uiOptions.sample_r, fill: uiOptions.sample_fill, stroke_width: uiOptions.sample_strokeWidth, stroke: uiOptions.sample_stroke} output.push(sample_loop.scatterMark(sample_options)); } else if (sampleView == "Traces"){ for (let i = 0; i < allTraces.length; i++) { output.push(Plot.lineY(allTraces[i].trace, {x: "0", y: "1", strokeWidth: 2, stroke: allTraces[i].color})); } } else if (sampleView == "Inpaint") { output = [...output, ...sample_loop.inpaintMarks(), Plot.ruleX([xobserved], {stroke: "firebrick", strokeWidth: 2}),] } networkObs; if (showCheckboxes.includes("True grid") || showCheckboxes.includes("Estimated grid")) { const grid = tf.tidy(() => { const xrange = tf.linspace(mainPlot.scalex.domain[0], mainPlot.scalex.domain[1], 20); const yrange = tf.linspace(mainPlot.scaley.domain[0], mainPlot.scaley.domain[1], 20); const grid = tf.meshgrid(xrange, yrange); return tf.stack([grid[0].reshape([-1]), grid[1].reshape([-1])], 1); }); if (showCheckboxes.includes("True grid")) { const grads = scheduler.trueGrad(grid, data.tensor(), sigma, true); const stacked = tf.tidy(() => { const n = tf.euclideanNorm(tf.sub(grads, grid), 1, 1); const s = tf.concat([grid, grads, n], 1); return s; }); grads.dispose(); output.push(Plot.arrow(stacked.arraySync(), {x1: "0", y1: "1", x2: "2", y2: "3", stroke: uiOptions.truegrad, strokeWidth: (d) => Number(d[4] < 10.), headLength: 3})); stacked.dispose(); } if (showCheckboxes.includes("Estimated grid")) { const [dev, scales, weights] = tf.tidy(() => scheduler.schedule(sigma)); const grads = tf.tidy(() => scheduler.predictMean(grid, network, dev, dev.mul(0.).add(sigma))); const stacked = tf.tidy(() => { const n = tf.euclideanNorm(tf.sub(grads, grid), 1, 1); const s = tf.concat([grid, grads, n], 1); return s; }); grads.dispose(); output.push(Plot.arrow(stacked.arraySync(), {x1: "0", y1: "1", x2: "2", y2: "3", strokeWidth: 1, stroke: uiOptions.estgrad, headLength: 3})); stacked.dispose(); scales.dispose(); weights.dispose(); dev.dispose(); } grid.dispose(); } return output; }4c:[] 41:{"id":"440ccfdd9a38b53d","slug":null,"title":"Understanding diffusion models in 2-D","thumbnail":"220a87dd987b70efbefedf1a3a55812a8f8d259a148a8ed2cf911773b19c2995","access_level":"public","create_time":"$D2024-06-02T20:55:07.853Z","update_time":"$D2026-09-10T21:39:35.920Z","view_time":null,"first_public_time":"$D2024-09-23T02:29:37.381Z","has_unpublished_edits":null,"fork":true,"stars":0,"starred":false,"owner":{"type":"team","id":"8770219c10fecf82","login":"gabehope","name":"Gabe Hope","avatar_url":"https://avatars-next.observableusercontent.com/f8f123418df8144a6c8b914408e78ff321a4ee014e356bcdeae201594eae30af","tier":"starter_2024","bio":""},"version":2552,"publish_version":null,"role":null,"authored":false,"authors":[{"type":"individual","id":"74942dd6d6baa727","login":"gabe-hope3","name":"Gabe Hope","avatar_url":"https://avatars-next.observableusercontent.com/d4aa0361a0af1ee7fd565ca01cfe4c6274d25c9d4ec3891b0e82e5307ad05c45","tier":null,"bio":""}],"members":[{"type":"team","id":"8770219c10fecf82","login":"gabehope","name":"Gabe Hope","avatar_url":"https://avatars-next.observableusercontent.com/f8f123418df8144a6c8b914408e78ff321a4ee014e356bcdeae201594eae30af","tier":"starter_2024","role":"editor"},{"type":"individual","id":"74942dd6d6baa727","login":"gabe-hope3","name":"Gabe Hope","avatar_url":"https://avatars-next.observableusercontent.com/d4aa0361a0af1ee7fd565ca01cfe4c6274d25c9d4ec3891b0e82e5307ad05c45","tier":null,"role":"editor"}],"comments":[],"custom_thumbnail":true,"trashed":false,"lapsed":null,"subscribed":null,"forks":4,"collection":{"id":"e46deae83dcca492","slug":"teaching-visualizations","title":"Teaching Visualizations","owner":{"id":"8770219c10fecf82","login":"gabehope","name":"Gabe Hope"}},"type":"notebook","body":{"titling":"auto","title":"Understanding diffusion models in 2-D","cells":[{"id":0,"value":"# Understanding diffusion models in 2-D","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2194,"value":"**Start here!** \n\nClick *start* to train a basic diffusion model. Let training run for at least 1000 steps, then click *stop* to pause training. You'll be able to come back and train new models later.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":1027,"value":"viewof trainCont = Inputs.button([\n [\"Start\", () => 1],\n [\"Stop\", () => 0],\n [\"Step\", () => 2],\n], {label: \"Train\", value: 0, disabled: (sampleView != \"None\")})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":747,"value":"viewof resetNet = Inputs.button(\"Reset Network\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":445,"value":"viewof lossvalue = {\n for (let iter = 0; iter < 50000; iter++){\n if (trainCont % 2) {\n yield trainStep(data.batch());\n } else {\n break;\n }\n }\n yield trainStep(data.batch(), (trainCont >= 0));\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2199,"value":"**Main view**\n\nHere's where we'll visualize the diffusion model we're training. The points represent our dataset, while the heatmap shows the distribution after diffusing the data to a time t. You can control the current time with the slider below!\n\nYou can mouseover to show the gradient of the diffusion distribution (in red) and the estimated gradient (in orange). You can also click to add new datapoints!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":14,"value":"viewof coords = mainPlot.element()","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":124,"value":"viewof sigma = Inputs.range([0., 1.], {label: \"Time\", step: 0.01, width: 500})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1000,"value":"viewof sampleCont = Inputs.button([\n [\"Start\", () => 1],\n [\"Stop\", () => 2],\n [\"Reset\", () => 0],\n [\"Clear traces\", () => -1],\n], {label: \"Sample\", value: 0, disabled: (sampleView == \"None\"), depends: samplerOptions})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":169,"value":"\nPlot.plot({\n height: 150,\n y: {\n grid: true,\n label: {sigma: \"σ (std. dev.)\", s: \"α (scale)\", f: \"f (SDE coefficient)\", g: \"g (SDE coefficient)\", snr: \"Signal-to-noise\", logsnr: \"Log signal-to-noise\", variance: \"σ^2 (Variance)\"}[lineplotY],\n },\n x: {\n label: \"Time\",\n },\n marks: [\n Plot.ruleY([0]),\n Plot.ruleX(sampler.t_steps, {stroke: \"lightgray\"}),\n Plot.ruleX([sigma], {stroke: \"firebrick\", strokeWidth: 2}),\n Plot.ruleX([sample_loop.t_next], {stroke: \"steelblue\", strokeWidth: 2}),\n Plot.line(scheduler.linePlot(lineplotY), {x: \"0\", y: \"1\"}),\n ]\n})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2203,"value":"Here you can view the variance/scale/signal-to-noise ratio for times 0 to 1. The red line indicates the current plotted time, while the blue line indicates the state of the sampler. Grey lines show the discretization we're using for sampling.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2207,"value":"The *mode* selector lets us move from training to sampling. Try changing the mode to *sampling* and then click *start* below to sample from our trained model!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":998,"value":"viewof sampleView = Inputs.select(new Map([[\"Training\", \"None\"], [\"Sampling\", \"Samples\"], [\"Draw sample traces\", \"Traces\"] , [\"Inpainting\", \"Inpaint\"], [\"Likelihood\", \"Likelihood\"]]), {label: \"Mode\", value: \"None\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2403,"value":"The *mode* selector also has a few other options:\n- *Draw sample traces:* samples a single point at a time, tracing its path\n- *Inpainting:* Observes the *x* dimension and samples *y* given *x*.\n- *Likelihood:* Solves a probability flow ODE to find the likelihood of each point.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2210,"value":"Here we can change the variance schedule for our model. You might need to go back and re-train if you change this!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":219,"value":"viewof select = Inputs.select([\"VP\", \"VE\", \"SubVP\", \"Cosine\", \"Linear\"], {label: \"Schedule\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1190,"value":"### View options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2212,"value":"Here we can change what's shown above, including the y-axis for the lineplot above. *True grid* and *estimated grid* will show the vector field of true and estimated gradients for the distribution, try turning them on during training!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2080,"value":"viewof lineplotY = Inputs.select(new Map([[\"σ (std. dev.)\", \"sigma\"], [\"σ^2 (Variance)\", \"variance\"], [\"α (scale)\", \"s\"], [\"f (SDE coefficient)\", \"f\"] , [\"g (SDE coefficient)\", \"g\"], [\"Signal-to-noise\", \"snr\"], [\"Log signal-to-noise\", \"logsnr\"]]), {label: \"Line plot\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1131,"value":"viewof showCheckboxes = Inputs.checkbox([\"True gradient\", \"Estimated gradient\", \"True grid\", \"Estimated grid\", \"Diffusion distribution\", \"Contours\", \"Data\"], {label: \"Show\", value: [\"True gradient\", \"Estimated gradient\", \"Diffusion distribution\", \"Contours\", \"Data\"]})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1177,"value":"### Sampler options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2214,"value":"We can also change the sampler we're using. This defaults to ancestral sampling, but you can also sample forward in time starting from the data (diffusing) or using a probability flow ODE.\n\nTry changing the number of discretization steps!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":1077,"value":"viewof samplerOptions = Inputs.form({\n sampler: Inputs.select([\"ancestral\", \"forward\", \"euler_ode\", \"huen_ode\", \"forward_ode\", \"blended\"], {label: \"Sampler type\"}),\n num_steps: Inputs.range([1, 1000], {label: \"Discretization steps\", step: 1, value: 100}),\n nsamples: Inputs.range([1, 1000], {label: \"Number of samples\", step: 1, value: 250}),\n trace_color: Inputs.color({label: \"Trace color\"}),\n});","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2216,"value":"This is the \"observed\" dimension for inpainting mode. When in inpainting mode, this will be indicated by a red line in the main plot.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":1375,"value":"viewof xobserved = Inputs.range(mainPlot.scalex.domain, {label: \"Observed X value\", step: 0.01})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2531,"value":"This is resolution to compute the likelihood at for likelihood mode:","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2517,"value":"viewof gridRes = Inputs.range([10, 100], {label: 'Likelihood resolution', value: 40, step: 1})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2537,"value":"viewof likelihoodView = Inputs.select(['Likelihood', 'Log-likelihood'], {label: 'Show:'})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1184,"value":"### Dataset options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2219,"value":"Here you can change the base dataset and the number of observations!","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":327,"value":"viewof datasetName = Inputs.select([\"Moons\", \"Spiral\", \"Circles\", \"Blobs\"], {label: \"Dataset\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":333,"value":"viewof datasize = Inputs.range([1, 200], {label: \"Dataset size\", step: 1, value: 50})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1180,"value":"### Training options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2221,"value":"Here you can change the training objective and optimizer. (continuous/discrete time, elbo/simple weighting, Adam/SGD optimizer)","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2042,"value":"viewof lossType = Inputs.select([\"Continuous\", \"Discrete\"], {label: \"Loss type\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2045,"value":"viewof lossWeighting = Inputs.select([\"vlb\", \"simple\"], {label: \"Loss weighting\", value: \"simple\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":400,"value":"viewof optAlgorithm = Inputs.select([\"SGD\", \"Adam\"], {label: \"Optimizer\", value: \"Adam\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":403,"value":"viewof lr = Inputs.range([0.00001, 1.], {label: \"Learning rate\", transform: Math.log, step: 0.00001, value: 0.001})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1981,"value":"### Network options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2223,"value":"Here are more options for the score network, including the architecture and the parameterization (i.e. *what* the network actually predicts)","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2231,"value":"viewof net_parameterization = Inputs.select([\"x0\", \"eps\", \"score\"], {label: \"Parameterization\", value: \"eps\"})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1977,"value":"viewof networkOptions = Inputs.form({\n num_layers: Inputs.range([1, 8], {label: \"Layers\", step: 1, value: 2}),\n num_units: Inputs.range([1, 1024], {label: \"Hidden units per layer\", step: 1, value: 256}),\n activation: Inputs.select([\"elu\", \"relu\", \"tanh\"], {label: \"Activation function\"}),\n residual: Inputs.range([0, 4], {label: \"Residual connection freq.\", step: 1, value: 2}),\n embedding_dim: Inputs.range([1, 1024], {label: \"Embedding dimension\", step: 1, value: 256}),\n embedding: Inputs.select([\"raw\", \"embedded\"], {label: \"Input preprocessing\", value: \"embedded\"}),\n});","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2333,"value":"### Render image","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2353,"value":"Saving the main plot as an image might not work well. Here you can render the plot as a PNG for publication. Use *scale* to scale the resolution. Once you've rendered a plot a preview will be shown, you can download the image using menu in the top left corner of the cell.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2335,"value":"{\n if (dorender) {\n return saveImage(imageScale);\n }\n return document.createElement('a');\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2342,"value":"viewof dorender = {\n clearImage;\n return Inputs.button(\"Render plot\");\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2357,"value":"viewof clearImage = Inputs.button(\"Clear\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2339,"value":"viewof imageScale = Inputs.range([0.1, 10], {label: \"Scale\", step: 0.1, value: 1})","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2000,"value":"### UI Options","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2246,"value":"If you'd like to customize the plot to make your own visualizations, you can change the settings here.","pinned":false,"hidden":false,"mode":"md","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"","database":"$undefined"},{"id":2004,"value":"viewof uiOptions = Inputs.form({\n colormap: Inputs.select(colorMapTexture.COLOR_SCALE_NAMES, {label: \"Colormap\"}),\n truegrad: Inputs.color({label: \"True gradient color\", value: \"#EE4B2B\"}),\n estgrad: Inputs.color({label: \"Estimated gradient color\", value: \"#F28C28\"}),\n sample_fill: Inputs.color({label: \"Sample color\", value: \"#FFFFFF\"}),\n sample_stroke: Inputs.color({label: \"Sample stroke color\", value: \"#000000\"}),\n sample_r: Inputs.range([0, 20], {label: \"Sample point size\", value: 3}),\n sample_strokeWidth: Inputs.range([0, 5], {label: \"Sample stroke width\", value: 1.}),\n data_fill: Inputs.color({label: \"Data color\", value: \"#F2D2BD\"}),\n data_stroke: Inputs.color({label: \"Data stroke color\", value: \"#000000\"}),\n data_r: Inputs.range([0, 20], {label: \"Data point size\", value: 3}),\n data_strokeWidth: Inputs.range([0, 5], {label: \"Data stroke width\", value: 1.}),\n});","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2467,"value":"uiOptions","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1330,"value":"{\n const lossPlot = sample_loop.watch();\n return lossPlot;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1052,"value":"allTraces = {\n sampleView;\n return [];\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":665,"value":"sample_loop = {\n if (sampler.finished && sampleCont == 1){\n sampler.reset(true);\n }\n while (!sampler.finished && sampleCont == 1){\n //await Promises.delay(1);\n if (sampleView == \"Inpaint\") {\n sampler.step([xobserved, NaN]);\n } else {\n sampler.step();\n }\n \n allTraces.pop();\n allTraces.push(sampler.getTrace());\n\n yield sampler;\n }\n\n if (sampleCont == -1){\n allTraces.length = 0;\n }\n if (sampleCont <= 0) {\n sampler.reset(true, sampleView == \"Inpaint\");\n if(allTraces.length >= 1 && allTraces[allTraces.length - 1].trace.length <= 1) {\n allTraces.pop();\n }\n allTraces.push(sampler.getTrace());\n }\n yield sampler;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":11,"value":"mainPlot = new heatmapPlot({levels: 8, \n contourWidth: 0., \n step: false, pixelRatio: 2,\n colormap: showCheckboxes.includes(\"Diffusion distribution\") ? uiOptions.colormap : \"White\", \n levels: !showCheckboxes.includes(\"Contours\") || 8});","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":228,"value":"$43","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":479,"value":"$44","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":251,"value":"trueGradPointer = {\n if (!showCheckboxes.includes(\"True gradient\")) {\n return;\n }\n const svg = d3.select(mainPlot.svg());\n \n svg.selectAll('#truegradpointer').remove();\n\n if (isFinite(coords[0]) && isFinite(coords[1])){\n const coordst = tf.tensor([coords], undefined, 'float32');\n const grad = scheduler.trueGrad(coordst, data.tensor(), sigma)[0];\n coordst.dispose();\n\n const scalex = mainPlot.scalex;\n const scaley = mainPlot.scaley;\n\n const normGrad = Math.sqrt(grad[0] * grad[0] + grad[1] * grad[1]);\n \n svg.append('line').attr('id', 'truegradpointer')\n .attr(\"x1\", scalex.apply(coords[0]))\n .attr(\"y1\", scaley.apply(coords[1]))\n .attr(\"x2\", scalex.apply(grad[0]))// + 2 * grad[0] / normGrad))\n .attr(\"y2\", scaley.apply(grad[1]))// + 2 * grad[1] / normGrad))\n .attr('stroke', uiOptions.truegrad)\n .attr('stroke-width', 2);\n return grad;\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2376,"value":"gridinput = tf.tidy(() => {\n const xrange = tf.linspace(mainPlot.scalex.domain[0], mainPlot.scalex.domain[1], gridRes);\n const yrange = tf.linspace(mainPlot.scaley.domain[0], mainPlot.scaley.domain[1], gridRes);\n const grid = tf.meshgrid(xrange, yrange);\n return tf.stack([grid[0].reshape([-1]), grid[1].reshape([-1])], 1);\n});","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":646,"value":"sampler = {\n const options = {...samplerOptions, batchsize: samplerOptions.nsamples, data: data};// batchsize: (sampleView == \"Samples\" ? samplerOptions.nsamples : 1)};\n //options.batchsize = (sampleView == \"Samples\" ? options.nsamples : 1);\n //return makeSampler((x, s, t) => scheduler.predict(x, network, s, t), scheduler, options);\n if (sampleView == 'Likelihood') {\n options.data = gridinput;\n options.sampler = 'forward_ode';\n }\n return makeSampler(network, scheduler, options); \n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":559,"value":"$45","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":65,"value":"$46","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1449,"value":"norm = {\n function cdf(x){\n const sqrt2 = Math.sqrt(2.);\n return tf.erf(tf.div(x, sqrt2)).add(1.).mul(0.5);\n }\n\n function kl_divergence(mu0, sigma0, mu1, sigma1) {\n const t1 = sigma0.div(sigma1).square();\n const t2 = (sigma1.square().div(sigma0.square())).log();\n const t3 = (mu1.sub(mu0)).square().div(sigma1.square());\n return t1.add(t2).add(t3).sub(1.).mul(0.5);\n }\n\n function logpdf(x, mu0, sigma0) {\n return (x.sub(mu0)).div(sigma0).square().mul(-0.5).sub(sigma0.log()).sub(0.5 * Math.log(2 * Math.PI));\n }\n\n return {cdf: cdf, logpdf: logpdf, kl_divergence: kl_divergence}\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1483,"value":"$47","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":1549,"value":"utils = {\n function getRandomInt(max) {\n return Math.floor(Math.random() * max);\n } \n\n function expand_size(x) {\n \n }\n\n return {\n getRandomInt: getRandomInt,\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":427,"value":"trainStep = makeTrainStep(network, lr, optAlgorithm, scheduler);","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":396,"value":"$48","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":26,"value":"$49","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":963,"value":"viewof networkObs = network.watch();","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":376,"value":"network = {\n resetNet;\n var hidden_layers = [];\n\n for (let i=0; i < networkOptions.num_layers; i++){\n if (networkOptions.residual) {\n var block = [];\n for (let j=0; j < networkOptions.residual; j++){\n block.push(networkOptions.num_units);\n }\n hidden_layers.push(block);\n }\n else {\n hidden_layers.push(networkOptions.num_units);\n }\n }\n\n const options = {\n hidden_layers: hidden_layers,\n activation: networkOptions.activation,\n embedding_dim: networkOptions.embedding_dim,\n embedding_2d: networkOptions.embedding == 'embedded',\n };\n \n return makeNetwork(options);\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":348,"value":"$4a","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":661,"value":"$4b","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":292,"value":"PointAdder = {\n d3.select(mainPlot.element()).on('mousedown', (event) => {\n const point = d3.pointer(event, mainPlot.svg());\n const realpoint = [mainPlot.scalex.invert(point[0]), mainPlot.scaley.invert(point[1])];\n data.add(realpoint);\n mainPlot.plotFunction(p => scheduler.heatmapPlot(p, data.tensor(), sigma));\n mainPlot.plot(data.marks());\n });\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":36,"value":"viewof data = dataset(getDataset(datasetName, datasize), {batchsize: 256, r: uiOptions.data_r, fill: uiOptions.data_fill, stroke_width: uiOptions.data_strokeWidth, stroke: uiOptions.data_stroke}).watch();","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":16,"value":"mainHeatmapPlot = {\n if (sampleView != \"Likelihood\") {\n mainPlot.plotFunction(p => scheduler.heatmapPlot(p, data.tensor(), sigma));\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2499,"value":"mainHeatmapPlotLikelihood = {\n if (sampleView == \"Likelihood\") {\n mainPlot.setHeatmapData(tf.tidy(() => {\n let heatmap = sample_loop.lp\n if (likelihoodView == 'Likelihood'){\n heatmap = heatmap.exp();\n }\n heatmap = heatmap.reshape([gridRes, gridRes, 1]);\n return tf.image.resizeBilinear(heatmap, [100, 100]).squeeze();\n }), {colormap: uiOptions.likelihood_colormap});\n }\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2051,"value":"mainMarks = mainPlot.plot(marks);","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":5,"value":"import {heatmapPlot, tf, colorMapTexture} from \"60f5a4ab96ab0672\"","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2316,"value":"html2canvas = require(\"https://html2canvas.hertzen.com/dist/html2canvas.min.js\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":2319,"value":"function saveImage(scale) {\n const image = html2canvas(mainPlot.element(), {scale: scale}).then((c) => {\n const img = new Image();\n img.src = c.toDataURL('image/png');\n return img;\n }\n );\n return image;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","tool":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"}],"files":"$W4c","stdlib":"1"}} 42:"$41:body"