iclr 2026 (tmp)

tmp/compare_fp16_vs_bf16.py

@@ -0,0 +1,59 @@
+# filepath: /home/tan/research/ctx-to-lora/tmp/compare_fp16_vs_bf16.py
+
+import torch
+
+
+def compare_fp16_vs_bf16():
+    """
+    Compare operations between float16 and bfloat16 tensors.
+    Computes difference and multiplication to see output dtypes.
+    """
+
+    # Create sample tensors with the same values but different dtypes
+    values = [[1.5, 2.3, 3.7], [4.1, 5.9, 6.2]]
+
+    # Create tensors with different dtypes
+    tensor_fp16 = torch.tensor(values, dtype=torch.float16)
+    tensor_bf16 = torch.tensor(values, dtype=torch.bfloat16)
+
+    print("Original tensors:")
+    print(f"FP16 tensor: {tensor_fp16}")
+    print(f"FP16 dtype: {tensor_fp16.dtype}")
+    print(f"BF16 tensor: {tensor_bf16}")
+    print(f"BF16 dtype: {tensor_bf16.dtype}")
+    print()
+
+    # Compute difference (subtraction)
+    diff_result = tensor_fp16 - tensor_bf16
+    print("Difference operation (FP16 - BF16):")
+    print(f"Result: {diff_result}")
+    print(f"Result dtype: {diff_result.dtype}")
+    print()
+
+    # Compute multiplication
+    mult_result = tensor_fp16 * tensor_bf16
+    print("Multiplication operation (FP16 * BF16):")
+    print(f"Result: {mult_result}")
+    print(f"Result dtype: {mult_result.dtype}")
+    print()
+
+    # Test the reverse order to see if it matters
+    diff_result_reverse = tensor_bf16 - tensor_fp16
+    mult_result_reverse = tensor_bf16 * tensor_fp16
+
+    print("Reverse order operations:")
+    print(f"BF16 - FP16 dtype: {diff_result_reverse.dtype}")
+    print(f"BF16 * FP16 dtype: {mult_result_reverse.dtype}")
+    print()
+
+    # Show the actual values to check precision differences
+    print("Value comparison:")
+    print(f"Original values: {values}")
+    print(f"FP16 values: {tensor_fp16.tolist()}")
+    print(f"BF16 values: {tensor_bf16.tolist()}")
+    print(f"Difference values: {diff_result.tolist()}")
+    print(f"Multiplication values: {mult_result.tolist()}")
+
+
+if __name__ == "__main__":
+    compare_fp16_vs_bf16()

tmp/display_self_gen_data.py

@@ -0,0 +1,27 @@
+import sys
+
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+def display_sample(ds, tk, i):
+    sample = ds[i]
+    print(f"ctx: {tk.decode(sample['ctx_ids'])}")
+    for ids, (start, end) in zip(sample["input_ids"], sample["response_start_end"]):
+        print(f"ids: {ids}")
+        print(f"input: {tk.decode(ids)}")
+        print(f"response: {tk.decode(ids[start:])}")
+        len_input_ids = len(ids)
+        pad_len_right = len_input_ids - end
+        print(f"{pad_len_right=}")
+
+
+data_files = sys.argv[1]
+
+ds = load_dataset("parquet", data_files=data_files, split="train")
+tk = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+
+for i in range(5):
+    display_sample(ds, tk, i)
+
+breakpoint()

tmp/fit_sigmoid_curve.py

@@ -0,0 +1,78 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.optimize import curve_fit
+
+
+def sigmoid(x, k, x0):
+    """
+    Sigmoid function with L fixed at 1.
+
+    Parameters:
+    x (array-like): Input data (independent variable).
+    k (float): The steepness or growth rate of the curve.
+    x0 (float): The x-value of the sigmoid's midpoint.
+    """
+    return 1 / (1 + np.exp(-k * (x - x0)))
+
+
+# --- User Input ---
+# Replace these with your actual x and y data points.
+# You can have 3 or more points.
+x_data = np.array([10, 20, 30])
+y_data = np.array([0.503, 0.554, 0.590])
+# --------------------
+
+# Check if the data is suitable for a sigmoid curve (monotonically increasing)
+if not np.all(np.diff(y_data) > 0):
+    print("Warning: The y-values of your data are not strictly increasing.")
+    print("A sigmoid fit may not be appropriate for this data.")
+
+try:
+    # Use curve_fit to find the best parameters
+    # p0 is an initial guess for the parameters [k, x0]
+    # We can estimate x0 as the mean of x_data as a starting point.
+    initial_guess = [1, np.mean(x_data)]
+
+    # The 'bounds' argument constrains the parameters. We'll keep k positive.
+    popt, pcov = curve_fit(
+        sigmoid, x_data, y_data, p0=initial_guess, bounds=(0, np.inf)
+    )
+
+    # Extract the optimal parameters
+    k_opt, x0_opt = popt
+    print("Optimal Parameters (L=1):")
+    print(f"k = {k_opt:.4f}")
+    print(f"x0 = {x0_opt:.4f}")
+
+    # --- Visualization ---
+    # Generate a smooth curve for plotting
+    x_curve = np.linspace(0, 300, 200)
+    y_curve = sigmoid(x_curve, k_opt, x0_opt)
+
+    # Plot the data and the fitted curve
+    plt.figure(figsize=(10, 6))
+    plt.plot(x_data, y_data, "o", label="Data Points", markersize=8, color="red")
+    plt.plot(
+        x_curve,
+        y_curve,
+        label=f"Fitted Sigmoid Curve\n(k={k_opt:.2f}, x₀={x0_opt:.2f})",
+        color="blue",
+    )
+    plt.plot([0, 300], [0.823, 0.823], label="Ours")
+    plt.plot([0, 300], [0.988, 0.988], label="CD (oracle)")
+    plt.plot([40, 40], [0, 1.2], "--", label="OOM")
+
+    plt.title("Overdetermined Sigmoid Curve Fit (L=1)")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    plt.legend()
+    plt.grid(True)
+    plt.ylim(0, 1.1)  # Set y-axis limits since L=1
+    plt.show()
+    plt.savefig("sigmoid_fit.png")
+
+except RuntimeError as e:
+    print(f"Error: Could not find optimal parameters. {e}")
+    print(
+        "This can happen if the data is not sigmoid-shaped or if the initial guess is poor."
+    )

tmp/plot_bar_perf_ctx_latency_mem.py

@@ -0,0 +1,375 @@
+from dataclasses import dataclass  # added
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Styling (reused)
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#FFB83D",
+        # "#555555",
+        "#9BC750",
+        # "#35775A",
+        "#5571EA",
+        # "#6CACFF",
+        # "#E84494",
+        # "#E47FB0",
+        # "#785EF0",
+        # "#FF924E",
+        # "#5BC5DB",
+        # "#AD6F50",
+        # "#A1A9AD",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+plt.rcParams["legend.fontsize"] = 16
+plt.rcParams["xtick.labelsize"] = 13
+plt.rcParams["ytick.labelsize"] = 13
+
+
+@dataclass
+class Point:
+    performance: float
+    context_ratio: float | None = None
+    latency: float | None = None
+    latency_std: float | None = None
+    peak_mem: float | None = None  # added
+    peak_mem_std: float | None = None  # added
+
+
+def combine_std_for_sum(std_list):
+    return np.sqrt(sum(std**2 for std in std_list))
+
+
+# Unified data structure (merged context + latency via Point)
+datasets = {
+    "2WikiMultihopQA\n(Long Context)": {
+        "Base model w/ context": [
+            Point(
+                performance=0.3387,
+                context_ratio=1.0,
+            ),
+        ],
+        "CD": [
+            Point(
+                performance=0.2385,
+                context_ratio=0.0,
+                latency=72.537,
+                latency_std=10,
+                peak_mem=79.371,
+            ),
+        ],
+        "Ours": [
+            Point(
+                performance=0.2857,
+                context_ratio=0.0,
+                latency=0.551,
+                latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]),
+                peak_mem=3.791,
+            ),
+        ],
+    },
+    # "SQuAD": {
+    #     "Base model w/ context": [
+    #         Point(performance=0.8692, context_ratio=1.0),
+    #     ],
+    #     # "Base model w/o context": [
+    #     #     Point(
+    #     #         performance=0.1866,
+    #     #         context_ratio=0.0,
+    #     #         latency=0.0,
+    #     #         latency_std=0.0,
+    #     #         peak_mem=0.0,
+    #     #     ),
+    #     # ],
+    #     # "CD (oracle)": [
+    #     #     Point(
+    #     #         performance=0.859,
+    #     #         context_ratio=0.0,
+    #     #         latency=39.974 + 456.989 * 1e-3,
+    #     #         latency_std=combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]),
+    #     #         peak_mem=855.762 / 1024,
+    #     #     ),
+    #     # ],
+    #     "CD": [
+    #         Point(
+    #             performance=0.5986,
+    #             context_ratio=0.0,
+    #             latency=21.254 * 4 + 49.358 + 7.450,
+    #             latency_std=10,
+    #             peak_mem=41.098,
+    #         ),
+    #     ],
+    #     # "Ours (batched)": [
+    #     #     Point(
+    #     #         performance=0.7174,
+    #     #         # context_ratio=0.0,
+    #     #         latency=85.954 * 1e-3 + 259.920 * 1e-6,
+    #     #         latency_std=combine_std_for_sum([31.057 * 1e-3, 740.286 * 1e-6]),
+    #     #         peak_mem=1.151,
+    #     #     ),
+    #     # ],
+    #     "Ours": [
+    #         Point(
+    #             performance=0.715,
+    #             context_ratio=0.0,
+    #             latency=426.849 * 1e-3 + 260.073 * 1e-6,
+    #             latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]),
+    #             peak_mem=328.134 / 1024,
+    #         ),
+    #     ],
+    #     # "T2L": [
+    #     #     Point(
+    #     #         performance=0.1761,
+    #     #         context_ratio=0.0,
+    #     #         latency=35.480 * 1e-3,
+    #     #         latency_std=22.193 * 1e-3,
+    #     #         peak_mem=1.793,
+    #     #     ),
+    #     # ],
+    #     # "LLMLingua-2": [
+    #     #     Point(performance=0.8492, context_ratio=0.9),
+    #     #     Point(performance=0.8278, context_ratio=0.8),
+    #     #     Point(performance=0.7749, context_ratio=0.6),
+    #     #     Point(performance=0.6624, context_ratio=0.4),
+    #     #     Point(performance=0.4271, context_ratio=0.2),
+    #     #     Point(performance=0.3192, context_ratio=0.1),
+    #     # ],
+    # },
+    # "DROP": {
+    # "Base model w/ context": [
+    #     Point(performance=0.4541, context_ratio=1.0),
+    # ],
+    # "Base model w/o context": [
+    #     Point(
+    #         performance=0.1417,
+    #         context_ratio=0.0,
+    #         latency=0.0,
+    #         latency_std=0.0,
+    #         peak_mem=0.0,
+    #     ),
+    # ],
+    # "CD (oracle)": [
+    #     Point(
+    #         performance=0.443,
+    #         context_ratio=0.0,
+    #         latency=39.974 + 422.820 * 1e-3,
+    #         latency_std=combine_std_for_sum([965.326 * 1e-3, 241.419 * 1e-3]),
+    #         peak_mem=1.381,
+    #     ),
+    # ],
+    # "CD": [
+    #     Point(
+    #         performance=0.2289,
+    #         context_ratio=0.0,
+    #         latency=23.883 * 4 + 9.755 + 41.301,
+    #         latency_std=14.8852904910,
+    #         peak_mem=44.305,
+    #     ),
+    # ],
+    # "Ours (batched)": [
+    #     Point(
+    #         performance=0.2972,
+    #         # context_ratio=0.0,
+    #         latency=84.132 * 1e-3 + 251.876 * 1e-6,
+    #         latency_std=combine_std_for_sum([30.503 * 1e-3, 22.829 * 1e-6]),
+    #         peak_mem=1.976,
+    #     ),
+    # ],
+    # "Ours": [
+    #     Point(
+    #         performance=0.3002,
+    #         context_ratio=0.0,
+    #         latency=419.462 * 1e-3 + 254.899 * 1e-6,
+    #         latency_std=combine_std_for_sum([40.211 * 1e-3, 29.784 * 1e-6]),
+    #         peak_mem=429.881 / 1024,
+    #     ),
+    # ],
+    # "T2L": [
+    #     Point(
+    #         performance=0.1468,
+    #         context_ratio=0.0,
+    #         latency=35.036 * 1e-3,
+    #         latency_std=18.144 * 1e-3,
+    #         peak_mem=3.004,
+    #     ),
+    # ],
+    # "LLMLingua-2": [  #
+    #     Point(performance=0.4513, context_ratio=0.9),
+    #     Point(performance=0.4488, context_ratio=0.8),
+    #     Point(performance=0.4076, context_ratio=0.6),
+    #     Point(performance=0.3488, context_ratio=0.4),
+    #     Point(performance=0.2617, context_ratio=0.2),
+    #     Point(performance=0.1712, context_ratio=0.1),
+    # ],
+    # },
+    # "ROPES": {
+    # "Base model w/ context": [
+    #     Point(performance=0.7457, context_ratio=1.0),
+    # ],
+    # "Base model w/o context": [
+    #     Point(
+    #         performance=0.4545,
+    #         context_ratio=0.0,
+    #         latency=0.0,
+    #         latency_std=0.0,
+    #         peak_mem=0.0,
+    #     ),
+    # ],
+    # "CD (oracle)": [
+    #     Point(
+    #         performance=0.7349,
+    #         context_ratio=0.0,
+    #         latency=39.633 + 285.467 * 1e-3,
+    #         latency_std=combine_std_for_sum([1.752, 72.791 * 1e-3]),
+    #         peak_mem=527.312 / 1024,
+    #     ),
+    # ],
+    # "CD": [
+    #     Point(
+    #         performance=0.5786,
+    #         context_ratio=0.0,
+    #         latency=23.714 * 4 + 11.473 + 59.660,
+    #         latency_std=11.7329040736,
+    #         peak_mem=43.271,
+    #     ),
+    # ],
+    # "Ours (batched)": [
+    #     Point(
+    #         performance=0.6759,
+    #         # context_ratio=0.0,
+    #         latency=85.652 * 1e-3 + 254.029 * 1e-6,
+    #         latency_std=combine_std_for_sum([30.852 * 1e-3, 51.127 * 1e-6]),
+    #         peak_mem=496.898 / 1024,
+    #     ),
+    # ],
+    # "Ours": [
+    #     Point(
+    #         performance=0.6786,
+    #         context_ratio=0.0,
+    #         latency=417.860 * 1e-3 + 257.422 * 1e-6,
+    #         latency_std=combine_std_for_sum([45.514 * 1e-3, 31.419 * 1e-6]),
+    #         peak_mem=273.867 / 1024,
+    #     ),
+    # ],
+    # "T2L": [
+    #     Point(
+    #         performance=0.4661,
+    #         context_ratio=0.0,
+    #         latency=37.913 * 1e-3,
+    #         latency_std=17.313 * 1e-3,
+    #         peak_mem=697.217 / 1024,
+    #     ),
+    # ],
+    # "LLMLingua-2": [
+    #     Point(performance=0.708, context_ratio=0.9),
+    #     Point(performance=0.7126, context_ratio=0.8),
+    #     Point(performance=0.7291, context_ratio=0.6),
+    #     Point(performance=0.677, context_ratio=0.4),
+    #     Point(performance=0.6188, context_ratio=0.2),
+    #     Point(performance=0.4894, context_ratio=0.1),
+    # ],
+    # },
+}
+
+
+def build_color_map(datasets):
+    methods = []
+    for data in datasets.values():
+        for m in data.keys():
+            if m not in methods:
+                methods.append(m)
+    palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+    return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
+
+
+def plot_bar_methods(datasets):
+    # Use first dataset found
+    dataset_name, methods = next(iter(datasets.items()))
+    color_map = build_color_map(datasets)
+
+    ref_perf = methods["Base model w/ context"][0].performance
+    # Aggregate metrics per method
+    stats_by_method = {}
+    for method, points in methods.items():
+        if method == "Base model w/ context":
+            continue
+        if not points:
+            continue
+
+        perf_vals = [p.performance for p in points if p.performance is not None]
+        lat_vals = [p.latency for p in points if p.latency is not None]
+        mem_vals = [p.peak_mem for p in points if p.peak_mem is not None]
+        perf = max(perf_vals) / ref_perf if perf_vals else np.nan
+        lat = float(np.mean(lat_vals)) if lat_vals else np.nan
+        mem = float(np.mean(mem_vals)) if mem_vals else np.nan
+        stats_by_method[method] = {"perf": perf, "latency": lat, "memory": mem}
+
+    # Ordering by performance (descending), NaNs last
+    with_perf = [m for m, s in stats_by_method.items() if not np.isnan(s["perf"])]
+    without_perf = [m for m, s in stats_by_method.items() if np.isnan(s["perf"])]
+    order = (
+        sorted(with_perf, key=lambda m: stats_by_method[m]["perf"], reverse=True)
+        + without_perf
+    )
+
+    # Build 3 horizontal bar plots: perf, latency, memory
+    fig, axes = plt.subplots(3, 1, figsize=(6, 4), sharex=False)
+    metric_specs = [
+        ("perf", "Normalized\nPerformance"),
+        ("latency", "Latency\n(s)"),
+        ("memory", "Memory\n(GB)"),
+    ]
+    for ax, (metric_key, y_label) in zip(axes, metric_specs):
+        names = [m for m in order if not np.isnan(stats_by_method[m][metric_key])]
+        values = [stats_by_method[m][metric_key] for m in names]
+        if not names:
+            ax.axis("off")
+            continue
+        y = np.arange(len(names))
+        colors = [color_map.get(m, "#888888") for m in names]
+        ax.barh(y, values, color=colors, edgecolor="black", alpha=0.9)
+        ax.set_yticks(y, labels=names)
+        ax.invert_yaxis()  # highest at top to match ordering
+        ax.set_ylabel(y_label, fontweight="bold", fontsize=12)
+        ax.grid(True, axis="x", alpha=0.3)
+        # Make latency and memory log-scale
+        if metric_key in ("latency", "memory"):
+            ax.set_xscale("log")
+
+    fig.suptitle(dataset_name, fontweight="bold", fontsize=20, y=0.91)
+    plt.tight_layout(rect=[0, 0, 1, 0.96])
+    return fig
+
+
+if __name__ == "__main__":
+    # fig = plot_ctx_and_latency(datasets)
+    fig = plot_bar_methods(datasets)
+    plt.show()
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem_bar.png",
+        dpi=300,
+        bbox_inches="tight",
+    )
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem_bar.pdf",
+        dpi=300,
+        bbox_inches="tight",
+    )

tmp/plot_niah_mem_latency_vs_ctx_len.py

@@ -0,0 +1,251 @@
+from dataclasses import dataclass  # added
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Styling (reused)
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#FFB83D",
+        "#5571EA",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+plt.rcParams["legend.fontsize"] = 16
+plt.rcParams["xtick.labelsize"] = 16
+plt.rcParams["ytick.labelsize"] = 16
+
+
+@dataclass
+class Point:
+    context_length: int
+    latency: float | None = None
+    latency_std: float | None = None
+    peak_mem: float | None = None
+
+
+def combine_std_for_sum(std_list):
+    return np.sqrt(sum(std**2 for std in std_list))
+
+
+# Unified data structure (merged context + latency via Point)
+datasets = {
+    "NIAH": {
+        "Base model w/ ICL": [
+            Point(
+                context_length=1024,
+                latency=375.689 * 1e-3,
+                latency_std=19.718 * 1e-3,
+                peak_mem=210.722 / 1024,
+            ),
+            Point(
+                context_length=2**11,
+                latency=375.428 * 1e-3,
+                latency_std=794.507 * 1e-6,
+                peak_mem=235.718 / 1024,
+            ),
+            Point(
+                context_length=2**12,
+                latency=384.812 * 1e-3,
+                latency_std=4.376 * 1e-3,
+                peak_mem=429.681 / 1024,
+            ),
+            Point(
+                context_length=2**13,
+                latency=438.886 * 1e-3,
+                latency_std=17.096 * 1e-3,
+                peak_mem=1.165,
+            ),
+            Point(
+                context_length=2**14,
+                # latency=13.877,
+                # latency_std=394.682,
+                peak_mem=1.740,
+            ),
+            Point(
+                context_length=2**15,
+                latency=13.985,
+                latency_std=66.566 * 1e-3,
+                peak_mem=3.223,
+            ),
+            Point(
+                context_length=2**16,
+                latency=14.953,
+                latency_std=131.882 * 1e-3,
+                peak_mem=5.278,
+            ),
+            Point(
+                context_length=2**17,
+                latency=18.489,
+                latency_std=227.430 * 1e-3,
+                peak_mem=12.072,
+            ),
+        ],
+        "Ours": [
+            Point(
+                context_length=1024,
+                latency=300.995 * 1e-3,
+                latency_std=24.009 * 1e-3,
+                peak_mem=30.988 / 1024,
+            ),
+            Point(
+                context_length=2048,
+                latency=296.856 * 1e-3,
+                latency_std=568.892 * 1e-6,
+                peak_mem=3.207 / 1024,
+            ),
+            Point(
+                context_length=2**12,
+                latency=297.037 * 1e-3,
+                latency_std=548.641 * 1e-6,
+                peak_mem=4.363 / 1024,
+            ),
+            Point(
+                context_length=2**13,
+                latency=296.507 * 1e-3,
+                latency_std=18.850 * 1e-3,
+                peak_mem=33.449 / 1024,
+            ),
+            Point(
+                context_length=2**14,
+                latency=297.774 * 1e-3,
+                latency_std=788.634 * 1e-6,
+                peak_mem=8.792 / 1024,
+            ),
+            Point(
+                context_length=2**15,
+                latency=297.553 * 1e-3,
+                latency_std=2.016 * 1e-3,
+                peak_mem=14.457 / 1024,
+            ),
+            Point(
+                context_length=2**16,
+                latency=296.913 * 1e-3,
+                latency_std=13.901 * 1e-3,
+                peak_mem=25.003 / 1024,
+            ),
+            Point(
+                context_length=2**17,
+                latency=274.299 * 1e-3,
+                latency_std=53.295 * 1e-3,
+                peak_mem=47.696 / 1024,
+            ),
+        ],
+    }
+}
+
+
+def build_color_map(datasets):
+    methods = []
+    for data in datasets.values():
+        for m in data.keys():
+            if m not in methods:
+                methods.append(m)
+    palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+    return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
+
+
+def plot_memory_and_latency(datasets):
+    """Plot memory and latency vs context length in two subplots."""
+    color_map = build_color_map(datasets)
+
+    fig, (ax1) = plt.subplots(1, 1, figsize=(9, 6))
+
+    # Plot 1: Memory vs Context Length
+    for dataset_name, methods in datasets.items():
+        for method_name, points in methods.items():
+            # Extract data with valid memory values
+            context_lengths = [
+                p.context_length for p in points if p.peak_mem is not None
+            ]
+            memories = [p.peak_mem for p in points if p.peak_mem is not None]
+
+            if context_lengths and memories:
+                ax1.plot(
+                    context_lengths,
+                    memories,
+                    marker="o",
+                    linewidth=2,
+                    markersize=10,
+                    alpha=1.0,
+                    markeredgecolor="black",
+                    markeredgewidth=1,
+                    markerfacecolor=None,
+                    label=method_name,
+                    color=color_map[method_name],
+                )
+
+    ax1.set_xlabel("Haystack Length (tokens)", fontweight="bold", fontsize=20)
+    ax1.set_ylabel("Additional Memory (GB)", fontweight="bold", fontsize=20)
+    ax1.set_title(
+        "Additional Memory Needed for Inference", fontweight="bold", fontsize=24
+    )
+    ax1.set_xscale("log", base=2)
+    # ax1.set_yscale("log")
+    ax1.grid(True, alpha=0.3)
+    # ax1.legend()
+
+    # # Plot 2: Latency vs Context Length
+    # for dataset_name, methods in datasets.items():
+    #     for method_name, points in methods.items():
+    #         # Extract data with valid latency values
+    #         context_lengths = [
+    #             p.context_length for p in points if p.latency is not None
+    #         ]
+    #         latencies = [p.latency for p in points if p.latency is not None]
+    #         latency_stds = [
+    #             p.latency_std if p.latency_std is not None else 0
+    #             for p in points
+    #             if p.latency is not None
+    #         ]
+
+    #         if context_lengths and latencies:
+    #             ax2.errorbar(
+    #                 context_lengths,
+    #                 latencies,
+    #                 yerr=latency_stds,
+    #                 marker="o",
+    #                 linewidth=2,
+    #                 markersize=6,
+    #                 capsize=5,
+    #                 label=method_name,
+    #                 color=color_map[method_name],
+    #             )
+
+    # ax2.set_xlabel("Context Length", fontweight="bold")
+    # ax2.set_ylabel("Latency (seconds)", fontweight="bold")
+    # ax2.set_title("Latency vs Context Length", fontweight="bold", fontsize=16)
+    # ax2.set_xscale("log", base=2)
+    # ax2.grid(True, alpha=0.3)
+    # ax2.legend()
+
+    plt.tight_layout()
+    return fig
+
+
+# Generate the plots
+if __name__ == "__main__":
+    fig = plot_memory_and_latency(datasets)
+    plt.show()
+
+    # Optionally save the figure
+    fig.savefig(
+        "niah_memory_latency_vs_context_length.pdf", dpi=300, bbox_inches="tight"
+    )

tmp/plot_perf_ctx_latency_mem.py

@@ -0,0 +1,786 @@
+from dataclasses import dataclass  # added
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Styling (reused)
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#FFB83D",
+        "#555555",
+        "#9BC750",
+        "#35775A",
+        "#5571EA",
+        "#6CACFF",
+        # "#E84494",
+        # "#E47FB0",
+        # "#785EF0",
+        # "#FF924E",
+        # "#5BC5DB",
+        "#AD6F50",
+        "#A1A9AD",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+plt.rcParams["legend.fontsize"] = 16
+plt.rcParams["xtick.labelsize"] = 13
+plt.rcParams["ytick.labelsize"] = 13
+
+
+@dataclass
+class Point:
+    performance: float
+    context_ratio: float | None = None
+    latency: float | None = None
+    latency_std: float | None = None
+    peak_mem: float | None = None  # added
+    peak_mem_std: float | None = None  # added
+
+
+def combine_std_for_sum(std_list):
+    return np.sqrt(sum(std**2 for std in std_list))
+
+
+# Unified data structure (merged context + latency via Point)
+datasets = {
+    "SQuAD": {
+        "Base model w/ context": [
+            Point(performance=0.8692, context_ratio=1.0),
+        ],
+        "Base model w/o context": [
+            Point(
+                performance=0.1866,
+                context_ratio=0.0,
+                latency=0.0,
+                latency_std=0.0,
+                peak_mem=0.0,
+            ),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.859,
+                context_ratio=0.0,
+                latency=39.974 + 456.989 * 1e-3,
+                latency_std=combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]),
+                peak_mem=855.762 / 1024,
+            ),
+        ],
+        "CD (generated queries)": [
+            Point(
+                performance=0.5986,
+                context_ratio=0.0,
+                latency=21.254 * 4 + 49.358 + 7.450,
+                latency_std=10,
+                peak_mem=41.098,
+            ),
+        ],
+        "Ours (batched)": [
+            Point(
+                # performance=0.7174,
+                performance=0.7077,
+                # context_ratio=0.0,
+                latency=85.954 * 1e-3 + 259.920 * 1e-6,
+                latency_std=combine_std_for_sum([31.057 * 1e-3, 740.286 * 1e-6]),
+                peak_mem=1.151,
+            ),
+        ],
+        "Ours (iterative)": [
+            Point(
+                performance=0.715,
+                context_ratio=0.0,
+                latency=426.849 * 1e-3 + 260.073 * 1e-6,
+                latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]),
+                peak_mem=328.134 / 1024,
+            ),
+        ],
+        "T2L": [
+            Point(
+                performance=0.1761,
+                context_ratio=0.0,
+                latency=35.480 * 1e-3,
+                latency_std=22.193 * 1e-3,
+                peak_mem=1.793,
+            ),
+        ],
+        "LLMLingua-2": [
+            Point(performance=0.8492, context_ratio=0.9),
+            Point(performance=0.8278, context_ratio=0.8),
+            Point(performance=0.7749, context_ratio=0.6),
+            Point(performance=0.6624, context_ratio=0.4),
+            Point(performance=0.4271, context_ratio=0.2),
+            Point(performance=0.3192, context_ratio=0.1),
+        ],
+    },
+    "DROP": {
+        "Base model w/ context": [
+            Point(performance=0.4541, context_ratio=1.0),
+        ],
+        "Base model w/o context": [
+            Point(
+                performance=0.1417,
+                context_ratio=0.0,
+                latency=0.0,
+                latency_std=0.0,
+                peak_mem=0.0,
+            ),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.443,
+                context_ratio=0.0,
+                latency=39.974 + 422.820 * 1e-3,
+                latency_std=combine_std_for_sum([965.326 * 1e-3, 241.419 * 1e-3]),
+                peak_mem=1.381,
+            ),
+        ],
+        "CD (generated queries)": [
+            Point(
+                performance=0.2289,
+                context_ratio=0.0,
+                latency=23.883 * 4 + 9.755 + 41.301,
+                latency_std=14.8852904910,
+                peak_mem=44.305,
+            ),
+        ],
+        "Ours (batched)": [
+            Point(
+                performance=0.2972,
+                # context_ratio=0.0,
+                latency=84.132 * 1e-3 + 251.876 * 1e-6,
+                latency_std=combine_std_for_sum([30.503 * 1e-3, 22.829 * 1e-6]),
+                peak_mem=1.976,
+            ),
+        ],
+        "Ours (iterative)": [
+            Point(
+                # performance=0.3002,
+                performance=0.3387,
+                context_ratio=0.0,
+                latency=419.462 * 1e-3 + 254.899 * 1e-6,
+                latency_std=combine_std_for_sum([40.211 * 1e-3, 29.784 * 1e-6]),
+                peak_mem=429.881 / 1024,
+            ),
+        ],
+        "T2L": [
+            Point(
+                performance=0.1468,
+                context_ratio=0.0,
+                latency=35.036 * 1e-3,
+                latency_std=18.144 * 1e-3,
+                peak_mem=3.004,
+            ),
+        ],
+        "LLMLingua-2": [  #
+            Point(performance=0.4513, context_ratio=0.9),
+            Point(performance=0.4488, context_ratio=0.8),
+            Point(performance=0.4076, context_ratio=0.6),
+            Point(performance=0.3488, context_ratio=0.4),
+            Point(performance=0.2617, context_ratio=0.2),
+            Point(performance=0.1712, context_ratio=0.1),
+        ],
+    },
+    "ROPES": {
+        "Base model w/ context": [
+            Point(performance=0.7457, context_ratio=1.0),
+        ],
+        "Base model w/o context": [
+            Point(
+                performance=0.4545,
+                context_ratio=0.0,
+                latency=0.0,
+                latency_std=0.0,
+                peak_mem=0.0,
+            ),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.7349,
+                context_ratio=0.0,
+                latency=39.633 + 285.467 * 1e-3,
+                latency_std=combine_std_for_sum([1.752, 72.791 * 1e-3]),
+                peak_mem=527.312 / 1024,
+            ),
+        ],
+        "CD (generated queries)": [
+            Point(
+                performance=0.5786,
+                context_ratio=0.0,
+                latency=23.714 * 4 + 11.473 + 59.660,
+                latency_std=11.7329040736,
+                peak_mem=43.271,
+            ),
+        ],
+        "Ours (batched)": [
+            Point(
+                performance=0.6759,
+                # context_ratio=0.0,
+                latency=85.652 * 1e-3 + 254.029 * 1e-6,
+                latency_std=combine_std_for_sum([30.852 * 1e-3, 51.127 * 1e-6]),
+                peak_mem=496.898 / 1024,
+            ),
+        ],
+        "Ours (iterative)": [
+            Point(
+                # performance=0.6786,
+                performance=0.6088,
+                context_ratio=0.0,
+                latency=417.860 * 1e-3 + 257.422 * 1e-6,
+                latency_std=combine_std_for_sum([45.514 * 1e-3, 31.419 * 1e-6]),
+                peak_mem=273.867 / 1024,
+            ),
+        ],
+        "T2L": [
+            Point(
+                performance=0.4661,
+                context_ratio=0.0,
+                latency=37.913 * 1e-3,
+                latency_std=17.313 * 1e-3,
+                peak_mem=697.217 / 1024,
+            ),
+        ],
+        "LLMLingua-2": [
+            Point(performance=0.708, context_ratio=0.9),
+            Point(performance=0.7126, context_ratio=0.8),
+            Point(performance=0.7291, context_ratio=0.6),
+            Point(performance=0.677, context_ratio=0.4),
+            Point(performance=0.6188, context_ratio=0.2),
+            Point(performance=0.4894, context_ratio=0.1),
+        ],
+    },
+}
+
+
+def build_color_map(datasets):
+    methods = []
+    for data in datasets.values():
+        for m in data.keys():
+            if m not in methods:
+                methods.append(m)
+    palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+    return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
+
+
+def plot_ctx_and_latency(datasets):
+    n = len(datasets)
+    # detect memory usage presence
+    memory_present = any(
+        p.peak_mem is not None
+        for methods in datasets.values()
+        for plist in methods.values()
+        for p in plist
+    )
+    num_cols = 3 if memory_present else 2
+    # Share x-axes across rows only when there are multiple datasets (n > 1)
+    if n > 1:
+        fig, axes = plt.subplots(
+            n,
+            num_cols,
+            figsize=(6 * num_cols, 5 * n),
+            sharey=True,
+            sharex="col",
+        )
+    else:
+        fig, axes = plt.subplots(
+            n, num_cols, figsize=(6 * num_cols, 5 * n), sharey=True
+        )
+    if n == 1:
+        axes = np.array([axes])
+
+    color_map = build_color_map(datasets)
+    global_methods = set()
+
+    label_text = "Efficient and Effective\nInternalization"
+    ours_methods = {"Ours (batched)", "Ours (iterative)"}  # added
+    annot_name = {
+        "Ours (batched)": "Ours (b)",
+        "Ours (iterative)": "Ours (i)",
+    }
+
+    # Keep track of axes groups per row for later separator lines (always store as list)
+    row_axes_groups = []
+    for row_idx, (dataset_name, methods) in enumerate(datasets.items()):
+        row_axes = axes[row_idx]
+        if memory_present:
+            ctx_ax, lat_ax, mem_ax = row_axes
+        else:
+            ctx_ax, lat_ax = row_axes
+        # Ensure each row group is a plain Python list of Axes objects
+        if isinstance(row_axes, (list, tuple)):
+            row_axes_groups.append(list(row_axes))
+        elif isinstance(row_axes, np.ndarray):
+            row_axes_groups.append(list(row_axes.ravel()))
+        else:  # single axis fallback (should not occur with current layout)
+            row_axes_groups.append([row_axes])
+        # Row subtitle (dataset name) centered across the row using the middle axis
+        title_ax = lat_ax  # use latency axis (middle) for centering
+        title_ax.set_title(dataset_name, fontweight="bold", fontsize=24, pad=20)
+
+        # Reference performance from ICL
+        ref_perf = methods["Base model w/ context"][0].performance
+        base_wo_ctx_perf_rel = (
+            next(
+                p.performance
+                for p in methods["Base model w/o context"]
+                if p.context_ratio is not None
+            )
+            / ref_perf
+        )
+
+        # --- Context subplot ---
+        ours_ctx_present = [  # added
+            m
+            for m in ours_methods
+            if any(p.context_ratio is not None for p in methods.get(m, []))
+        ]
+        for method, points in methods.items():
+            ctx_points = [p for p in points if p.context_ratio is not None]
+            if not ctx_points:
+                continue
+            xs = [p.context_ratio for p in ctx_points]
+            ys = [p.performance / ref_perf for p in ctx_points]
+            marker = (
+                "D"
+                if method == "CD (oracle)"
+                else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o")
+            )
+            markersize = 17 if "Ours" in method else 12
+
+            if len(xs) == 1:
+                ctx_ax.scatter(
+                    xs,
+                    ys,
+                    marker=marker,
+                    s=markersize**2,  # scatter uses area, so square the size
+                    label=method,
+                    color=color_map[method],
+                    edgecolors="black",
+                    linewidths=1.5,
+                    alpha=0.9,
+                )
+            else:
+                ctx_ax.plot(
+                    xs,
+                    ys,
+                    marker=marker,
+                    linewidth=2,
+                    markersize=markersize,
+                    label=method,
+                    color=color_map[method],
+                    markeredgecolor="black",
+                    markeredgewidth=1.5,
+                    alpha=0.9,
+                )
+            if method in {"Ours (batched)", "Ours (iterative)"}:
+                yoff = (
+                    (10 if method == "Ours (batched)" else -10)
+                    if len(ours_ctx_present) == 2
+                    else 6
+                )
+
+                ctx_ax.annotate(  # fixed to ctx_ax
+                    annot_name[method],
+                    (xs[-1], ys[-1]),
+                    xytext=(8, yoff),
+                    textcoords="offset points",
+                    fontsize=12,
+                    color=color_map[method],
+                    fontweight="bold",
+                    bbox=None,
+                )
+            global_methods.add(method)
+
+        # Add x-axis label on every context subplot (reverted per-user preference)
+        ctx_ax.set_xlabel("Context Length Ratio", fontweight="bold", fontsize=16)
+        # if row_idx == 0:
+        ctx_ax.set_ylabel(
+            "Normalized Performance",
+            fontweight="bold",
+            fontsize=16,
+        )
+        ctx_ax.grid(True, alpha=0.3)
+        ctx_ax.set_xlim(-0.05, 1.05)
+        ctx_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
+        ctx_ax.axhline(
+            base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
+        )
+        # Removed highlight rectangle and label for context subplot
+        # --- Latency subplot ---
+        ours_lat_present = [  # added
+            m
+            for m in ours_methods
+            if any(p.latency is not None for p in methods.get(m, []))
+        ]
+        all_latencies = []
+        for method, points in methods.items():
+            lat_points = [p for p in points if p.latency is not None]
+            if not lat_points:
+                continue
+            lat_x = [p.latency for p in lat_points]
+            perf_rel = [p.performance / ref_perf for p in lat_points]
+            xerr = [p.latency_std for p in lat_points]
+            marker = (
+                "D"
+                if method == "CD (oracle)"
+                else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o")
+            )
+            markersize = 17 if "Ours" in method else 12
+
+            if len(lat_x) == 1 and all(err == 0 for err in xerr):
+                lat_ax.scatter(
+                    lat_x,
+                    perf_rel,
+                    marker=marker,
+                    s=markersize**2,
+                    label=method,
+                    color=color_map[method],
+                    edgecolors="black",
+                    linewidths=1.5,
+                    alpha=0.9,
+                )
+            else:
+                lat_ax.errorbar(
+                    lat_x,
+                    perf_rel,
+                    xerr=xerr,
+                    fmt=marker + "-",
+                    linewidth=2,
+                    markersize=markersize,
+                    capsize=4,
+                    capthick=1,
+                    label=method,
+                    color=color_map[method],
+                    markeredgecolor="black",
+                    markeredgewidth=1.5,
+                    alpha=0.9,
+                )
+            # Add floating annotation for Ours
+            if method in {"Ours (batched)", "Ours (iterative)"}:
+                yoff = (
+                    (-20 if method == "Ours (batched)" else 10)
+                    if len(ours_lat_present) == 2
+                    else 6
+                )
+                lat_ax.annotate(
+                    annot_name[method],
+                    (lat_x[-1], perf_rel[-1]),
+                    xytext=(-25, yoff),
+                    textcoords="offset points",
+                    fontsize=12,
+                    color=color_map[method],
+                    fontweight="bold",
+                    bbox=None,
+                )
+
+            all_latencies.extend(lat_x)
+            global_methods.add(method)
+        lat_ax.set_xscale("log")
+        if all_latencies:
+            positive = [x for x in all_latencies if x > 0]
+            if positive:
+                min_pos = min(positive)
+                max_x = max(all_latencies)
+                lat_ax.set_xlim(min_pos * 0.8, max_x * 1.2)
+        lat_ax.set_xlabel("Update Latency (seconds)", fontweight="bold", fontsize=16)
+        lat_ax.grid(True, which="both", alpha=0.3)
+        lat_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
+        lat_ax.axhline(
+            base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
+        )
+        # Add vertical line at 1 second
+        lat_ax.axvline(1.0, color="gray", linestyle="--", alpha=0.7, zorder=-1)
+        lat_ax.text(
+            1.1,
+            0.5,
+            "Sub-second internalization",
+            transform=lat_ax.get_xaxis_transform(),
+            ha="left",
+            va="center",
+            fontsize=8,
+            color="gray",
+            fontstyle="italic",
+            rotation=90,
+        )
+        # # Highlight rectangle
+        #     mpl.patches.Rectangle(
+        #         (0.0, 0.5),
+        #         0.5,
+        #         0.5,
+        #         transform=lat_ax.transAxes,
+        #         facecolor="green",
+        #         edgecolor="none",
+        #         alpha=0.15,
+        #         zorder=-2,
+        #     )
+        # )
+        # lat_ax.text(
+        #     0.02,
+        #     0.62,
+        #     label_text,
+        #     transform=lat_ax.transAxes,
+        #     fontsize=8.5,
+        #     ha="left",
+        #     va="bottom",
+        #     color="green",
+        #     fontweight="bold",
+        # )
+
+        # --- Memory subplot (new) ---
+        if memory_present:
+            ours_mem_present = [  # added
+                m
+                for m in ours_methods
+                if any(p.peak_mem is not None for p in methods.get(m, []))
+            ]
+            all_mems = []
+            for method, points in methods.items():
+                mem_points = [p for p in points if p.peak_mem is not None]
+                if not mem_points:
+                    continue
+                mem_x = [p.peak_mem for p in mem_points]
+                perf_rel = [p.performance / ref_perf for p in mem_points]
+                xerr = [
+                    p.peak_mem_std if p.peak_mem_std is not None else 0
+                    for p in mem_points
+                ]
+                marker = (
+                    "D"
+                    if method == "CD (oracle)"
+                    else (
+                        "*" if method in {"Ours (batched)", "Ours (iterative)"} else "o"
+                    )
+                )
+                markersize = 17 if "Ours" in method else 12
+
+                if len(mem_x) == 1 and all(err == 0 for err in xerr):
+                    mem_ax.scatter(
+                        mem_x,
+                        perf_rel,
+                        marker=marker,
+                        s=markersize**2,
+                        label=method,
+                        color=color_map[method],
+                        edgecolors="black",
+                        linewidths=1.5,
+                        alpha=0.9,
+                    )
+                else:
+                    mem_ax.errorbar(
+                        mem_x,
+                        perf_rel,
+                        xerr=xerr,
+                        fmt=marker + "-",
+                        linewidth=2,
+                        markersize=markersize,
+                        capsize=4,
+                        capthick=1,
+                        label=method,
+                        color=color_map[method],
+                        markeredgecolor="black",
+                        markeredgewidth=1.5,
+                        alpha=0.9,
+                    )
+                # Add floating annotation for Ours
+                if method in {"Ours (batched)", "Ours (iterative)"}:
+                    yoff = (
+                        (-20 if method == "Ours (batched)" else 15)
+                        if len(ours_mem_present) == 2
+                        else 6
+                    )
+                    mem_ax.annotate(
+                        annot_name[method],
+                        (mem_x[-1], perf_rel[-1]),
+                        xytext=(-15, yoff),
+                        textcoords="offset points",
+                        fontsize=12,
+                        color=color_map[method],
+                        fontweight="bold",
+                        bbox=None,
+                    )
+
+                all_mems.extend(mem_x)
+            if all_mems:
+                max_mem = max(all_mems)
+                min_mem = min(all_mems)
+                span = (
+                    max_mem - min_mem
+                    if max_mem > min_mem
+                    else max_mem
+                    if max_mem != 0
+                    else 1
+                )
+                mem_ax.set_xlim(min_mem - 0.05 * span, max_mem * 1.05)
+            mem_ax.set_xlabel(
+                "Additional Memory Needed\nfor Model Updates (GB)",
+                fontweight="bold",
+                fontsize=16,
+            )
+            mem_ax.grid(True, alpha=0.3)
+            mem_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
+            mem_ax.axhline(
+                base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
+            )
+            # # Highlight rectangle (memory) fixed 20% width in axes coords
+            # mem_ax.add_patch(
+            #     mpl.patches.Rectangle(
+            #         (0.0, 0.6),
+            #         0.4,
+            #         0.4,
+            #         transform=mem_ax.transAxes,
+            #         facecolor="green",
+            #         edgecolor="none",
+            #         alpha=0.15,
+            #         zorder=-2,
+            #     )
+            # )
+            # mem_ax.text(
+            #     0.02,
+            #     0.62,
+            #     label_text,
+            #     transform=mem_ax.transAxes,
+            #     fontsize=8.5,
+            #     ha="left",
+            #     va="bottom",
+            #     color="green",
+            #     fontweight="bold",
+            # )
+
+    # REMOVE old single legend block
+    # (deleted fig.legend(...) that was on the right)
+
+    # Build grouped legends
+    in_context = [
+        m for m in ["Base model w/ context", "LLMLingua-2"] if m in global_methods
+    ]
+    in_param_order = [
+        "CD (oracle)",
+        "CD (generated queries)",
+        "Ours (iterative)",
+        "Ours (batched)",
+        "T2L",
+        "Base model w/o context",
+    ]
+    in_param = [m for m in in_param_order if m in global_methods]
+
+    def make_handles(names):
+        return [
+            mpl.lines.Line2D(
+                [],
+                [],
+                label=m,
+                marker=(
+                    "D"
+                    if m == "CD (oracle)"
+                    else ("*" if m in {"Ours (batched)", "Ours (iterative)"} else "o")
+                ),
+                linestyle="-",
+                linewidth=2,
+                color=plt.rcParams["axes.prop_cycle"].by_key()["color"][
+                    list(build_color_map(datasets).keys()).index(m)
+                    % len(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+                ],
+                markeredgecolor="black",
+                markeredgewidth=1.5,
+                markersize=20 if "Ours" in m else 10,
+            )
+            for m in names
+        ]
+
+    ic_handles = make_handles(in_context)
+    ip_handles = make_handles(in_param)
+
+    # Extra bottom margin for legends
+    # (reduced wspace & hspace to tighten gaps)
+    # Increased hspace to add more vertical separation between dataset rows
+    plt.subplots_adjust(right=0.88, wspace=0.12, hspace=0.8, bottom=0.25)
+
+    # Re-enable x tick labels for all rows (matplotlib hides them for shared axes not on bottom)
+    if n > 1:
+        for ax in axes.reshape(-1):
+            ax.tick_params(labelbottom=True)
+
+    # Set legend title fontweight before creating legends
+    plt.rcParams["legend.title_fontsize"] = 14
+
+    fig.legend(
+        ic_handles,
+        [h.get_label() for h in ic_handles],
+        title="In-context knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.25, 0.1),
+        ncol=len(ic_handles),
+        frameon=True,
+        fancybox=True,
+        fontsize=12,
+        title_fontproperties={"weight": "bold", "size": 14},
+    )
+    fig.legend(
+        ip_handles,
+        [h.get_label() for h in ip_handles],
+        title="In-parameter knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.6, 0.1),
+        ncol=min(len(ip_handles), 3),
+        frameon=True,
+        fancybox=True,
+        fontsize=12,
+        title_fontproperties={"weight": "bold", "size": 14},
+    )
+
+    # Removed global suptitle referencing only the last dataset; per-row titles now provided.
+
+    # Add horizontal separator lines between dataset rows (if more than one dataset)
+    if len(row_axes_groups) > 1:
+        # Use figure coordinate system
+        for upper_group, lower_group in zip(row_axes_groups[:-1], row_axes_groups[1:]):
+            # Reference axes (first axis in each group)
+            upper_ax = upper_group[0]
+            lower_ax = lower_group[0]
+            upper_bottom = upper_ax.get_position().y0
+            lower_top = lower_ax.get_position().y1
+            line_y = (upper_bottom + lower_top) / 2.0
+            combined = upper_group + lower_group
+            x0 = min(ax.get_position().x0 for ax in combined)
+            x1 = max(ax.get_position().x1 for ax in combined)
+            fig.add_artist(
+                mpl.lines.Line2D(
+                    [x0, x1],
+                    [line_y, line_y],
+                    transform=fig.transFigure,
+                    color="#cccccc",
+                    linewidth=1.2,
+                    alpha=0.8,
+                )
+            )
+    return fig
+
+
+if __name__ == "__main__":
+    fig = plot_ctx_and_latency(datasets)
+    plt.show()
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.png",
+        dpi=300,
+        bbox_inches="tight",
+    )
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.pdf",
+        dpi=300,
+        bbox_inches="tight",
+    )

tmp/plot_perf_gen_mem.py

@@ -0,0 +1,453 @@
+from dataclasses import dataclass
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Styling (mirroring the other plot for visual consistency)
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#FFB83D",
+        "#EEE684",
+        "#555555",
+        "#9BC750",
+        "#35775A",
+        "#5571EA",
+        "#6CACFF",
+        # "#E84494",
+        # "#E47FB0",
+        # "#785EF0",
+        # "#FF924E",
+        # "#5BC5DB",
+        "#AD6F50",
+        "#A1A9AD",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+# plt.rcParams["legend.fontsize"] = 20
+plt.rcParams["xtick.labelsize"] = 20
+plt.rcParams["ytick.labelsize"] = 20
+
+
+@dataclass
+class Point:
+    performance: float
+    peak_mem: float | None = None  # in GB
+
+
+# Dataset focused only on memory vs performance (relative performance computed later)
+datasets = {
+    "2WikiMultihopQA": {
+        # We keep Base model w/ context only as a performance reference (no point plotted if memory absent)
+        "Base model w/ context": [
+            Point(performance=0.242, peak_mem=2.099 * 1024),
+        ],
+        "Base model w/ truncated context": [
+            Point(performance=0.3387, peak_mem=1.194 * 1024),
+        ],
+        "Base model w/o context": [
+            Point(performance=0.2138, peak_mem=65.162),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.305,
+                peak_mem=18.196,
+            ),
+        ],
+        "CD (5 generated queries)": [
+            Point(
+                performance=0.2385,
+                peak_mem=32.501,
+            ),
+        ],
+        "Ours": [
+            Point(
+                performance=0.2901,
+                peak_mem=35.777,
+            ),
+        ],
+        "Ours + truncated context": [
+            Point(
+                performance=0.3585,
+                peak_mem=1.367 * 1024,
+            )
+        ],
+        # "Ours (8k chunks)": [
+        #     Point(
+        #         performance=0.2901,
+        #         peak_mem=35.777,
+        #     ),
+        # ],
+        "T2L": [
+            Point(
+                performance=0.2474,
+                peak_mem=37.776,
+            ),
+        ],
+        "LLMLingua-2": [
+            # Point(performance=0.3141, peak_mem=1.111 * 1024),
+            Point(performance=0.3169, peak_mem=1.018 * 1024),
+            Point(performance=0.3558, peak_mem=807.366),
+            Point(performance=0.3591, peak_mem=592.584),
+            Point(performance=0.2767, peak_mem=305.298),
+            Point(performance=0.2343, peak_mem=158.760),
+        ],
+    },
+    "MultiFieldQA": {
+        "Base model w/ context": [
+            Point(performance=0.303, peak_mem=1.211 * 1024),
+        ],
+        "Base model w/ truncated context": [
+            Point(performance=0.3938, peak_mem=1.084 * 1024),
+        ],
+        "Base model w/o context": [
+            Point(performance=0.0948, peak_mem=65.853),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.4101,
+                peak_mem=7.836,
+            ),
+        ],
+        "CD (5 generated queries)": [
+            Point(
+                performance=0.165,
+                peak_mem=31.587,
+            ),
+        ],
+        "Ours": [
+            Point(
+                performance=0.1895,
+                peak_mem=35.769,
+            ),
+        ],
+        "Ours + truncated context": [
+            Point(
+                performance=0.4293,
+                peak_mem=1.233 * 1024,
+            ),
+        ],
+        "T2L": [
+            Point(
+                performance=0.083,
+                peak_mem=38.907,
+            ),
+        ],
+        "LLMLingua-2": [
+            Point(performance=0.4222, peak_mem=909.663),
+            Point(performance=0.3898, peak_mem=716.981),
+            Point(performance=0.3773, peak_mem=490.478),
+            Point(performance=0.2885, peak_mem=256.169),
+            Point(performance=0.2298, peak_mem=146.394),
+        ],
+    },
+    "QASPER": {
+        # We keep Base model w/ context only as a performance reference (no point plotted if memory absent)
+        "Base model w/ context": [
+            Point(performance=0.346, peak_mem=2.217 * 1024),
+        ],
+        "Base model w/ truncated context": [
+            Point(performance=0.3839, peak_mem=788.697),
+        ],
+        "Base model w/o context": [
+            Point(performance=0.0859, peak_mem=64.989),
+        ],
+        "CD (oracle)": [
+            Point(
+                performance=0.3627,
+                peak_mem=19.611,
+            ),
+        ],
+        "CD (5 generated queries)": [
+            Point(
+                performance=0.1712,
+                peak_mem=36.767,
+            ),
+        ],
+        "Ours": [
+            Point(
+                performance=0.2051,
+                peak_mem=46.130,
+            ),
+        ],
+        "Ours + truncated context": [
+            Point(
+                performance=0.3751,
+                peak_mem=873.099,
+            ),
+        ],
+        "T2L": [
+            Point(
+                performance=0.1259,
+                peak_mem=54.176,
+            ),
+        ],
+        "LLMLingua-2": [
+            # Point(performance=0.3734, peak_mem=710.176),
+            Point(performance=0.3849, peak_mem=647.385),
+            Point(performance=0.3838, peak_mem=487.781),
+            Point(performance=0.3328, peak_mem=386.447),
+            Point(performance=0.3329, peak_mem=219.062),
+            Point(performance=0.2614, peak_mem=132.576),
+        ],
+    },
+}
+
+
+def build_color_map(datasets):
+    # Collect all unique methods in a deterministic order
+    methods = []
+    for data in datasets.values():
+        for m in data.keys():
+            if m not in methods:
+                methods.append(m)
+
+    # Sort methods to ensure consistent ordering across runs
+    # methods = sorted(methods)
+
+    palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+    return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
+
+
+def plot_memory_vs_performance(datasets):
+    n = len(datasets)
+    # If multiple datasets, arrange horizontally; otherwise single standard figure
+    if n == 1:
+        fig, single_ax = plt.subplots(1, 1, figsize=(7, 5))
+        axes = [single_ax]
+    else:
+        fig, axes_arr = plt.subplots(1, n, figsize=(7 * n, 5))
+        # Ensure axes is a simple list
+        axes = list(np.atleast_1d(axes_arr))
+
+    color_map = build_color_map(datasets)
+    global_methods = set()
+    ours_methods = {"Ours", "Ours + truncated context"}  # added
+    annot_name = {  # added
+        "Ours": "Ours",
+        "Ours + truncated context": "Ours+ctx",
+    }
+
+    highlight_label = "Efficient and Effective\nInternalization"
+
+    for idx, (dataset_name, methods) in enumerate(datasets.items()):
+        ax = axes[idx]
+        ref_perf = methods["Base model w/ truncated context"][0].performance
+        base_wo_ctx_perf_rel = (
+            methods["Base model w/o context"][0].performance / ref_perf
+            if "Base model w/o context" in methods
+            else None
+        )
+
+        all_mems = []
+        all_rel_perfs = []
+        ours_present = [  # added: detect both ours points for staggering
+            m
+            for m in ours_methods
+            if m in methods and any(p.peak_mem is not None for p in methods[m])
+        ]
+        for method, points in methods.items():
+            mem_points = [p for p in points if p.peak_mem is not None]
+            if not mem_points:
+                continue  # skip reference method if no memory value
+            xs = [p.peak_mem for p in mem_points]
+            ys = [p.performance / ref_perf for p in mem_points]
+            # marker mapping (oracle=diamond, ours=star, others=circle)
+            marker = (
+                "D"
+                if method == "CD (oracle)"
+                else ("*" if method in ours_methods else "o")
+            )
+            size = 20 if "Ours" in method else 12
+            ax.plot(
+                xs,
+                ys,
+                linestyle="-",
+                marker=marker,
+                linewidth=2,
+                markersize=size,
+                label=method,
+                color=color_map[method],
+                markeredgecolor="black",
+                markeredgewidth=1.5,
+                alpha=0.9,
+            )
+            # Add annotation for Ours, stagger if both are present
+            if method in ours_methods:
+                yoff = (
+                    (10 if method == "Ours" else -10) if len(ours_present) == 2 else 6
+                )
+                place_left = method == "Ours + truncated context" and dataset_name in {
+                    "2WikiMultihopQA",
+                    "MultiFieldQA",
+                }
+                if place_left:
+                    yoff = 10  # move annotation above the marker for these two plots
+                xoff = -5 if place_left else 8
+                ha = "right" if place_left else "left"
+                ax.annotate(
+                    annot_name[method],
+                    (xs[-1], ys[-1]),
+                    xytext=(xoff, yoff),
+                    textcoords="offset points",
+                    fontsize=12,
+                    color=color_map[method],
+                    fontweight="bold",
+                    ha=ha,
+                    bbox=None,
+                )
+
+            all_mems.extend(xs)
+            all_rel_perfs.extend(ys)
+            global_methods.add(method)
+
+        if all_mems:
+            max_mem = max(all_mems)
+            min_mem = min(all_mems)
+            span = (
+                (max_mem - min_mem)
+                if max_mem > min_mem
+                else (max_mem if max_mem != 0 else 1)
+            )
+            ax.set_xlim(5, max_mem * 1.25)
+
+        # Set individual y-axis limits for each subplot
+        if all_rel_perfs:
+            min_y = min(all_rel_perfs)
+            max_y = max(all_rel_perfs)
+            y_range = max_y - min_y
+            margin = y_range * 0.1 if y_range > 0 else 0.1
+            ax.set_ylim(min_y - margin, max_y + margin)
+
+        ax.set_xscale("log")
+        ax.set_xlabel(
+            "Additional Memory Needed\nfor Generation (MB)",
+            fontweight="bold",
+            fontsize=22,
+        )
+        if idx == 0:  # Only show y-label on first plot
+            ax.set_ylabel("Normalized Performance", fontweight="bold", fontsize=22)
+        ax.grid(True, alpha=0.3)
+        ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
+        if base_wo_ctx_perf_rel is not None:
+            ax.axhline(
+                base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
+            )
+
+        ax.set_title(dataset_name, fontweight="bold", fontsize=30)
+
+    # Build grouped legends
+    in_context = [
+        m
+        for m in [
+            "Base model w/ context",
+            "Base model w/ truncated context",
+            "LLMLingua-2",
+        ]
+        if m in global_methods
+    ]
+    in_param_order = [
+        "CD (oracle)",
+        "CD (5 generated queries)",
+        "Ours",
+        "Ours + truncated context",
+        "T2L",
+        "Base model w/o context",
+    ]
+    in_param = [m for m in in_param_order if m in global_methods]
+
+    def make_handles(names):
+        return [
+            mpl.lines.Line2D(
+                [],
+                [],
+                label=m,
+                marker=(
+                    "D"
+                    if m == "CD (oracle)"
+                    else ("*" if m in {"Ours", "Ours + truncated context"} else "o")
+                ),
+                linestyle="-",
+                linewidth=2,
+                color=color_map[m],
+                markeredgecolor="black",
+                markeredgewidth=1.5,
+                markersize=22 if "Ours" in m else 14,
+            )
+            for m in names
+        ]
+
+    ic_handles = make_handles(in_context)
+    ip_handles = make_handles(in_param)
+
+    # Extra bottom margin for legends
+    # (reduced wspace & hspace to tighten gaps)
+    # Increased hspace to add more vertical separation between dataset rows
+    plt.subplots_adjust(right=0.88, wspace=0.12, hspace=1.0, bottom=0.25)
+
+    # Set legend title fontweight before creating legends
+    # plt.rcParams["legend.title_fontsize"] = 14
+
+    fig.legend(
+        ic_handles,
+        [h.get_label() for h in ic_handles],
+        title="In-context knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.2, -0.05),
+        ncol=2,
+        frameon=True,
+        fancybox=True,
+        fontsize=18,
+        title_fontproperties={"weight": "bold", "size": 20},
+    )
+    fig.legend(
+        ip_handles,
+        [h.get_label() for h in ip_handles],
+        title="In-parameter knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.7, -0.05),
+        ncol=min(len(ip_handles), 3),
+        frameon=True,
+        fancybox=True,
+        fontsize=18,
+        title_fontproperties={"weight": "bold", "size": 20},
+    )
+
+    # Adjust spacing: horizontal layout uses wspace, single layout unaffected
+    if n == 1:
+        plt.subplots_adjust(bottom=0.18, left=0.12, right=0.98)
+    else:
+        # Tighter horizontal spacing and margins
+        plt.subplots_adjust(bottom=0.18, left=0.06, right=0.99, wspace=0.15)
+    return fig
+
+
+if __name__ == "__main__":
+    fig = plot_memory_vs_performance(datasets)
+    plt.show()
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.png",
+        dpi=300,
+        bbox_inches="tight",
+    )
+    fig.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.pdf",
+        dpi=300,
+        bbox_inches="tight",
+    )

tmp/plot_perf_vs_ctx_len.py

@@ -0,0 +1,272 @@
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#648FFF",
+        "#E84494",
+        "#9BC750",
+        # "#DC4BDC",
+        "#229487",
+        "#785EF0",
+        "#FF924E",
+        "#5BC5DB",
+        "#F0434F",
+        "#AD6F50",
+        # "#1982C4",
+        # "#6A4C93",
+        # "#56CCF2",
+        # "#2F52E0",
+        # "#FF006E",
+        # "#FF8A00",
+        # "#FFD300",
+        # "#00A6ED",
+        # "#FF7A5A",
+        # "#FFAA92",
+        # "#FFB7B2",
+        # "#FFDAC1",
+        # "#E2F0CB",
+        # "#B5EAD7",
+        # "#C7CEEA",
+        # "#D4A5A5",
+        # "#E9B0DF",
+        # "#F8D3DA",
+        # "#FBE0C3",
+        # "#FFF0B8",
+        # "#E6E2AF",
+        # "#A2D2FF",
+        # "#B4F8C8",
+        # "#FFD6E0",
+        # "#FFEFB5",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+
+# Reproducibility
+rng = np.random.default_rng(42)
+
+# X-axis: context length as a fraction of full ICL
+x_inparam = np.array([0.0])  # in-parameter updates
+x_icl = np.array([1.0])  # full ICL
+
+# Define mock performances (F1) per dataset, consistent with the paper narrative:
+# - ICL highest
+# - Compression approaches ICL as ratio increases
+
+datasets = {
+    "SQuAD": {
+        "Base model w/ ICL": [(1.0, 0.8692)],
+        "Base model w/o context": [(0.0, 0.1866)],
+        "CD (oracle)": [(0.0, 0.859)],
+        "CD (generated queries)": [(0.0, 0.5986)],
+        "Ours": [(0.0, 0.7174)],
+        "T2L": [(0.0, 0.1761)],
+        "LLMLingua-2": [
+            (0.9, 0.8492),
+            (0.8, 0.8278),
+            (0.6, 0.7749),
+            (0.4, 0.6624),
+            (0.2, 0.4271),
+            (0.1, 0.3192),
+        ],
+    },
+    # "DROP": {
+    #     "Base model w/ ICL": [(1.0, 0.4541)],
+    #     "Base model w/o context": [(0.0, 0.1417)],
+    #     "CD (oracle)": [(0.0, 0.443)],
+    #     "CD (generated queries)": [(0.0, X)],
+    #     "Ours": [(0.0, 0.7174)],
+    #     "T2L": [(0.0, 0.1761)],
+    #     "LLMLingua-2": [
+    #         (0.9, 0.8492),
+    #         (0.8, 0.8278),
+    #         (0.6, 0.7749),
+    #         (0.4, 0.6624),
+    #         (0.2, 0.4271),
+    #         (0.1, 0.3192),
+    #     ],
+    # },
+}
+
+
+def plot_performance_comparison(datasets):
+    """Plot performance of methods relative to Base model w/ ICL reference."""
+    n_datasets = len(datasets)
+    fig, axes = plt.subplots(1, n_datasets, figsize=(8 * n_datasets, 6))
+
+    # Handle single dataset case
+    if n_datasets == 1:
+        axes = [axes]
+
+    global_methods = set()  # collect methods for global legends
+
+    for idx, (dataset_name, methods) in enumerate(datasets.items()):
+        ax = axes[idx]
+
+        # Get reference performance from "Base model w/ ICL"
+        ref_performance = methods["Base model w/ ICL"][0][1]
+
+        # Get base model w/o context performance for reference line
+        no_context_performance = (
+            methods["Base model w/o context"][0][1] / ref_performance
+        )
+
+        # Define method groups
+        in_context_methods = {"Base model w/ ICL", "LLMLingua-2"}
+        in_param_methods = {
+            "Base model w/o context",
+            "Ours",
+            "T2L",
+            "CD (oracle)",
+            "CD (generated queries)",
+        }
+
+        # Plot each method
+        for method_name, data_points in methods.items():
+            x_vals = [point[0] for point in data_points]
+            y_vals = [
+                point[1] / ref_performance for point in data_points
+            ]  # Relative to reference
+
+            ax.plot(
+                x_vals,
+                y_vals,
+                label=method_name,
+                linewidth=2,
+                markersize=8,
+                marker="o",
+                alpha=1,
+                markeredgecolor="black",
+                markeredgewidth=1,
+                markerfacecolor=None,
+            )
+            global_methods.add(method_name)
+
+        # Formatting
+        ax.set_xlabel("Context Length Ratio vs ICL", fontweight="bold")
+        ax.set_ylabel("Relative Performance vs ICL)", fontweight="bold")
+        ax.grid(True, alpha=0.3)
+        ax.set_xlim(-0.05, 1.05)
+
+        # Add horizontal line at y=1 (reference performance)
+        ax.axhline(
+            y=1.0,
+            color="gray",
+            linestyle="--",
+            alpha=0.5,
+            zorder=-1,
+        )
+
+        # Add horizontal line for base model w/o context
+        ax.axhline(
+            y=no_context_performance,
+            color="red",
+            linestyle=":",
+            alpha=0.7,
+            zorder=-1,
+        )
+
+        # Add highlight rectangle: fixed 20% of axis width (axes coords), y from 0.8 to 1.0
+        ax.add_patch(
+            mpl.patches.Rectangle(
+                (0.0, 0.8),
+                0.2,
+                0.2,
+                transform=ax.transAxes,
+                facecolor="green",
+                edgecolor="none",
+                alpha=0.15,
+                zorder=-2,
+            )
+        )
+
+    # Global legends (after all subplots)
+    in_context = [
+        m for m in ["Base model w/ ICL", "LLMLingua-2"] if m in global_methods
+    ]
+    in_param_order = [
+        "CD (oracle)",
+        "CD (generated queries)",
+        "Ours",
+        "T2L",
+        "Base model w/o context",
+    ]
+    in_param = [m for m in in_param_order if m in global_methods]
+
+    def make_handles(names):
+        return [
+            mpl.lines.Line2D(
+                [],
+                [],
+                label=m,
+                marker="o",
+                linestyle="-",
+                linewidth=2,
+                markeredgecolor="black",
+                markeredgewidth=1,
+                markersize=8,
+            )
+            for m in names
+        ]
+
+    ic_handles = make_handles(in_context)
+    ip_handles = make_handles(in_param)
+
+    # Adjust bottom for legends
+    plt.subplots_adjust(right=0.75, bottom=0.22)
+
+    # Place legends below figure, centered in two groups
+    fig.legend(
+        ic_handles,
+        [h.get_label() for h in ic_handles],
+        title="In-context knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.25, -0.05),
+        ncol=len(ic_handles),
+        frameon=True,
+        fancybox=True,
+        fontsize=10,
+        title_fontsize=11,
+    )
+    fig.legend(
+        ip_handles,
+        [h.get_label() for h in ip_handles],
+        title="In-parameter knowledge",
+        loc="upper center",
+        bbox_to_anchor=(0.72, -0.05),
+        ncol=len(ip_handles),
+        frameon=True,
+        fancybox=True,
+        fontsize=10,
+        title_fontsize=11,
+    )
+    return fig
+
+
+# Create the plot
+fig = plot_performance_comparison(datasets)
+plt.show()
+
+# Optionally save the figure
+plt.savefig(
+    "/home/tan/research/ctx-to-lora/tmp/performance_comparison.png",
+    dpi=300,
+    bbox_inches="tight",
+)

tmp/plot_perf_vs_update_latency.py

@@ -0,0 +1,204 @@
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#648FFF",
+        "#E84494",
+        "#9BC750",
+        "#229487",
+        "#785EF0",
+        "#FF924E",
+        "#5BC5DB",
+        "#F0434F",
+        "#AD6F50",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+plt.rcParams["legend.fontsize"] = 16
+plt.rcParams["xtick.labelsize"] = 13
+plt.rcParams["ytick.labelsize"] = 13
+
+
+# Reproducibility
+rng = np.random.default_rng(42)
+
+# Define mock performances vs update latency (in seconds)
+# X-axis: update latency in seconds (log scale)
+# Y-axis: F1 performance
+
+
+def combine_std_for_sum(std_list):
+    """
+    Compute combined standard deviation when adding independent random variables.
+
+    For independent random variables X1, X2, ..., Xn:
+    Var(X1 + X2 + ... + Xn) = Var(X1) + Var(X2) + ... + Var(Xn)
+    Therefore: std(sum) = sqrt(std1² + std2² + ... + stdn²)
+
+    Args:
+        std_list: List of standard deviations
+
+    Returns:
+        Combined standard deviation for the sum
+    """
+    return np.sqrt(sum(std**2 for std in std_list))
+
+
+datasets = {
+    "SQuAD": {
+        "Base model w/o context": [(0.0, 0.0, 0.1866)],  # No update needed
+        "CD (oracle)": [
+            (
+                39.974 + 456.989 * 1e-3,
+                combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]),
+                0.859,
+            )
+        ],
+        "CD (generated queries)": [(21.254 * 4 + 49.358 + 7.450, 10, 0.5986)],
+        "Ours": [(85.954 * 1e-3, 31.057 * 1e-3, 0.7174)],
+        "T2L": [(35.480 * 1e-3, 22.193 * 1e-3, 0.1761)],
+    },
+    # "DROP": {
+    #     "Base model w/o context": [(0.0, 0.1417)],
+    #     "CD (oracle)": [(300, 0.443)],
+    #     "CD (generated queries)": [(600, 0.3986)],
+    #     "Ours": [(120, 0.4174)],
+    #     "T2L": [(180, 0.1561)],
+    # },
+}
+
+
+def plot_performance_comparison(datasets):
+    """Plot performance of methods vs update latency."""
+    n_datasets = len(datasets)
+    fig, axes = plt.subplots(1, n_datasets, figsize=(8 * n_datasets, 6))
+
+    # Handle single dataset case
+    if n_datasets == 1:
+        axes = [axes]
+
+    for idx, (dataset_name, methods) in enumerate(datasets.items()):
+        ax = axes[idx]
+
+        # Use a fixed reference performance for ICL (since we removed it from the plot)
+        ref_performance = 0.8692 if dataset_name == "SQuAD" else 0.4541
+
+        # Get base model w/o context performance for reference line
+        no_context_performance = (
+            methods["Base model w/o context"][0][2] / ref_performance
+        )
+
+        # All methods are now in-parameter methods
+        in_param_methods = {
+            "Base model w/o context",
+            "Ours",
+            "T2L",
+            "CD (oracle)",
+            "CD (generated queries)",
+        }
+
+        # Plot each method
+        for method_name, data_points in methods.items():
+            x_vals = [point[0] for point in data_points]
+            y_vals = [
+                point[2] / ref_performance for point in data_points
+            ]  # Relative to reference (third element)
+            std_vals = [
+                point[1] / ref_performance for point in data_points
+            ]  # Standard deviation (second element)
+
+            ax.errorbar(
+                x_vals,
+                y_vals,
+                xerr=std_vals,
+                label=method_name,
+                linewidth=2,
+                markersize=8,
+                marker="o",
+                alpha=1,
+                markeredgecolor="black",
+                markeredgewidth=1,
+                markerfacecolor=None,
+                capsize=4,
+                capthick=1,
+            )
+
+        # Calculate max x value for dynamic xlim
+        max_x = max(
+            point[0] for data_points in methods.values() for point in data_points
+        )
+
+        # Formatting
+        ax.set_xlabel("Update Latency (seconds)", fontweight="bold")
+        ax.set_ylabel("Relative Performance vs ICL", fontweight="bold")
+        ax.set_title(f"{dataset_name}", fontweight="bold", fontsize=14)
+        ax.grid(True, alpha=0.3)
+        ax.set_xlim(-5, max_x * 1.2)
+
+        # Add horizontal line at y=1 (reference performance)
+        ax.axhline(
+            y=1.0,
+            color="gray",
+            linestyle="--",
+            alpha=0.5,
+            zorder=-1,
+        )
+
+        # Add horizontal line for base model w/o context
+        ax.axhline(
+            y=no_context_performance,
+            color="red",
+            linestyle="--",
+            alpha=0.7,
+            zorder=-1,
+        )
+
+        # Create single legend for in-parameter methods
+        handles, labels = ax.get_legend_handles_labels()
+
+        ax.legend(
+            handles,
+            labels,
+            title="In-parameter knowledge",
+            bbox_to_anchor=(1.05, 0.6),
+            loc="upper left",
+            ncol=1,
+            fontsize=10,
+            title_fontsize=11,
+            frameon=True,
+            fancybox=True,
+        ).get_title().set_fontweight("bold")
+
+    plt.tight_layout()
+    plt.subplots_adjust(right=0.75)  # Make room for legend on the right
+    return fig
+
+
+# Create the plot
+fig = plot_performance_comparison(datasets)
+plt.show()
+
+# Optionally save the figure
+plt.savefig(
+    "/home/tan/research/ctx-to-lora/tmp/performance_vs_latency.png",
+    dpi=300,
+    bbox_inches="tight",
+)

tmp/test_ce.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""
+Test cross entropy loss shape with reduction=None
+"""
+
+import torch
+from torch import nn
+
+# Create sample data
+batch_size, seq_len, vocab_size = 2, 5, 1000
+logits = torch.randn(batch_size, seq_len, vocab_size)
+labels = torch.randint(0, vocab_size, (batch_size, seq_len))
+
+# Flatten for cross entropy (expects 2D input)
+logits_flat = logits.view(-1, vocab_size)  # (batch_size * seq_len, vocab_size)
+labels_flat = labels.view(-1)  # (batch_size * seq_len,)
+
+print(f"Input shapes:")
+print(f"  logits_flat: {logits_flat.shape}")
+print(f"  labels_flat: {labels_flat.shape}")
+
+loss = nn.functional.cross_entropy(logits_flat, labels_flat, reduction="none")
+
+print(f"Output shape:")
+print(f"  loss: {loss.shape}")
+print(f"  loss values (first 10): {loss[:10]}")

tmp/test_distillation_sparse_vs_dense.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python
+"""
+Compare the current dense KL (teacher top-k scatter into full vocab) vs a sparse gather-based
+implementation used in distillation.
+
+Metrics:
+- Exact numerical equality (max |diff|) across many random seeds.
+- Runtime (per iteration, averaged).
+- Peak memory (GPU if available, else best-effort RSS delta on CPU).
+
+By default we mimic the current code semantics (NO renormalization of the provided teacher
+log-probs). That is, we treat absent vocab items as probability 0 and do:
+  loss_i = - sum_j exp(target_logp[i,j]) * log q(indices[i,j])
+which matches the existing dense scatter approach that sets -inf everywhere else.
+
+We also optionally show (flag --normalized) a variant where teacher log-probs are re-normalized
+over the provided top-k indices (useful if you want a true cross-entropy on the partial dist).
+
+Usage (from repo root):
+  python tmp/test_distillation_sparse_vs_dense.py --seeds 100 --n 2048 --vocab 32000 --k 8 --device cuda
+
+Adjust sizes to fit your GPU memory. The dense method allocates an (N, V) float tensor twice
+(outputs_logits + teacher_logp). For large (N,V) this can OOM.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import statistics
+import time
+import tracemalloc
+from dataclasses import dataclass
+
+import torch
+
+try:
+    import psutil  # optional for CPU RSS
+except ImportError:  # pragma: no cover
+    psutil = None
+
+
+def dense_loss(
+    outputs_logits: torch.Tensor, indices: torch.Tensor, target_logp: torch.Tensor
+) -> torch.Tensor:
+    """Replicates current implementation: build full teacher_logp filled with -inf, scatter top-k logp.
+    outputs_logits: (N, V)
+    indices: (N, K) long
+    target_logp: (N, K) float (NOT assumed normalized)
+    Returns: (N,) per-position losses.
+    """
+    N, V = outputs_logits.shape
+    # Create full teacher log-probs
+    teacher_logp = torch.full_like(outputs_logits, -torch.inf)
+    teacher_logp.scatter_(1, indices, target_logp)
+    p = teacher_logp.exp()  # zeros for -inf
+    # logq = torch.log_softmax(outputs_logits, dim=-1)
+    logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True)  # (N,1)
+    logq = outputs_logits - logq_full_denom
+    loss = -(p * logq).sum(dim=-1)  # (N,)
+    return loss
+
+
+def sparse_loss(
+    outputs_logits: torch.Tensor,
+    indices: torch.Tensor,
+    target_logp: torch.Tensor,
+    *,
+    normalize=False,
+) -> torch.Tensor:
+    """Sparse equivalent to dense_loss without constructing a full (N,V) teacher matrix.
+    If normalize=True we renormalize teacher log-probs over provided indices.
+    Returns (N,) per-position losses with SAME semantics as dense_loss when normalize=False.
+    """
+    # Gather student log-probs only at teacher indices.
+    logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True)  # (N,1)
+    selected_logits = outputs_logits.gather(1, indices)  # (N,K)
+    logq_selected = selected_logits - logq_full_denom  # log softmax at selected indices
+
+    if normalize:
+        # True cross-entropy over provided subset.
+        teacher_logp_norm = target_logp - torch.logsumexp(
+            target_logp, dim=-1, keepdim=True
+        )
+        teacher_p = teacher_logp_norm.exp()
+    else:
+        teacher_p = target_logp.exp()
+
+    # NOT EFFICIENT?
+    # logq = torch.nn.functional.log_softmax(outputs_logits, dim=-1)
+    # logq_selected = logq.gather(1, indices)  # (N,K)
+    # teacher_p = target_logp.exp()
+
+    loss = -(teacher_p * logq_selected).sum(dim=-1)
+    return loss
+
+
+@dataclass
+class Metrics:
+    time_ms: float
+    peak_mem_mb: float | None
+    max_abs_diff: float | None = None
+
+
+def measure(
+    method_fn, *args, repeat=1, device="cpu", measure_memory=True, **kwargs
+) -> tuple[torch.Tensor, Metrics]:
+    """Run method_fn (returns tensor), measure avg wall time and peak memory.
+    For GPU memory we use torch.cuda.reset_peak_memory_stats / max_memory_allocated.
+    For CPU fallback we sample RSS before/after (approx)."""
+    is_cuda = device.startswith("cuda") and torch.cuda.is_available()
+
+    if is_cuda:
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+    else:
+        if measure_memory and psutil is not None:
+            process = psutil.Process(os.getpid())
+            rss_before = process.memory_info().rss
+        elif measure_memory:
+            tracemalloc.start()
+
+    start = time.perf_counter()
+    out = None
+    for _ in range(repeat):
+        out = method_fn(*args, **kwargs)
+    if is_cuda:
+        torch.cuda.synchronize()
+    elapsed = (time.perf_counter() - start) / repeat * 1000.0
+
+    peak_mb = None
+    if is_cuda:
+        peak_mb = torch.cuda.max_memory_allocated(device=device) / (1024**2)
+    else:
+        if measure_memory and psutil is not None:
+            rss_after = process.memory_info().rss
+            peak_mb = (rss_after - rss_before) / (1024**2)
+        elif measure_memory:
+            current, peak = tracemalloc.get_traced_memory()
+            tracemalloc.stop()
+            peak_mb = peak / (1024**2)
+
+    return out, Metrics(time_ms=elapsed, peak_mem_mb=peak_mb)
+
+
+def run_experiment(args):
+    device = args.device
+    torch.manual_seed(0)
+
+    # Pre-generate shapes once per seed to reduce variance from dynamic graph allocations.
+    dense_times = []
+    dense_mems = []
+    sparse_times = []
+    sparse_mems = []
+    diffs = []
+    allclose_fail_seeds = []
+
+    N = args.n
+    V = args.vocab
+    K = args.k
+
+    if K > V:
+        raise ValueError("k cannot exceed vocab size")
+
+    dtype = torch.float16 if args.fp16 else torch.float32
+
+    for seed in range(args.seeds):
+        torch.manual_seed(seed)
+        # Student logits (simulate model outputs for label positions)
+        outputs_logits = 1 + 10 * torch.randn(N, V, device=device, dtype=dtype)
+        # Ensure uniqueness per row: sample without replacement via randperm
+        indices = torch.stack([torch.randperm(V, device=device)[:K] for _ in range(N)])
+        # Teacher log-probs (could be any real numbers). Simulate by sampling logits then subtract logsumexp to get log-probs,
+        # then optionally add random temperature scaling so they are not normalized (matching current dense semantics).
+        teacher_logits = 1 + 10 * torch.randn(
+            N, K, device=device, dtype=torch.float32
+        )  # keep higher precision for stability
+        target_logp = (
+            teacher_logits - torch.logsumexp(teacher_logits, dim=-1, keepdim=True)
+        ).to(dtype)
+
+        # Measure dense
+        dense_out, dense_metrics = measure(
+            dense_loss, outputs_logits, indices, target_logp, device=device
+        )
+        # Measure sparse (semantics-matching version)
+        sparse_out, sparse_metrics = measure(
+            sparse_loss,
+            outputs_logits,
+            indices,
+            target_logp,
+            device=device,
+            normalize=args.normalize_sparse,
+        )
+
+        # Compare
+        # Move to float32 for diff to reduce underflow when fp16
+        diff = (dense_out.float() - sparse_out.float()).abs().max().item()
+        diffs.append(diff)
+        dense_times.append(dense_metrics.time_ms)
+        dense_mems.append(dense_metrics.peak_mem_mb or math.nan)
+        sparse_times.append(sparse_metrics.time_ms)
+        sparse_mems.append(sparse_metrics.peak_mem_mb or math.nan)
+
+        # Per-seed torch.allclose check when semantics should match (no normalization flags)
+        if not args.normalize_sparse:
+            if not torch.allclose(
+                dense_out, sparse_out, rtol=args.rtol, atol=args.atol
+            ):
+                allclose_fail_seeds.append(seed)
+
+    print("=== Distillation Dense vs Sparse Comparison ===")
+    print(f"Device: {device}  dtype: {dtype}  seeds: {args.seeds}")
+    print(f"Shape: N={N} positions, V={V} vocab, K={K} top-k per position")
+    print(f"Sparse renormalize flag (normalize_sparse): {args.normalize_sparse}")
+    print("-- Timing (ms per iteration) --")
+    print(
+        f"Dense mean: {statistics.mean(dense_times):.3f}  median: {statistics.median(dense_times):.3f}"
+    )
+    print(
+        f"Sparse mean: {statistics.mean(sparse_times):.3f}  median: {statistics.median(sparse_times):.3f}"
+    )
+    speedup = statistics.mean(dense_times) / statistics.mean(sparse_times)
+    print(f"Speedup (dense / sparse): {speedup:.2f}x")
+
+    if all(not math.isnan(m) for m in dense_mems + sparse_mems):
+        print("-- Peak Memory (MB) --")
+        print(
+            f"Dense mean: {statistics.mean(dense_mems):.1f}  median: {statistics.median(dense_mems):.1f}"
+        )
+        print(
+            f"Sparse mean: {statistics.mean(sparse_mems):.1f}  median: {statistics.median(sparse_mems):.1f}"
+        )
+        mem_reduction = statistics.mean(dense_mems) / max(
+            1e-9, statistics.mean(sparse_mems)
+        )
+        print(f"Memory reduction (dense / sparse): {mem_reduction:.2f}x")
+    else:
+        print("(Memory stats not fully available on this platform / configuration)")
+
+    print("-- Numerical Differences --")
+    max_diff = max(diffs)
+    mean_diff = statistics.mean(diffs)
+    print(f"Max |dense - sparse| over seeds: {max_diff:.3e}")
+    print(f"Mean |dense - sparse|: {mean_diff:.3e}")
+
+    tol = 1e-5
+    if args.normalize_sparse:
+        print(
+            "NOTE: Normalization flags enabled; numerical results may intentionally differ from dense baseline."
+        )
+    else:
+        # Summary of allclose results
+        if not allclose_fail_seeds:
+            print(
+                f"SUCCESS: Sparse matches dense within tolerance (max diff {max_diff:.3e} < {tol}) and allclose passed for all seeds (rtol={args.rtol}, atol={args.atol})."
+            )
+        else:
+            if allclose_fail_seeds:
+                print(
+                    f"WARNING: torch.allclose failed for {len(allclose_fail_seeds)} / {args.seeds} seeds. Failed seeds: {allclose_fail_seeds[:10]}{'...' if len(allclose_fail_seeds) > 10 else ''}"
+                )
+            # if max_diff >= tol:
+            #     print(
+            #         f"WARNING: Max diff {max_diff:.3e} exceeds heuristic dtype tolerance {tol}."
+            #     )
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--seeds", type=int, default=100, help="Number of random seeds / trials"
+    )
+    p.add_argument("--n", type=int, default=1024, help="Number of label positions (N)")
+    p.add_argument("--vocab", type=int, default=32000, help="Vocabulary size (V)")
+    p.add_argument("--k", type=int, default=8, help="Top-k teacher tokens (K)")
+    p.add_argument(
+        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu"
+    )
+    p.add_argument(
+        "--fp16", action="store_true", help="Use float16 for student logits/log-probs"
+    )
+    p.add_argument(
+        "--normalize-sparse",
+        dest="normalize_sparse",
+        action="store_true",
+        help="Renormalize teacher log-probs inside sparse path",
+    )
+    return p.parse_args()
+
+
+if __name__ == "__main__":
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    args = parse_args()
+
+    # Set default tolerances if not provided
+    args.rtol = 1.3e-6
+    args.atol = 1e-5
+    run_experiment(args)

tmp/test_element_wise_mult_einsum.py

@@ -0,0 +1,33 @@
+import torch
+
+# Setup from your example
+x = torch.rand(16, 26, 8, 2048)
+a = torch.rand(1, 26, 8, 1)
+
+# Original operation for comparison
+y = x * a
+
+# --- Einsum Solutions ---
+
+# 1. Explicit form (best for clarity)
+# We assign letters to each dimension and specify the exact output format.
+y_einsum_explicit = torch.einsum("ijkl,ijkl->ijkl", x, a)
+
+# 2. Implicit form (more concise)
+# For element-wise products, you can omit the output part.
+# Einsum infers the output should contain all unique indices.
+y_einsum_implicit = torch.einsum("ijkl,ijkl", x, a)
+
+
+# --- Verification ---
+print("Original shape:", y.shape)
+print("Einsum shape:", y_einsum_explicit.shape)
+
+# Check if the results are numerically the same
+print("\nAre the results the same?")
+print("Explicit form match:", torch.allclose(y, y_einsum_explicit))
+print("Implicit form match:", torch.allclose(y, y_einsum_implicit))
+
+print(f"{y}=")
+print(f"{y_einsum_explicit=}")
+print(f"{y_einsum_implicit=}")

tmp/test_llm_lingua.py

@@ -0,0 +1,19 @@
+from llmlingua import PromptCompressor
+
+with open("src/ctx_to_lora/modeling/test_ctx.txt") as f:
+    prompt = f.read()
+
+llm_lingua = PromptCompressor(
+    model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
+    use_llmlingua2=True,  # Whether to use llmlingua-2
+)
+compressed_prompt = llm_lingua.compress_prompt(
+    prompt, rate=0.25, force_tokens=["\n", "?"]
+)
+print(compressed_prompt)
+
+## Or use LLMLingua-2-small model
+# llm_lingua = PromptCompressor(
+#     model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
+#     use_llmlingua2=True,  # Whether to use llmlingua-2
+# )

tmp/test_logits_processor_vllm.py

@@ -0,0 +1,39 @@
+import os
+
+import torch
+from vllm import LLM, SamplingParams
+
+os.environ["VLLM_USE_V1"] = "0"  # vLLM V1 does not support logits processors.
+
+
+class LogitsSpy:
+    def __init__(self):
+        self.processed_logits: list[torch.Tensor] = []
+
+    def __call__(self, token_ids: list[int], logits: torch.Tensor):
+        self.processed_logits.append(logits)
+        return logits
+
+
+if __name__ == "__main__":
+    llm = LLM(model="google/gemma-2-2b-it", trust_remote_code=True)
+    prompts = [
+        "What is the capital of France?",
+        "Explain the theory of relativity in simple terms.",
+        "Write a short story about a robot learning to love.",
+    ]
+    logits_spy = LogitsSpy()
+    sampling_params = SamplingParams(
+        temperature=0.7, top_p=0.95, max_tokens=5, logits_processors=[logits_spy]
+    )
+    outputs = llm.generate(prompts, sampling_params)
+    logits = logits_spy.processed_logits
+    print(f"Logits: {logits}")
+    print(len(logits))
+    print(logits[0].shape)
+    print(logits[1].shape)
+    print(logits[2].shape)
+    print(logits[3].shape)
+    print(logits[4].shape)
+
+    breakpoint()

tmp/test_needle_haystack_plot.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""
+Test script to create a line plot showing needle-in-a-haystack task accuracy vs input length.
+"""
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+plt.style.use(["ggplot", latte_style])
+plt.rcParams["axes.prop_cycle"] = mpl.cycler(
+    color=[
+        "#FFB83D",
+        "#5571EA",
+    ]
+)
+plt.rcParams["axes.facecolor"] = "white"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+
+# plt.rcParams["font.family"] = "Ubuntu"
+# plt.rcParams["font.size"] = 14
+# plt.rcParams["font.weight"] = "bold"
+plt.rcParams["axes.labelweight"] = "bold"
+plt.rcParams["axes.facecolor"] = "F7FBFC"
+plt.rcParams["figure.facecolor"] = "white"
+plt.rcParams["savefig.facecolor"] = "white"
+plt.rcParams["grid.color"] = "cccccc"
+plt.rcParams["grid.linewidth"] = 1
+plt.rcParams["axes.edgecolor"] = "ccd0da"
+plt.rcParams["legend.facecolor"] = "white"
+plt.rcParams["legend.fontsize"] = 16
+plt.rcParams["xtick.labelsize"] = 16
+plt.rcParams["ytick.labelsize"] = 16
+
+
+def main():
+    # Token bins from the user's specification
+    tok_bins = [1024 * (i + 1) for i in range(16)]
+    tok_bins += [2**14 + 2**12 * (i + 1) for i in range(4)]
+    tok_bins += [2**15 + 2**13 * (i + 1) for i in range(12)]
+    x_values = tok_bins
+    print(x_values)
+
+    y_values = {
+        "Base model w/ context": [
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            0.99,
+            0.80,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+            0.00,
+        ],
+        "Ours": [
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            1.00,
+            0.95,
+            0.87,
+            0.80,
+            0.65,
+            0.53,
+            0.43,
+            0.37,
+            0.28,
+            0.24,
+            0.17,
+            0.15,
+        ],
+    }
+
+    # Create the plot
+    cleaner_x_vals = [
+        2**10,
+        2**11,
+        2**12,
+        2**13,
+        2**13 + 2**10,
+        2**13 + 2**11,
+        2**14,
+        2**15,
+        2**15 + 2**13,
+    ]
+    plt.figure(figsize=(9, 6))
+    for name, ys in y_values.items():
+        xys = [
+            (x, y)
+            for x, y in zip(x_values, ys)
+            if x in cleaner_x_vals or x > 2**15 + 2**13
+        ]
+        plt.plot(
+            [xy[0] for xy in xys],
+            [xy[1] for xy in xys],
+            marker="o",
+            linewidth=2,
+            markersize=8,
+            alpha=1.0,
+            label=name,
+            markeredgecolor="black",
+            markeredgewidth=1,
+            markerfacecolor=None,
+        )
+
+    # Customize the plot
+    plt.xlabel("Haystack Length (tokens)", fontsize=20)
+    plt.ylabel("ROUGE-L F1 Score", fontsize=20)
+    plt.title(
+        "NIAH Retrieval Performance",
+        fontsize=24,
+        fontweight="bold",
+    )
+    plt.grid(True, alpha=0.3)
+
+    # Set log scale for x-axis
+    plt.xscale("log", base=2)
+
+    # Format x-axis labels with K suffix (representing 1024 multiples)
+    # x_labels = [
+    #     f"{x // 1024}K" if x >= 1024 else str(x)
+    #     for x in x_values
+    #     if x in cleaner_x_vals or x > 2**15 + 2**13
+    # ]
+    # plt.xticks(
+    #     [x for x in x_values if x in cleaner_x_vals or x > 2**15 + 2**13],
+    #     x_labels,
+    #     rotation=90,
+    #     fontsize=8,
+    # )
+
+    # Remove minor ticks for cleaner appearance
+    plt.gca().xaxis.set_minor_locator(plt.NullLocator())
+
+    # Add some styling
+    plt.ylim(-0.05, 1.05)
+    plt.xlim(min(x_values) * 0.8, max(x_values) * 1.2)
+
+    # Add horizontal lines for reference
+
+    # Add legend
+    plt.legend(loc="lower left")
+
+    # Tight layout for better appearance
+    plt.tight_layout()
+
+    # Save the plot
+    plt.savefig(
+        "/home/tan/research/ctx-to-lora/tmp/needle_haystack_accuracy_plot.pdf",
+        format="pdf",
+        dpi=300,
+        bbox_inches="tight",
+    )
+
+    # Show statistics
+    print(f"Generated plot with {len(x_values)} data points")
+    print(f"Input length range: {min(x_values):.0f} - {max(x_values):.0f} tokens")
+    print(
+        f"Plot saved to: /home/tan/research/ctx-to-lora/tmp/needle_haystack_accuracy_plot.pdf"
+    )
+
+    # Display the plot
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

tmp/test_per_ctx_loss.py

@@ -0,0 +1,243 @@
+"""Test / diagnostic script for comparing current per_ctx_loss implementation
+against a proposed fixed version that detects label spans via mask transitions.
+
+Run:
+  python tmp/test_per_ctx_loss.py
+
+It will:
+  * Construct synthetic batches with known label span structure.
+  * Call the existing per_ctx_loss (imported from ctx_to_lora.trainer).
+  * Call a fixed implementation (fixed_per_ctx_loss) using label_mask transitions.
+  * Show detected spans, per-query losses, per-context losses, and scaling ratios.
+  * Highlight cases where the current implementation fails (empty spans / mismatch).
+
+Note: The current per_ctx_loss assumes the provided `loss` vector is ordered
+exactly as the concatenation of all label-token losses (no unlabeled tokens),
+matching DistillationTrainer's usage, but NOT CrossEntropyTrainer's raw CE output.
+"""
+
+from __future__ import annotations
+
+import math
+
+# Ensure "src" is on path so we can import project modules when running from project root
+import sys
+import traceback
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from ctx_to_lora.trainer import per_ctx_loss  # type: ignore
+
+
+@dataclass
+class SyntheticCase:
+    name: str
+    label_mask: list[bool]
+    n_queries_per_ctx: list[int]
+    description: str
+
+
+def per_ctx_loss(inputs, labels, loss):
+    n_queries_per_ctx = inputs["n_queries"].tolist()
+
+    position_ids = inputs["position_ids"].squeeze(0)
+    # account only label positions
+    label_mask = labels.squeeze(0) != -100
+    label_pos_ids = label_mask * position_ids
+    label_pos_ids_diff = label_pos_ids.diff(
+        append=torch.tensor([0], device=position_ids.device)
+    )
+    # assumes the input starts with non-assistant tokens
+    start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
+    end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
+
+    label_seq_lens = end_label_pos - start_label_pos
+    cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
+    start_indices = torch.cat(
+        (
+            torch.tensor([0], device=cu_label_seq_lens.device),
+            cu_label_seq_lens[:-1],
+        )
+    )
+
+    # these stack and split can be optimized but let's keep it simple
+    # mean across tokens of each q
+    qa_losses = torch.stack(
+        [loss[start:end].mean() for start, end in zip(start_indices, cu_label_seq_lens)]
+    )
+
+    # mean across queries of each ctx
+    per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)]
+
+    # per-ctx loss
+    loss = torch.stack(per_ctx_losses)
+    return loss
+
+
+def fixed_per_ctx_loss(inputs, labels, loss_vec: torch.Tensor) -> torch.Tensor:
+    """Reference implementation using label_mask transitions.
+
+    Args:
+        inputs: dict with keys: n_queries, position_ids
+        labels: [1, L]
+        loss_vec: tensor of shape (#label_tokens,) in label-token order.
+    Returns:
+        Tensor of shape (#contexts,) containing mean loss per context.
+    """
+    n_queries_per_ctx = inputs["n_queries"].tolist()
+    label_mask = labels.squeeze(0) != -100
+
+    lm = label_mask.to(torch.int)
+    if lm.numel() == 0:
+        return torch.empty(0)
+
+    if lm.numel() == 1:
+        spans = [(0, 1)] if lm[0] == 1 else []
+    else:
+        transition = lm[1:] - lm[:-1]
+        starts = (transition == 1).nonzero(as_tuple=False).flatten() + 1
+        ends = (transition == -1).nonzero(as_tuple=False).flatten() + 1
+        if lm[0] == 1:
+            starts = torch.cat([torch.tensor([0], device=starts.device), starts])
+        if lm[-1] == 1:
+            ends = torch.cat([ends, torch.tensor([lm.numel()], device=ends.device)])
+        assert starts.numel() == ends.numel()
+        spans = list(zip(starts.tolist(), ends.tolist()))
+
+    # Extract per-query mean losses (loss_vec is already label-token sequential).
+    lengths = [e - s for s, e in spans]
+    assert sum(lengths) == loss_vec.numel(), (
+        "Mismatch: total span label tokens != loss vector length."
+    )
+    cu_lengths = torch.tensor(lengths).cumsum(0)
+    starts_loss = torch.cat([torch.tensor([0]), cu_lengths[:-1]])
+    qa_means = [loss_vec[a:b].mean() for a, b in zip(starts_loss, cu_lengths)]
+    qa_means_t = torch.stack(qa_means)
+
+    # Group into contexts
+    split_q = torch.split(qa_means_t, n_queries_per_ctx)
+    per_ctx = torch.stack([ctx.mean() for ctx in split_q])
+    return per_ctx
+
+
+def build_labels_from_mask(mask: Sequence[bool]) -> torch.Tensor:
+    # Assign dummy token id 1 where True, -100 elsewhere.
+    data = [1 if m else -100 for m in mask]
+    return torch.tensor([data])  # shape [1, L]
+
+
+def build_position_ids(L: int) -> torch.Tensor:
+    return torch.arange(L).unsqueeze(0)
+
+
+def make_loss_vector(mask: Sequence[bool]) -> torch.Tensor:
+    # Provide a smoothly varying loss value per label token to visualize averaging.
+    # Example: losses increase linearly → weighting differences are obvious.
+    values = []
+    counter = 0
+    for m in mask:
+        if m:
+            # synthetic harder tokens later
+            values.append(0.5 + 0.1 * counter)
+            counter += 1
+    return torch.tensor(values)
+
+
+def run_case(case: SyntheticCase):
+    print(f"\n=== Case: {case.name} ===")
+    print(case.description)
+    mask = case.label_mask
+    labels = build_labels_from_mask(mask)
+    position_ids = build_position_ids(len(mask))
+    loss_vec = make_loss_vector(mask)
+    inputs = {
+        "n_queries": torch.tensor(case.n_queries_per_ctx),
+        "position_ids": position_ids,
+    }
+    print(f"Label mask: {''.join('1' if m else '.' for m in mask)}")
+    print(f"Loss vector (per label token order): {loss_vec.tolist()}")
+
+    # Attempt current implementation
+    try:
+        current = per_ctx_loss(inputs, labels, loss_vec)
+        print(
+            "Current per_ctx_loss output:",
+            current.tolist(),
+            "shape=",
+            list(current.shape),
+        )
+    except Exception as e:
+        print("Current per_ctx_loss raised exception:")
+        traceback.print_exc()
+        current = None
+
+    # Fixed implementation
+    try:
+        fixed = fixed_per_ctx_loss(inputs, labels, loss_vec)
+        print(
+            "Fixed per_ctx_loss output:  ", fixed.tolist(), "shape=", list(fixed.shape)
+        )
+    except Exception as e:
+        print("Fixed implementation raised exception:")
+        traceback.print_exc()
+        fixed = None
+
+    # Compare magnitudes if both succeeded
+    if current is not None and fixed is not None and current.numel() == fixed.numel():
+        ratio = (
+            (current.mean() / fixed.mean()).item() if fixed.mean() != 0 else math.nan
+        )
+        print(f"Mean(current)/Mean(fixed) ratio: {ratio:.4f}")
+    print("---")
+
+
+def main():
+    cases = [
+        # SyntheticCase(
+        #     name="Two spans, single context",
+        #     label_mask=[False, False, True, True, True, False, False, True, True],
+        #     n_queries_per_ctx=[2],
+        #     description="Two answer/query spans of lengths 3 and 2 in one context.",
+        # ),
+        SyntheticCase(
+            name="Single span",
+            label_mask=[False, True, True, True, False],
+            n_queries_per_ctx=[1],
+            description="One continuous labeled span.",
+        ),
+        # SyntheticCase(
+        #     name="Three contexts (2,1 queries)",
+        #     label_mask=[False, True, True, False, True, False, True, True, True, False],
+        #     n_queries_per_ctx=[1, 1, 1],
+        #     description="Three separate single-query contexts.",
+        # ),
+        # SyntheticCase(
+        #     name="Edge case: starts labeled",
+        #     label_mask=[True, True, False, True, False],
+        #     n_queries_per_ctx=[2],
+        #     description="Sequence begins with a labeled span then another later.",
+        # ),
+        # SyntheticCase(
+        #     name="Edge case: ends labeled",
+        #     label_mask=[False, True, False, True, True],
+        #     n_queries_per_ctx=[2],
+        #     description="Sequence ends with a labeled span.",
+        # ),
+    ]
+
+    for case in cases:
+        run_case(case)
+
+    print("\nDone. Inspect above for discrepancies.\n")
+
+
+if __name__ == "__main__":
+    main()

tmp/test_perf_vs_latency_plot.py

@@ -0,0 +1,89 @@
+import matplotlib.patches as patches
+import matplotlib.pyplot as plt
+import numpy as np
+
+# --- Data Points ---
+# Each tuple represents (Update Latency, Performance)
+data = {"CD": (27, 92), "TTT": (28, 65), "T2L": (4, 25), "P2L": (6, 80)}
+
+labels = list(data.keys())
+x_coords = [val[0] for val in data.values()]
+y_coords = [val[1] for val in data.values()]
+
+# --- Plotting Setup ---
+# Using a modern and clean style for the plot
+plt.style.use("seaborn-v0_8-whitegrid")
+fig, ax = plt.subplots(figsize=(10, 7), dpi=100)
+
+# --- Create the Scatter Plot ---
+# Using a single color for points and making them large for visibility
+colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(labels)))
+scatter = ax.scatter(
+    x_coords, y_coords, s=200, c=colors, edgecolors="black", alpha=0.8, zorder=5
+)
+
+# --- Annotate Each Point ---
+# Add text labels next to each point for clarity
+for i, label in enumerate(labels):
+    # Adjust text position slightly for better aesthetics
+    ax.annotate(
+        label,
+        (x_coords[i], y_coords[i]),
+        textcoords="offset points",
+        xytext=(15, 0),  # Offset text to the right of the point
+        ha="left",
+        fontsize=12,
+        fontweight="bold",
+        zorder=6,
+    )
+
+# --- Add the "Efficient Internalization" Zone ---
+# Create a transparent rectangle in the top-left corner
+efficient_zone = patches.Rectangle(
+    (0, 70),  # (x, y) coordinate of the bottom-left corner
+    12,  # Width of the rectangle
+    30,  # Height of the rectangle
+    edgecolor="green",
+    facecolor="green",
+    alpha=0.15,  # Transparency level
+    linewidth=1.5,
+    linestyle="--",
+)
+ax.add_patch(efficient_zone)
+
+# Add text label for the zone
+ax.text(
+    1,
+    85,  # (x, y) position for the text
+    "Efficient Internalization",
+    fontsize=14,
+    fontweight="bold",
+    color="darkgreen",
+    rotation=0,  # Can be rotated if needed, e.g., rotation=90
+)
+
+
+# --- Final Touches & Labels ---
+# Set plot title and axis labels
+ax.set_title(
+    "Performance vs. Update Latency Analysis", fontsize=18, fontweight="bold", pad=20
+)
+ax.set_xlabel("Update Latency (seconds)", fontsize=14, labelpad=15)
+ax.set_ylabel("Performance Score", fontsize=14, labelpad=15)
+
+# Set axis limits to provide some padding
+ax.set_xlim(0, 32)
+ax.set_ylim(0, 105)
+
+# Customize tick parameters for a cleaner look
+ax.tick_params(axis="both", which="major", labelsize=12)
+
+# Add a subtle grid
+ax.grid(True, which="both", linestyle="--", linewidth=0.5)
+
+# Ensure the layout is tight and clean
+plt.tight_layout()
+
+# --- Display the Plot ---
+plt.show()
+plt.savefig("perf_vs_latency_plot.png", dpi=300)  # Save as high-res PNG

tmp/viz_gen_lora.py

@@ -0,0 +1,113 @@
+from collections import defaultdict
+
+import torch
+
+# Store activations
+activations = defaultdict(list)
+
+
+def apply_forward_hook(module, name):
+    def hook_fn(m, input_args, output_tensor):
+        activations[name].append(output_tensor.cpu().detach())
+
+    module.register_forward_hook(hook_fn)
+
+
+def last_token_pool(
+    outputs: dict[str, torch.Tensor], attention_mask: torch.Tensor
+) -> torch.Tensor:
+    last_hidden_states = (
+        outputs["hidden_states"][-1].detach()
+        if "hidden_states" in outputs
+        else outputs["last_hidden_state"].detach()
+    )
+    left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[
+            torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths
+        ]
+
+
+if __name__ == "__main__":
+    from functools import partial
+
+    import torch
+    from peft import get_peft_model
+    from torch.utils.data import DataLoader
+    from tqdm import tqdm
+
+    from ctx_to_lora.data.collator import generation_collator
+    from ctx_to_lora.data.processing import get_tokenized_dataset
+    from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
+    from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
+
+    model_name_or_path = "google/gemma-2-2b-it"
+    base_model, tokenizer = get_model_and_tokenizer(
+        model_name_or_path, train=False, requires_grad=False
+    )
+    ds = get_tokenized_dataset(
+        ds_name="squad",
+        split="test",
+        max_qas_len=-1,
+        max_qas_per_sample=1,
+        base_model_max_len=8192,
+        tokenizer=tokenizer,
+        ctx_model_max_len=8192,
+        ctx_tokenizer=tokenizer,
+        max_ctx_chunk_len=-1,
+        min_ctx_chunk_len=-1,
+        num_chunk_probs=None,
+        max_ctx_chunk_num=None,
+        add_ctx_to_chat=False,
+        use_kl_loss=True,
+        max_new_tokens=0,
+        add_self_distill_template=False,
+        set_format="pt",
+    )
+    ds = ds.take(100)
+
+    peft_config = get_lora_config(
+        model_name_or_path,
+        lora_r=8,
+        lora_dropout=0,
+        target_modules=["down_proj"],
+    )
+    peft_config.lora_alpha = 16
+    peft_model = get_peft_model(base_model, peft_config)
+
+    checkpoint_path = "train_outputs/runs/Sep06_12-27-02_slurm0-a3nodeset-11_84523_bbcb67ca/checkpoint-40000/pytorch_model.bin"
+    model = ModulatedPretrainedModel.from_state_dict(
+        state_dict=torch.load(checkpoint_path, weights_only=False),
+        train=False,
+        base_model_kwargs=dict(attn_implementation="flash_attention_2"),
+        use_flash_attn=True,
+        use_sequence_packing=False,  # for generation
+    )
+    dataloader = DataLoader(
+        ds,
+        batch_size=1,
+        shuffle=False,
+        sampler=None,
+        batch_sampler=None,
+        num_workers=0,
+        collate_fn=partial(generation_collator, tokenizer=tokenizer),
+        pin_memory=False,
+        drop_last=False,
+        timeout=0,
+        persistent_workers=False,
+    )
+
+    apply_forward_hook(model.ctx_encoder, "ctx_encoder")
+    apply_forward_hook(model.layers, "hypernet_layers")
+
+    for sample_idx, sample in tqdm(enumerate(dataloader)):
+        for k, v in sample.items():
+            sample[k] = v.to(model.device)
+        with torch.no_grad():
+            lora_state, _ = model.generate_weights(
+                sample["ctx_ids"], sample["ctx_attn_mask"]
+            )

tmp/viz_lora_weights.py

@@ -0,0 +1,338 @@
+if __name__ == "__main__":
+    # New imports for analysis
+    from collections import defaultdict
+    from functools import partial
+
+    import torch
+    from peft import get_peft_model
+    from torch.utils.data import DataLoader
+    from tqdm import tqdm
+
+    try:
+        from sklearn.manifold import TSNE
+    except ImportError:
+        TSNE = None
+    try:
+        import umap
+    except ImportError:
+        umap = None
+    # Added plotting imports
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        plt = None
+    try:
+        import seaborn as sns
+    except ImportError:
+        sns = None
+
+    from ctx_to_lora.data.collator import generation_collator
+    from ctx_to_lora.data.definitions import CTX_AFFIXES
+    from ctx_to_lora.data.processing import get_tokenized_dataset
+    from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
+    from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
+
+    # from ctx_to_lora.modeling.context_distillation import CtxDistillModel
+    from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
+
+    model_name_or_path = "google/gemma-2-2b-it"
+    base_model, tokenizer = get_model_and_tokenizer(
+        model_name_or_path, train=False, requires_grad=False
+    )
+    ds = get_tokenized_dataset(
+        ds_name="squad",
+        split="test",
+        max_qas_len=-1,
+        max_qas_per_sample=1,
+        base_model_max_len=8192,
+        tokenizer=tokenizer,
+        ctx_model_max_len=8192,
+        ctx_tokenizer=tokenizer,
+        max_ctx_chunk_len=-1,
+        min_ctx_chunk_len=-1,
+        num_chunk_probs=None,
+        max_ctx_chunk_num=None,
+        add_ctx_to_chat=False,
+        use_kl_loss=True,
+        max_new_tokens=0,
+        add_self_distill_template=False,
+        set_format="pt",
+    )
+    ds = ds.take(100)
+
+    peft_config = get_lora_config(
+        model_name_or_path,
+        lora_r=8,
+        lora_dropout=0,
+        target_modules=["down_proj"],
+    )
+    peft_config.lora_alpha = 16
+    peft_model = get_peft_model(base_model, peft_config)
+    sep_seq = (
+        tokenizer(
+            SELF_QA_INTX.strip("\n"),
+            add_special_tokens=False,
+            return_tensors="pt",
+        )
+        .input_ids[0]
+        .to(base_model.device)
+    )
+    ctx_distill_kwargs = dict(
+        prefix_tokens=torch.tensor(
+            CTX_AFFIXES[model_name_or_path]["prefix"], device=base_model.device
+        ),
+        ctx_inp_sep_seq=sep_seq,
+        pad_token_id=tokenizer.pad_token_id,
+        update_iterations=100,
+        tokenizer=tokenizer,
+        reprompt_ctx=False,
+    )
+
+    # if args.cd_use_gen_q:
+    #     q_model, q_tokenizer = get_model_and_tokenizer(
+    #         "google/gemma-3-4b-it",
+    #         train=False,
+    #         requires_grad=False,
+    #     )
+    #     ctx_distill_kwargs["q_model"] = q_model
+    #     ctx_distill_kwargs["q_tokenizer"] = q_tokenizer
+    # model = CtxDistillModel(peft_model, **ctx_distill_kwargs)
+    checkpoint_path = "train_outputs/runs/Sep06_12-27-02_slurm0-a3nodeset-11_84523_bbcb67ca/checkpoint-40000/pytorch_model.bin"
+    model = ModulatedPretrainedModel.from_state_dict(
+        state_dict=torch.load(checkpoint_path, weights_only=False),
+        train=False,
+        base_model_kwargs=dict(attn_implementation="flash_attention_2"),
+        use_flash_attn=True,
+        use_sequence_packing=False,  # for generation
+    )
+    dataloader = DataLoader(
+        ds,
+        batch_size=1,
+        shuffle=False,
+        sampler=None,
+        batch_sampler=None,
+        num_workers=0,
+        collate_fn=partial(generation_collator, tokenizer=tokenizer),
+        pin_memory=False,
+        drop_last=False,
+        timeout=0,
+        persistent_workers=False,
+    )
+
+    # Control / config variables for analysis
+    collect_max_samples = None  # set int to limit samples processed
+    modules_filter_substr = None  # e.g. "down_proj" to restrict; or None for all
+    perform_analysis = True  # set False to only collect
+    save_dir = "./lora_weight_analysis"  # will be created if not exists
+    use_fp32_storage = False  # store ΔW in float32 (else float16 if original)
+    apply_pca_before_tsne = True
+    pca_dim = 64
+    random_seed = 42
+    perform_plotting = True  # NEW: set False to skip figure generation
+
+    # Data structures:
+    # delta_w_flat[module] -> list[Tensor(num_params,)]  (one per sample)
+    delta_w_flat = defaultdict(list)
+    sample_ids = []  # overall sample index order
+    # Keep raw shapes (out_features, in_features)
+    delta_w_shapes = {}
+
+    # Precompute scaling factor (α / r)
+    lora_alpha = peft_config.lora_alpha
+    lora_r = peft_config.r
+    lora_scaling = lora_alpha / lora_r
+
+    for sample_idx, sample in tqdm(enumerate(dataloader)):
+        if collect_max_samples is not None and sample_idx >= collect_max_samples:
+            break
+        for k, v in sample.items():
+            sample[k] = v.to(model.device)
+        with torch.no_grad():
+            lora_state, _ = model.generate_weights(
+                sample["ctx_ids"], sample["ctx_attn_mask"]
+            )
+
+        # # NEW: lora_state is now a nested dict:
+        # # { module_name: { "A": [1, n_layers, r, d_in], "B": [1, n_layers, r, d_out] } }
+        for module_name, parts in lora_state.items():
+            if "A" not in parts or "B" not in parts:
+                continue
+            A_all = parts["A"].squeeze(0).float()  # [n_layers, r, d_in]
+            B_all = parts["B"].squeeze(0).float()  # [n_layers, r, d_out]
+            n_layers_local = A_all.size(0)
+            for layer_idx in range(n_layers_local):
+                A = A_all[layer_idx]  # [r, d_in]
+                B = B_all[layer_idx]  # [r, d_out]
+                # Original expectation was B: [d_out, r]; now transpose
+                delta_w = B.T @ A  # [d_out, d_in]
+                mod_key = f"{module_name}.layer{layer_idx}"
+                if modules_filter_substr and modules_filter_substr not in mod_key:
+                    continue
+                delta_w_shapes.setdefault(mod_key, delta_w.shape)
+                if use_fp32_storage:
+                    delta_w = delta_w.float()
+                delta_w_flat[mod_key].append(delta_w.reshape(-1).cpu())
+        sample_ids.append(sample_idx)
+
+    if perform_analysis and len(delta_w_flat):
+        import os
+
+        os.makedirs(save_dir, exist_ok=True)
+        torch.manual_seed(random_seed)
+
+        # Optional PCA helper
+        def run_pca(x, k):
+            # x: (n, d)
+            if k >= x.shape[1]:
+                return x
+            # Center
+            mean = x.mean(0, keepdim=True)
+            x0 = x - mean
+            # SVD economy
+            try:
+                U, S, Vt = torch.linalg.svd(x0, full_matrices=False)
+            except RuntimeError:
+                # fallback CPU
+                U, S, Vt = torch.linalg.svd(x0.cpu(), full_matrices=False)
+                U, S, Vt = U.to(x0.device), S.to(x0.device), Vt.to(x0.device)
+            comp = Vt[:k]
+            return torch.matmul(x0, comp.T)
+
+        analysis_summary = {}
+
+        for module_name, vec_list in delta_w_flat.items():
+            X = torch.stack(vec_list)  # (num_samples, num_params)
+            # Cosine similarity matrix
+            X_norm = torch.nn.functional.normalize(X, dim=1)
+            cosine_sim = X_norm @ X_norm.T  # (N,N)
+            print(cosine_sim)
+
+            # Dimensionality reduction inputs
+            X_dr = X
+            if apply_pca_before_tsne and TSNE is not None:
+                # Keep at most pca_dim or (num_samples - 1, num_params)
+                k = min(pca_dim, X.shape[1])
+                X_dr = run_pca(X.to(torch.float32), k)
+
+            embeddings = {}
+            if TSNE is not None and X.shape[0] > 2:
+                try:
+                    perplexity = max(2, min(30, X.shape[0] - 1))
+                    tsne_2d = TSNE(
+                        n_components=2,
+                        perplexity=perplexity,
+                        random_state=random_seed,
+                        init="pca",
+                        learning_rate="auto",
+                    ).fit_transform(X_dr.numpy())
+                    embeddings["tsne"] = torch.from_numpy(tsne_2d)
+                except Exception as e:
+                    print(f"[WARN] TSNE failed for {module_name}: {e}")
+
+            if umap is not None and X.shape[0] > 2:
+                try:
+                    reducer = umap.UMAP(
+                        n_components=2,
+                        n_neighbors=min(15, max(2, X.shape[0] - 1)),
+                        random_state=random_seed,
+                        metric="cosine",
+                    )
+                    umap_2d = reducer.fit_transform(X_dr.numpy())
+                    embeddings["umap"] = torch.from_numpy(umap_2d)
+                except Exception as e:
+                    print(f"[WARN] UMAP failed for {module_name}: {e}")
+
+            # Save tensors
+            base_fname = module_name.replace(".", "_")
+            torch.save(
+                {
+                    "module": module_name,
+                    "sample_ids": sample_ids,
+                    "delta_W_flat": X,  # shape (N, P)
+                    "shape_matrix": delta_w_shapes[module_name],
+                    "cosine_sim": cosine_sim,  # (N, N)
+                    "embeddings": embeddings,  # dict
+                    "config": {
+                        "lora_alpha": lora_alpha,
+                        "lora_r": lora_r,
+                        "scaling": lora_scaling,
+                        "pca_applied": apply_pca_before_tsne and TSNE is not None,
+                    },
+                },
+                os.path.join(save_dir, f"{base_fname}.pt"),
+            )
+
+            analysis_summary[module_name] = {
+                "num_samples": X.shape[0],
+                "num_params": X.shape[1],
+                "saved_file": os.path.join(save_dir, f"{base_fname}.pt"),
+                "have_tsne": "tsne" in embeddings,
+                "have_umap": "umap" in embeddings,
+            }
+
+            # -------- NEW PLOTTING SECTION --------
+            if perform_plotting and plt is not None:
+                # Cosine similarity heatmap
+                try:
+                    fig, ax = plt.subplots(figsize=(4, 4))
+
+                    im = ax.imshow(cosine_sim.numpy(), cmap="viridis")
+                    ax.set_xticks(range(len(sample_ids)))
+                    ax.set_yticks(range(len(sample_ids)))
+                    ax.set_xticklabels(sample_ids, fontsize=6)
+                    ax.set_yticklabels(sample_ids, fontsize=6)
+                    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+                    ax.set_title(
+                        f"Cosine Sim: {module_name.split('.')[-3:]}".replace(",", "")
+                    )
+                    ax.set_xlabel("Sample")
+                    ax.set_ylabel("Sample")
+                    fig.tight_layout()
+                    fig.savefig(
+                        os.path.join(save_dir, f"{base_fname}_cosine.png"), dpi=300
+                    )
+                    plt.close(fig)
+                except Exception as e:
+                    print(f"[WARN] Cosine plot failed for {module_name}: {e}")
+
+                # TSNE / UMAP scatter
+                def _scatter(emb_key):
+                    emb = embeddings.get(emb_key)
+                    if emb is None:
+                        return
+                    try:
+                        emb_np = emb.numpy()
+                        fig, ax = plt.subplots(figsize=(4, 4))
+                        ax.scatter(
+                            emb_np[:, 0],
+                            emb_np[:, 1],
+                            s=30,
+                            c=range(len(emb_np)),
+                            cmap="tab10",
+                        )
+                        if len(emb_np) <= 20:
+                            for i, (xv, yv) in enumerate(emb_np):
+                                ax.text(xv, yv, str(sample_ids[i]), fontsize=6)
+                        ax.set_title(f"{emb_key.upper()} {module_name.split('.')[-3:]}")
+                        ax.set_xticks([])
+                        ax.set_yticks([])
+                        fig.tight_layout()
+                        fig.savefig(
+                            os.path.join(save_dir, f"{base_fname}_{emb_key}.png"),
+                            dpi=150,
+                        )
+                        plt.close(fig)
+                    except Exception as e:
+                        print(f"[WARN] {emb_key} plot failed for {module_name}: {e}")
+
+                _scatter("tsne")
+                _scatter("umap")
+            # -------- END NEW PLOTTING SECTION --------
+
+        # Print a concise summary
+        print("LoRA ΔW analysis summary:")
+        for m, info in analysis_summary.items():
+            print(
+                f"{m}: samples={info['num_samples']} params={info['num_params']} "
+                f"tsne={info['have_tsne']} umap={info['have_umap']} -> {info['saved_file']}"
+            )

tmp/vllm_logprobs.py

@@ -0,0 +1,36 @@
+from datasets import load_dataset
+
+from ctx_to_lora.data.processing import get_tokenized_dataset
+from ctx_to_lora.model_loading import get_tokenizer
+
+if __name__ == "__main__":
+    ds = load_dataset(
+        "parquet",
+        data_files="./data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact/train/ds.parquet",
+        split="train",
+    )
+
+    # tokenizer = ctx_tokenizer = AutoTokenizer.from_pretrained(
+    #     "google/gemma-2-2b-it",
+    # )
+    tokenizer = ctx_tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
+
+    tokenized_ds = get_tokenized_dataset(
+        "self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact",
+        split="train",
+        base_model_max_len=2**13,
+        tokenizer=tokenizer,
+        tokenizer_kwargs={},
+        ctx_model_max_len=2**13,
+        ctx_tokenizer=ctx_tokenizer,
+        ctx_tokenizer_kwargs={},
+        max_qas_len=2048,
+        max_qas_per_sample=1,
+        add_ctx_to_chat=False,
+        add_repeat_prompt=False,
+        add_negative_prompt=False,
+        use_kl_loss=True,
+    )
+    print(ds[0])
+    print(tokenized_ds[0])
+    breakpoint()