From 18dd665f39dad2201c5097b9c860c390ea6f01f8 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 7 Apr 2026 23:18:58 -0700 Subject: [PATCH 1/4] Add configurable learned-spline init activation and exploration sweep --- .../default_inf_lsa_init_mlp_swiglu.yaml | 59 +++++++++++++++++++ gpt_conf.py | 1 + train_args.py | 11 ++++ variations/activation_variations.py | 37 +++++++++--- 4 files changed, 101 insertions(+), 7 deletions(-) create mode 100644 explorations/default_inf_lsa_init_mlp_swiglu.yaml diff --git a/explorations/default_inf_lsa_init_mlp_swiglu.yaml b/explorations/default_inf_lsa_init_mlp_swiglu.yaml new file mode 100644 index 0000000000..1433118c4b --- /dev/null +++ b/explorations/default_inf_lsa_init_mlp_swiglu.yaml @@ -0,0 +1,59 @@ +# default_inf_lsa_init_mlp_swiglu.yaml +--- + +named_static_groups: + # QK Norm + - named_group: "qk_norm" + use_qk_norm: [true] + use_qk_norm_scale: [true] + + # Norm Type + - named_group: "peri_ln" + use_pre_ln: [true] + use_peri_ln: [true] + use_post_ln: [false] + + # Position Embeddings + - named_group: "rotary" + use_rotary_embeddings: [true] + use_abs_pos_embeddings: [false] + + # Relu2Max + - named_group: "relu2max" + softmax_variant_attn: ["relu2max"] + + # Infinite Attention + - named_group: "infinite" + attention_variant: ["infinite"] + use_concat_heads: [true] + + # MQA + - named_group: "mqa" + n_kv_group: [1] + + # Head Dimension + - named_group: "hd_100" + n_qk_head_dim: [100] + n_v_head_dim: [100] + +common_group: + dataset: ["minipile"] + eval_interval: [2500] + max_iters: [10000] + never_save_checkpoint: [true] + compile: [true] + log_rankme: [true] + log_areq: [true] + +parameter_groups: + - named_group_static: + - "qk_norm" + - "peri_ln" + - "rotary" + - "relu2max" + - "infinite" + - "mqa" + - "hd_100" + mlp_variant: ["mlp", "swiglu"] + activation_variant: ["learned_spline"] + lsa_init_activation: ["mish", "silu", "gelu", "squared_relu"] diff --git a/gpt_conf.py b/gpt_conf.py index 5b1323e3fe..fed13416aa 100644 --- a/gpt_conf.py +++ b/gpt_conf.py @@ -442,6 +442,7 @@ class GPTConfig: ## LearnedSplineActivation - lsa lsa_num_knots: int = 30 + lsa_init_activation: str = "gelu" # Linear Alternatives diff --git a/train_args.py b/train_args.py index 9553b7d92b..6c4096e641 100644 --- a/train_args.py +++ b/train_args.py @@ -835,6 +835,10 @@ def parse_args(): "tanh", "identity", ] + lsa_init_activation_variations = [ + variation for variation in activation_variations + if variation != "learned_spline" + ] ## DynamicActivations model_group.add_argument("--dact_activation", type=str, default="tanh", choices=activation_variations) @@ -873,6 +877,13 @@ def parse_args(): ## LearnedSplineActivation - lsa model_group.add_argument("--lsa_num_knots", type=int, default=30) + model_group.add_argument( + "--lsa_init_activation", + type=str, + default="gelu", + choices=lsa_init_activation_variations, + help="Activation used to initialize learned_spline knot y-values.", + ) # Attention Variations diff --git a/variations/activation_variations.py b/variations/activation_variations.py index e253c2ce48..64267d785b 100644 --- a/variations/activation_variations.py +++ b/variations/activation_variations.py @@ -173,17 +173,40 @@ class LearnedSplineActivation(nn.Module): def __init__(self, config, num_knots=10, init_x_range=(-5, 5)): super().__init__() - self.num_knots = config.lsa_num_knots + self.num_knots = getattr(config, "lsa_num_knots", num_knots) + self.init_activation_name = getattr(config, "lsa_init_activation", "gelu") - # Initialize learnable x_vals and y_vals - x_init = torch.linspace(init_x_range[0], init_x_range[1], num_knots) + if self.init_activation_name == "learned_spline": + raise ValueError( + "lsa_init_activation cannot be 'learned_spline' because it would recurse during initialization." + ) - # Create an instance of GELU - gelu = nn.GELU() - y_init = gelu(x_init) # Use the GELU instance here + init_activation_cls = activation_dictionary.get(self.init_activation_name) + if init_activation_cls is None: + raise ValueError( + f"Unknown lsa_init_activation '{self.init_activation_name}'. " + f"Expected one of: {sorted(activation_dictionary.keys())}" + ) + + # Initialize learnable x_vals and y_vals + x_init = torch.linspace(init_x_range[0], init_x_range[1], self.num_knots) + init_activation = init_activation_cls(config) + + with torch.no_grad(): + if self.init_activation_name == "glu": + # GLU halves the last dimension, so feed a 2-channel input and squeeze back. + y_init = init_activation(torch.stack((x_init, x_init), dim=-1)).squeeze(-1) + else: + y_init = init_activation(x_init) + + if y_init.shape != x_init.shape: + raise ValueError( + f"lsa_init_activation='{self.init_activation_name}' must preserve shape. " + f"Got output shape {tuple(y_init.shape)} for input shape {tuple(x_init.shape)}." + ) self.x_vals = nn.Parameter(x_init) - self.y_vals = nn.Parameter(y_init) + self.y_vals = nn.Parameter(y_init.detach().clone()) def forward(self, x): # Compute spline coefficients From c6ca55369faff329001bf76792c74a02c5529dee Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Wed, 8 Apr 2026 18:32:35 -0700 Subject: [PATCH 2/4] Add Plotly learned-spline layer visualizer script --- .../plot_learned_splines.py | 250 ++++++++++++++++++ 1 file changed, 250 insertions(+) create mode 100644 analysis/activation_analysis/plot_learned_splines.py diff --git a/analysis/activation_analysis/plot_learned_splines.py b/analysis/activation_analysis/plot_learned_splines.py new file mode 100644 index 0000000000..db26eb53f9 --- /dev/null +++ b/analysis/activation_analysis/plot_learned_splines.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""Plot learned spline activations from one or more checkpoints. + +This script finds learned-spline knot parameters (`x_vals`, `y_vals`) in checkpoint +state dicts, reconstructs the spline curves, and writes a single interactive Plotly +HTML figure with one trace per layer/path. +""" + +import argparse +import re +from pathlib import Path + +import plotly.graph_objects as go +import torch + + +def _load_state_dict(checkpoint_path: Path): + ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + + # Common checkpoint layouts in this repo: + # 1) {'model': state_dict, ...} + # 2) raw state_dict + if isinstance(ckpt, dict) and "model" in ckpt and isinstance(ckpt["model"], dict): + return ckpt["model"] + if isinstance(ckpt, dict): + return ckpt + + raise ValueError(f"Unsupported checkpoint format for: {checkpoint_path}") + + +def _natural_cubic_spline_eval(x_vals: torch.Tensor, y_vals: torch.Tensor, x_query: torch.Tensor): + """Mirror LearnedSplineActivation's natural cubic spline evaluation.""" + x_vals = x_vals.to(dtype=torch.float64) + y_vals = y_vals.to(dtype=torch.float64) + x_query = x_query.to(dtype=torch.float64) + + n = x_vals.numel() + if n < 2: + raise ValueError("Need at least 2 knots to evaluate spline.") + + h = x_vals[1:] - x_vals[:-1] + delta = (y_vals[1:] - y_vals[:-1]) / h + + a_mat = torch.zeros((n, n), dtype=torch.float64) + rhs = torch.zeros((n,), dtype=torch.float64) + + a_mat[0, 0] = 1.0 + a_mat[-1, -1] = 1.0 + + for i in range(1, n - 1): + a_mat[i, i - 1] = h[i - 1] + a_mat[i, i] = 2 * (h[i - 1] + h[i]) + a_mat[i, i + 1] = h[i] + rhs[i] = 3 * (delta[i] - delta[i - 1]) + + m_vals = torch.linalg.solve(a_mat, rhs) + + # Interval coefficients + a_coef = y_vals[:-1] + b_coef = delta - h * (2 * m_vals[:-1] + m_vals[1:]) / 3 + c_coef = m_vals[:-1] / 2 + d_coef = (m_vals[1:] - m_vals[:-1]) / (6 * h) + + indices = torch.searchsorted(x_vals, x_query, right=False) - 1 + indices = torch.clamp(indices, 0, n - 2) + + x_k = x_vals[indices] + dx = x_query - x_k + + return ( + a_coef[indices] + + b_coef[indices] * dx + + c_coef[indices] * dx**2 + + d_coef[indices] * dx**3 + ) + + +def _extract_layer_index(param_name: str): + match = re.search(r"\.h\.(\d+)\.", param_name) + if match: + return int(match.group(1)) + return None + + +def _collect_learned_splines(state_dict: dict, checkpoint_path: Path): + traces = [] + + # LearnedSplineActivation in this codebase stores parameters as: + # ...activation_variant.x_vals and ...activation_variant.y_vals + x_keys = [k for k in state_dict if k.endswith(".x_vals") and "activation_variant" in k] + + for x_key in x_keys: + y_key = x_key[:-len(".x_vals")] + ".y_vals" + if y_key not in state_dict: + continue + + x_vals = state_dict[x_key] + y_vals = state_dict[y_key] + if not isinstance(x_vals, torch.Tensor) or not isinstance(y_vals, torch.Tensor): + continue + if x_vals.ndim != 1 or y_vals.ndim != 1: + continue + if x_vals.numel() != y_vals.numel() or x_vals.numel() < 4: + continue + + layer_idx = _extract_layer_index(x_key) + traces.append( + { + "checkpoint": str(checkpoint_path), + "layer_idx": layer_idx, + "x_key": x_key, + "x_vals": x_vals.detach().cpu(), + "y_vals": y_vals.detach().cpu(), + } + ) + + traces.sort(key=lambda t: (t["checkpoint"], -1 if t["layer_idx"] is None else t["layer_idx"], t["x_key"])) + return traces + + +def _build_html_with_range_controls(fig, html_path: Path, default_xmin: float, default_xmax: float): + div = fig.to_html(full_html=False, include_plotlyjs="cdn", div_id="spline_plot") + + html = f""" + + + + Learned spline curves + + + +

Learned spline curves

+
+ + + + + + + Legend labels include checkpoint path + layer + parameter path. +
+ {div} + + + + +""" + + html_path.write_text(html, encoding="utf-8") + + +def main(): + parser = argparse.ArgumentParser(description="Plot learned spline activations from checkpoints.") + parser.add_argument( + "checkpoints", + nargs="+", + type=Path, + help="Checkpoint path(s), e.g. out/run1/ckpt.pt out/run2/ckpt.pt", + ) + parser.add_argument( + "--output_html", + type=Path, + default=Path("analysis/activation_analysis/learned_splines.html"), + help="Output HTML path.", + ) + parser.add_argument("--x_min", type=float, default=-10.0, help="Default plot x-axis minimum.") + parser.add_argument("--x_max", type=float, default=10.0, help="Default plot x-axis maximum.") + parser.add_argument("--num_points", type=int, default=1200, help="Points per spline trace.") + args = parser.parse_args() + + if args.x_min >= args.x_max: + raise ValueError("x_min must be smaller than x_max") + if args.num_points < 8: + raise ValueError("num_points should be >= 8") + + x_query = torch.linspace(args.x_min, args.x_max, args.num_points, dtype=torch.float64) + + figure = go.Figure() + total_splines = 0 + + for checkpoint_path in args.checkpoints: + state_dict = _load_state_dict(checkpoint_path) + splines = _collect_learned_splines(state_dict, checkpoint_path) + + for spline in splines: + y_query = _natural_cubic_spline_eval(spline["x_vals"], spline["y_vals"], x_query) + layer_label = f"layer {spline['layer_idx']}" if spline["layer_idx"] is not None else "layer ?" + label = f"{spline['checkpoint']} | {layer_label} | {spline['x_key']}" + + figure.add_trace( + go.Scatter( + x=x_query.tolist(), + y=y_query.tolist(), + mode="lines", + name=label, + hovertemplate="x=%{x:.3f}
y=%{y:.3f}%{fullData.name}", + ) + ) + total_splines += 1 + + if total_splines == 0: + raise RuntimeError( + "No learned spline parameters were found. " + "Expected keys like '...activation_variant.x_vals' + '...activation_variant.y_vals'." + ) + + figure.update_layout( + title=f"Learned spline curves ({total_splines} traces)", + xaxis_title="x", + yaxis_title="activation(x)", + template="plotly_white", + legend=dict(orientation="v"), + ) + figure.update_xaxes(range=[args.x_min, args.x_max]) + + args.output_html.parent.mkdir(parents=True, exist_ok=True) + _build_html_with_range_controls(figure, args.output_html, args.x_min, args.x_max) + print(f"Saved: {args.output_html} ({total_splines} traces)") + + +if __name__ == "__main__": + main() From 93aa021e217e9243adca9ed9d904449fa98f4007 Mon Sep 17 00:00:00 2001 From: klei22 Date: Wed, 8 Apr 2026 17:42:25 -0700 Subject: [PATCH 3/4] Add back z feature --- plot_view.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/plot_view.py b/plot_view.py index 0845e2a97c..b04b9baccf 100644 --- a/plot_view.py +++ b/plot_view.py @@ -430,6 +430,23 @@ def plot_multi_bars_trim( # NEW fig.write_image(f"{safe}.png", scale=2) fig.show() +def plot_bars_trim( + rows: List[Dict[str, Any]], + *, + y: str, + label_cols: List[str], +) -> None: + """ + Backward-compatible single-metric Δ-bar helper used by monitor ``z`` hotkey. + + This intentionally mirrors ``plot_bars`` args while delegating to the + multi-metric trimmed implementation. + """ + if not label_cols: + raise ValueError("Need ≥1 label column for bar chart") + plot_multi_bars_trim(rows, y_cols=[y], label_cols=label_cols) + + # ───────────────────────────── CLI wrapper ────────────────────────── From 5d8097d50fe1d5761ad5ec354ab559377dda5eba Mon Sep 17 00:00:00 2001 From: klei22 Date: Wed, 8 Apr 2026 18:53:30 -0700 Subject: [PATCH 4/4] Update yaml file for learned spline sweep --- explorations/default_inf_lsa_init_mlp_swiglu.yaml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/explorations/default_inf_lsa_init_mlp_swiglu.yaml b/explorations/default_inf_lsa_init_mlp_swiglu.yaml index 1433118c4b..00a04cdfc7 100644 --- a/explorations/default_inf_lsa_init_mlp_swiglu.yaml +++ b/explorations/default_inf_lsa_init_mlp_swiglu.yaml @@ -40,7 +40,7 @@ common_group: dataset: ["minipile"] eval_interval: [2500] max_iters: [10000] - never_save_checkpoint: [true] + never_save_checkpoint: [false] compile: [true] log_rankme: [true] log_areq: [true] @@ -50,10 +50,8 @@ parameter_groups: - "qk_norm" - "peri_ln" - "rotary" - - "relu2max" - "infinite" - - "mqa" - "hd_100" mlp_variant: ["mlp", "swiglu"] activation_variant: ["learned_spline"] - lsa_init_activation: ["mish", "silu", "gelu", "squared_relu"] + lsa_init_activation: ["mish", "silu", "gelu", "squared_relu", "identity", "relu"]