| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np, math, os |
|
|
| OUT = "/root/alblit/diffusion/figures" |
| os.makedirs(OUT, exist_ok=True) |
| plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9, |
| "axes.spines.top": False, "axes.spines.right": False, |
| "axes.edgecolor": "#444", "figure.dpi": 200}) |
| AR, DF = "#9aa5b1", "#c2410c" |
|
|
| def ci(p, n): return 1.96 * math.sqrt(p * (1 - p) / n) * 100 |
|
|
| fig, ax = plt.subplots(1, 2, figsize=(9, 3.2), gridspec_kw={"width_ratios": [1.35, 1]}) |
| names = ["MMLU\n(n=200)", "GSM8K\n(n=50)"] |
| ar = [42.0, 21.0]; df = [45.5, 4.0]; ns = [200, 50] |
| x = np.arange(2) |
| ax[0].bar(x - .19, ar, .37, color=AR, label="AR abliterated") |
| ax[0].bar(x + .19, df, .37, color=DF, label="Diffusion") |
| ax[0].errorbar(x - .19, ar, yerr=[ci(a/100, n) for a, n in zip(ar, ns)], fmt="none", ecolor="#475569", capsize=3) |
| ax[0].errorbar(x + .19, df, yerr=[ci(d/100, n) for d, n in zip(df, ns)], fmt="none", ecolor="#7c2d12", capsize=3) |
| ax[0].axhline(25, ls=":", color="#94a3b8", lw=1) |
| ax[0].text(1.42, 26, "chance", fontsize=7, color="#64748b") |
| ax[0].set_xticks(x); ax[0].set_xticklabels(names) |
| ax[0].set_ylabel("accuracy (%)"); ax[0].set_ylim(0, 65) |
| ax[0].legend(frameon=False, loc="upper right") |
| ax[0].set_title("Knowledge survives, multi-step reasoning does not", loc="left", fontsize=10) |
| for i, (a, d) in enumerate(zip(ar, df)): |
| ax[0].text(i - .19, a + 2.5, f"{a:.0f}", ha="center", fontsize=7, color="#475569") |
| ax[0].text(i + .19, d + 2.5, f"{d:.0f}", ha="center", fontsize=7, color=DF) |
|
|
| sp = [27.3, 61.7, 124.2, 218.9] |
| lab = ["AR", "64", "32", "16"] |
| cols = [AR, DF, DF, DF] |
| ax[1].bar(range(4), sp, color=cols, width=.65) |
| ax[1].set_xticks(range(4)); ax[1].set_xticklabels(lab) |
| ax[1].set_xlabel("diffusion steps (AR = baseline)") |
| ax[1].set_ylabel("tokens/s") |
| ax[1].set_title("Parallel unmasking wins on throughput", loc="left", fontsize=10) |
| for i, v in enumerate(sp): |
| ax[1].text(i, v + 6, f"{v:.0f}", ha="center", fontsize=7, color="#475569") |
| ax[1].set_ylim(0, 250) |
| fig.tight_layout(); fig.savefig(f"{OUT}/fig1_results.png", bbox_inches="tight"); plt.close(fig) |
|
|
| steps = [16, 32, 64] |
| tps = [218.9, 124.2, 61.7] |
| fig, ax = plt.subplots(figsize=(6.4, 2.9)) |
| ax.plot(steps, tps, "o-", color=DF, lw=1.6, ms=6, label="diffusion") |
| ax.axhline(27.3, ls="--", color=AR, lw=1.2, label="AR, KV cache") |
| ax.set_xlabel("denoising steps for 128 tokens") |
| ax.set_ylabel("tokens/s") |
| ax.set_xticks(steps) |
| ax.legend(frameon=False) |
| ax.set_title("Throughput is set by the step count", loc="left", fontsize=10) |
| ax.annotate("8.0x", xy=(16, 218.9), xytext=(24, 205), fontsize=8, color=DF, |
| arrowprops=dict(arrowstyle="-", color="#cbd5e1", lw=.8)) |
| ax.annotate("2.3x", xy=(64, 61.7), xytext=(46, 95), fontsize=8, color=DF, |
| arrowprops=dict(arrowstyle="-", color="#cbd5e1", lw=.8)) |
| fig.tight_layout(); fig.savefig(f"{OUT}/fig2_speed.png", bbox_inches="tight"); plt.close(fig) |
|
|
| st = [100, 500, 1000, 2000, 3000, 4000] |
| loss = [0.4137, 0.3258, 0.3482, 0.3147, 0.3018, 0.2953] |
| fig, ax = plt.subplots(figsize=(6.4, 2.7)) |
| ax.plot(st, loss, "o-", color=DF, lw=1.6, ms=5) |
| ax.set_xlabel("training step") |
| ax.set_ylabel("loss (EMA)") |
| ax.set_ylim(0, 0.5) |
| ax.set_title("65M tokens, 55 minutes on one H200", loc="left", fontsize=10) |
| fig.tight_layout(); fig.savefig(f"{OUT}/fig3_loss.png", bbox_inches="tight"); plt.close(fig) |
| print("ok", os.listdir(OUT)) |
|
|
