-
Notifications
You must be signed in to change notification settings - Fork 0
Home
TASM (The Alignment Stress Map) is a web-based analysis platform for measuring alignment signals in transformer language models at inference time. It implements three complementary signal families — ASM, LTP, and SFD — that together characterize the intensity, directionality, and dimensionality of the corrections introduced by instruction-tuning (RLHF/SFT) relative to a base model. The theoretical framework is described in two companion papers by Ostrander (2026).
The application is a single-server Python system built on FastAPI, serving a single-page web frontend over HTTP. It is designed for interactive single-prompt analysis, CSV-driven batch experiments, and conversational probing via a built-in chat interface.
TASM is structured as a monolithic FastAPI application (app.py, ~1,300 lines) backed by a modular engine layer of eight specialized Python modules:
The frontend consists of a single-page HTML application (static/index.html, ~159 KB) and a separate chat interface (static/chat.html). All rendering and interaction logic is client-side JavaScript; the server provides a JSON API and base64-encoded plot images.
The fundamental operation is weight delta projection. At model load time, TASM downloads the base and instruct variants of a model pair from HuggingFace, then computes the weight delta ΔW = W_instruct − W_base for six projection matrices per transformer layer (Q, K, V, O, gate, up). Base model weights are read directly from safetensors files on disk one tensor at a time, avoiding full base-model instantiation and halving peak memory.
At analysis time, a single forward pass through the instruct model with registered hooks captures hidden states and attention weights at monitored layers. All three signal families are extracted from these cached activations without additional forward passes:
ASM (Alignment Stress Map) computes per-token signed attribution by projecting hidden states through ΔW_V, weighting by attention patterns, and aggregating across signal layers (the middle third of the network). This yields a scalar stress score, per-token attribution vectors, distribution metrics (entropy, Gini, top-2 share, interior CV), and — when full trajectory mode is enabled — a layer-by-layer amplitude trace and a token×layer heatmap.
LTP (Lateral Tension Profile) probes the alignment field perpendicular to the generation path. For each token position, it identifies the top-k counterfactual tokens the model considered but did not select, computes unembedding directions for each, and measures the lateral tension (the component of the alignment correction toward each alternative) via ΔW_Vprojection. This produces per-token ranked tension profiles, profile shape classifications (steep/flat/inverted), PCA-projected dual trajectories (semantic vs. tension), and four summary statistics: M (offset magnitude), C (offset consistency), V (offset variance), and L (lateral coverage). Two optional enhancements are available: SVD truncation of ΔW_V to isolate the dominant safety subspace, and tuned-lens calibration of per-layer unembedding probes.
SFD (Spectral Field Density) measures how many dimensions of the QK routing subspace each token's activation engages. At load time, it SVD-decomposes the concatenated [ΔW_Q; ΔW_K] per layer and caches the right singular vectors. At inference, it projects each token's hidden state through this basis, yielding per-token energy, spectral entropy, and density ratio. Prompt-level aggregates (mean, max, variance, p90) summarize the dimensionality axis.
Behavioral comparison optionally loads the base model for a separate forward pass to compute KL(instruct ‖ base) divergence at the output distribution and capture top-k next-token predictions from both models. A rank displacement metric compares counterfactual token orderings between base and instruct models.
Single prompt analysis accepts:
-
prompt(string, ≤5,000 characters) — the text to analyze -
category(string) — one ofbenign,mild,harmful,jailbreak,adversarial,dual-use, or arbitrary - Boolean flags:
compute_kl,compute_trajectory,capture_responses,full_capture,compute_ltp,compute_sfd - LTP parameters:
ltp_k(counterfactual depth: 4, 6, or 8),ltp_layer_strategy(signalorlate),ltp_svd_rank(0 = raw, or truncation rank),ltp_tuned_lens(boolean)
Batch analysis accepts a CSV file with columns prompt and category, plus the same boolean and LTP parameter flags applied uniformly to all prompts.
Chat interface accepts a message history (JSON array of {role, content} objects), max_tokens (≤512), and optional analyze/analyze_response flags that trigger ASM/LTP/SFD analysis of the user message and/or the model's generated response.
The model registry (models.json) ships with four Qwen 2.5 pairs (0.5B, 1.5B, 3B, 7B). Custom HuggingFace base/instruct pairs can be added at runtime. The engine auto-detects layer count, attention head configuration, GQA grouping, and hidden dimensionality.
-
prompts.csv— a library of categorized prompts for the sidebar prompt picker -
baselines.csv— benign prompt bank used for length-normalized signal baselines -
punctuation_probes.csv— specialized probes for punctuation-sensitivity testing
Each analysis produces a PromptResult serialized to a JSON dictionary with the following field groups:
Identity and tokenization: prompt, category, tokens (list of decoded token strings), seq_len
ASM scalars: stress_score (float), net_correction (float), entropy, gini, top2_share, middle_share, interior_cv (all floats, distribution metrics), n_negative_tokens (int), has_negative_tokens (bool)
ASM arrays: per_token_stress (float array, length = seq_len), signed_attr (float array, length = seq_len), amplitude_trajectory (float array, length = 2 × n_layers for attn+MLP sublayers), amplitude_normalized (same shape), heatmap (2D array: sublayers × seq_len)
Behavioral divergence: kl_divergence (float or null), per_token_kl (float array or null), instruct_topk and base_topk (lists of [token_string, probability] pairs, up to 10 each), base_counterfactual_tokens (per-position top-k from base model)
LTP sub-object (ltp): profiles (list of numpy arrays per token), tension_magnitudes (float list), profile_shapes(string list: "steep"/"flat"/"inverted"), counterfactual_tokens (per-position ranked alternatives with probabilities), offset_magnitude/offset_consistency/offset_variance/lateral_coverage (per-layer dictionaries), mean_M/mean_C/mean_V/mean_L (summary scalars), max_prc (peak rank concentration), n_directional (count of tokens with PRC above threshold), semantic_trajectory_2d/tension_trajectory_2d (PCA projections), k, layer_strategy, svd_rank, tuned_lens (configuration echo)
SFD sub-object (sfd): per_token_energy/per_token_entropy/per_token_density (float arrays), 12 prompt-level aggregates (energy_mean/max/var/p90, entropy_mean/max/var/p90, density_mean/max/var/p90), global_erank, n_layers_monitored, k
Rank displacement (rank_displacement): mean_matched, mean_replacement, mean_concentration, mean_tau(Kendall's tau), mean_overlap, per-position detail arrays
Full capture extras (when enabled): per_token_coherence (cross-layer direction agreement), per_token_spectral_rank, attn_frac (attention vs. MLP contribution ratio), token_similarity (token×token cosine similarity matrix)
Classification: classifiers (dictionary of classifier outputs), classification (predicted category and confidence)
Each analysis generates up to 14 plot types as base64-encoded PNGs:
ASM plots: signed attribution bar chart, focused stress bar chart, distribution metrics panel, amplitude trajectory line plot, token×layer heatmap.
LTP plots: lateral tension profiles (stacked magnitude), tension magnitude bar chart, dual trajectory (PCA), summary statistics panel, profile heatmap (token × rank).
SFD plots: density bar chart, energy bar chart, entropy bar chart, rank displacement chart.
Summary CSV (summary.csv): one row per analyzed prompt, ~40 columns of scalar metrics including all ASM, LTP, SFD, and rank displacement aggregates, plus classification results.
Full results JSON (results.json): complete per-prompt result objects with all per-token arrays, profiles, and trajectories. For 271 prompts with full capture enabled, this file is approximately 17 MB.
Aggregate statistics JSON (aggregate_statistics.json): batch-level computed analytics including per-category bootstrapped means with confidence intervals, pairwise separability (Cohen's d with bootstrap CIs for every metric between category pairs), length-correlation analysis, and cross-metric correlations.
Export ZIP (tasm_session_<timestamp>.zip): configurable archive containing any combination of summary CSV, full results JSON, aggregate statistics JSON, per-prompt PNG plots, comparative/dashboard plots, and a PDF report.
PDF report (via ReportLab): formatted batch analysis report with summary tables, separability analysis, and embedded visualizations.
All API endpoints return JSON. The single-analysis endpoint returns:
{
"ok": true,
"result": { /* full PromptResult dict */ },
"plot_keys": ["signed_attribution", "stress_per_token", ...],
"session_n": 42,
"cache_size_bytes": 1048576
}
Plots are served separately via GET /api/plots/{plot_key} and GET /api/plots/individual/{index}/{plot_key}as base64-encoded PNG strings.
The server maintains global mutable state (loaded models, activation caches, session data) protected by two threading locks: _analysis_lock serializes forward passes and session writes to prevent activation cache corruption; _loading_lock makes model loading state transitions atomic. Batch analysis runs in a background daemon thread, with progress communicated via a shared log list polled by the frontend.
- Language: Python 3.10+
- Framework: FastAPI + Uvicorn
- ML stack: PyTorch, HuggingFace Transformers, Safetensors, Accelerate
- Visualization: Matplotlib
- Statistics: NumPy, SciPy
- Reporting: ReportLab (PDF)
- Compute: CPU by default; GPU optional. The instruct model remains in memory; the base model is loaded on-demand for KL/LTP computation and immediately unloaded afterward.
- Memory: 4–32 GB depending on model scale (0.5B–7B parameters).