When Does Safety Activate? Temporal Analysis of LLM Alignment via Logit-Margin Score

Official code and data for "When Does Safety Activate? Temporal Analysis of LLM Alignment via Logit-Margin Score".

Overview

Standard safety evaluations (ASR) tell you what happened but not when or how safety mechanisms engage during generation. We introduce the Logit-Margin Score ($S_t$), a lightweight temporal monitor that tracks compliance–refusal balance at each decoding step using only standard logits.

Key contributions:

1. Temporal safety monitoring via $S_t$ trajectories (best ROC-AUC = 0.786)
2. Failure-mode taxonomy: Type A (pre-generation bias) / Type B (decoding-time interference) / Undetected, validated via permutation tests with BH-FDR correction
3. Construct validity: per-step alignment with hidden-state refusal directions (|r| up to 0.62)
4. Taxonomy validation: a threshold-free $t_{\text{cross}}$-gated early stopping policy confirms the taxonomy predicts intervention precision (Type A RP = 98.7% vs. Undetected 44.7%)

---

Repository Structure

Python Scripts

Core Generation & S_t Computation

Script	Description
logits_v6_updated.py	Main S_t extraction pipeline (v6). Runs multi-turn MCM attacks, computes per-step logit-margin scores, and saves per-sample trajectories and summary metrics. Supports Llama, Qwen, Mistral, Gemma.
logits_v7_extended.py	Extended pipeline (v7). Adds additional models (Mistral, Gemma) and lexicon perturbation variants (no_sorry, extended_refusal, extended_compliance, minimal).
logits_nanogcg.py	GCG attack pipeline. Runs nanoGCG suffix optimization and computes S_t trajectories under gradient-based adversarial suffixes.
logits_deepinception.py	DeepInception attack pipeline. Runs nested fictional scenario prompting (Li et al., 2023) and extracts S_t trajectories.
logits_benign_baseline.py	Benign baseline generation. Generates responses to attack-matched benign queries (same conversational format, safe content) for format-controlled comparison.
logits_decoy_eval.py	Decoy attack evaluation. Tests robustness against prefix-decoy (A), token-decoy (B), and pivot-decoy (C) manipulations that attempt to game early S_t values.
logits_stochastic_eval_v2.py	Stochastic decoding evaluation. Generates 10 trials per prompt across temperature (T∈{0.0, 0.3, 0.7, 1.0}) and nucleus sampling (p∈{0.9, 0.95, 1.0}) settings.
logits_local_perturb_eval.py	Lexicon perturbation evaluation. Tests S_t robustness under modified compliance/refusal lexicons.
compute_st_trajectory_metrics_v1.py	Metric computation from raw trajectories. Reads per-step S_t JSON files and computes summary metrics (early_k_mean, t_cross, sign_reversal, etc.).
attach_traj_to_csv.py	Trajectory attachment utility. Merges per-step S_t trajectory data into the main CSV for downstream analysis.

Analysis & Statistics

Script	Description
analyze_v9_comprehensive.py	Main analysis pipeline (v9). Aggregates all experimental conditions, computes ROC-AUC with bootstrap CIs, effect sizes (Cohen's d), Type A/B classification with RP normalization, permutation tests, and BH-FDR correction. Produces the authoritative all_samples.csv.
analyze_v9_multijudge.py	Multi-judge analysis. Compares Llama-Guard, HarmBench, and GPT-4o labels: inter-judge agreement, ASR comparison, Type A/B stability, and AUC consistency across judges.
analyze_decoy_results.py	Decoy robustness analysis. Evaluates AUC degradation under prefix/token/pivot decoy attacks (Appendix D).
analyze_stochastic_results_v2.py	Stochastic decoding analysis. Computes prompt-level AUC (median aggregation), ICC variance decomposition, response diversity, and greedy vs. stochastic comparison (Appendix G).
analyze_lexicon_topk.py	Top-k aggregation sensitivity. Evaluates S_t stability across k∈{5, 10, 25, 50} (Appendix F).
analyze_hidden_vs_logit.py	Hidden-state vs. logit comparison. Bridge analysis comparing logit-based S_t, hidden-state baselines (B1 refusal direction, B2 linear probe), and score-level fusion (Appendix E).
hidden_baselines.py	Hidden-state baseline extraction. Extracts refusal direction projections (B1) and trains linear probes (B2) on hidden states for comparison with logit-based S_t.
entropy_analysis.py	Entropy cross-validation. Computes generation-time entropy as a non-lexical baseline and validates Undetected classification (Appendix S).
entropy_harmful_vs_benign.py	Entropy harmful vs. benign analysis. Compares entropy distributions between harmful and benign responses.
learned_baseline_and_permtest.py	Learned baseline & permutation tests. Logistic regression on temporal features (5-fold CV), and permutation tests for Type A/B with BH-FDR correction (Appendices K, I).
judge_annotation_selector.py	Human annotation sample selector. Selects 100 borderline samples (high judge disagreement) for targeted human annotation.

LLM-as-Judge

Script	Description
llama_guard_judge.py	Llama-Guard-3-8B judge pipeline. Evaluates generated responses for safety violations across 14 categories (S1–S14). Outputs is_success_llm labels. Uses JudgeResponseLogger for comprehensive JSONL output.

Visualization

Script	Description
paper_figures_final_v5.py	Paper figure generation (v5). Generates Figure 1 (intro: state space, early-k AUC, sign reversal), Figure 3 (results: state space, ROC, Type A/B diverging bar), and standalone analysis plots.
draw_method_figure.py	Method overview figure. Generates figure2_method.png illustrating the S_t computation pipeline and Type A/B decomposition.

Utilities

Script	Description
a.py	Quick utility/test script.

Data Directories

Raw Experimental Results (Results_*)

Each Results_* directory contains per-sample JSON files with full S_t trajectories (64 decoding steps) and summary metrics.

Directory	Model	Attack	Content	Notes
Results_v6_llama_harmful_manual	Llama	MCM	Harmful	Main experiment
Results_v6_llama_nanogcg_harmful	Llama	GCG	Harmful	Main experiment
Results_di_llama_harmful	Llama	DeepInception	Harmful	Main experiment
Results_v6_llama_benign_manual	Llama	MCM	Benign	Format-matched baseline
Results_v6_llama_benign	Llama	GCG	Benign	Format-matched baseline
Results_v6_llama_decoyA	Llama	MCM+DecoyA	Harmful	Prefix-decoy robustness
Results_v6_llama_decoyB	Llama	MCM+DecoyB	Harmful	Token-decoy robustness
Results_v6_llama_decoy_base	Llama	MCM+Base	Harmful	Decoy control
Results_v6_llama_pivotC	Llama	MCM+PivotC	Harmful	Pivot-decoy robustness
Results_v6_llama_manual_k5/k25/k50_*	Llama	MCM	Harmful	Top-k ablation (k=5,25,50)
Results_v6_qwen_*	Qwen	MCM/GCG	Both	Main + decoy experiments
Results_v6_mistral_*	Mistral	GCG	Both	Main experiments
Results_v6_gemma_*	Gemma	GCG	Both	Main experiments
Results_v7_gemma/mistral_*_manual	Gemma/Mistral	MCM	Both	Extended model coverage
Results_v7_llama_baselines	Llama	MCM	Harmful	Hidden-state baseline data
Results_v7_llama_stoch_*	Llama	MCM	Harmful	Stochastic decoding (4 conditions)
Results_v7_lexicon	Llama	MCM	Harmful	Lexicon perturbation variants
Results_stochastic_temporal_v7	Llama	MCM	Harmful	Temporal metrics under stochastic decoding
Results_topk_sim_llama_*	Llama	All 3	Harmful	API top-k truncation simulation

Aggregated Analysis Results (analysis_*)

Directory	Description
analysis_v9_llamaGuard/	Primary: All 660+660 samples with Llama-Guard-3-8B labels. Contains all_samples.csv and analysis figures.
analysis_v9_HarmBench/	Validation judge: HarmBench classifier labels.
analysis_decoy_all/	Decoy robustness analysis results (all models).
analysis_decoy_qwen/	Decoy robustness analysis (Qwen-specific).
analysis_entropy_*	Entropy cross-validation results.
analysis_hidden_baselines/	Hidden-state baseline comparison (Llama).
analysis_hidden_baselines_qwen/	Hidden-state baseline comparison (Qwen).
analysis_hidden_vs_logit/	Bridge analysis: logit vs. hidden-state (Llama).
analysis_hidden_vs_logit_qwen/	Bridge analysis (Qwen).
analysis_quick_llama_gcg/	Quick diagnostic analysis (Llama+GCG).
analysis_quick_qwen_gcg/	Quick diagnostic analysis (Qwen+GCG).
analysis_stochastic_v2/	Stochastic decoding analysis results.

Other

Directory	Description
Data/	Input data (JailbreakBench prompts, GCG suffixes).
scripts/	Shell scripts for batch job submission.
logs/	Experiment execution logs.
figures_early_k_v3/	Early-k AUC progression figures.
earlyk_json_plots_v2/	Per-condition early-k trajectory plots.
induced_lexicon_v1/	Data-driven lexicon induction experiment.
lexicon_induction_results/	Induced lexicon transfer analysis results.
paper_figures_final_v5_LlamaGuard/	Generated paper figures (Llama-Guard labels).
paper_figures_final_v5_HarmBench/	Generated paper figures (HarmBench labels).
stochastic_question_level_results/	Prompt-level stochastic decoding results.
summary/	Summary tables and statistics.
_legacy/	Deprecated code and results.

Other Files

File	Description
figure2_method.png	Method overview figure (used in paper Figure 2).
labels_all_conditions.csv	Consolidated outcome labels across all conditions.
llm_judge_cache.json	Cached LLM-as-judge responses (avoids re-evaluation).

---

Quick Start

1. Generate S_t trajectories

# MCM attack on Llama
python logits_v6_updated.py --model llama --mode harmful --memory_mode manual

# GCG attack on Llama
python logits_nanogcg.py --model llama

# DeepInception attack on Llama
python logits_deepinception.py --model llama

# Benign baseline
python logits_benign_baseline.py --model llama

2. Run LLM-as-Judge

python llama_guard_judge.py --results_dir Results_v6_llama_harmful_manual

3. Aggregate and Analyze

python analyze_v9_comprehensive.py

4. Generate Paper Figures

python paper_figures_final_v5.py \
    --csv-path analysis_v9_llamaGuard/all_samples.csv \
    --output-dir paper_figures_final_v5_LlamaGuard

---

Models

Model	Parameters	Reference
Llama-3.1-8B-Instruct	8B	arXiv:2407.21783
Mistral-7B-Instruct-v0.3	7B	arXiv:2310.06825
Qwen2.5-7B-Instruct	7B	arXiv:2412.15115
Gemma-2-9B-Instruct	9B	arXiv:2408.00118

Requirements

torch
transformers
numpy
pandas
matplotlib
scikit-learn
scipy

Citation

License