"""Architecture-aware preset defaults for optimal abliteration. Detects the model's architecture class (dense vs MoE, standard vs reasoning) and returns research-grounded parameter overrides that maximize refusal removal while preserving coherence. Research grounding: - SAFEx (NeurIPS 2025): Safety in MoE concentrated in <0.2% of experts - Cracken AI (2025): Global abliteration fails on large MoE; domain-specific works - Korinsky (2025): MoE abliteration damages reasoning; dense does not - L3 (Feb 2026): Expert silencing <20% achieves 70.4% ASR on MoE - Rannaberg (2025): Abliteration fails on DeepSeek R1 distills - Young (Dec 2025): Single-pass projection preserves GSM8K better than iterative - DECCP: -0.13pp GSM8K avg vs Heretic: -7.81pp (single-pass wins) """ from __future__ import annotations import logging import re from dataclasses import dataclass, field from enum import Enum from typing import Any logger = logging.getLogger(__name__) class ArchitectureClass(Enum): """Detected architecture classification.""" DENSE = "dense" SMALL_MOE = "small_moe" # <100B total params (e.g. Qwen3-30B-A3B, Mixtral-8x7B) LARGE_MOE = "large_moe" # >=100B total (e.g. DeepSeek-V3, Kimi K2, Qwen3-235B) class ReasoningClass(Enum): """Whether the model has chain-of-thought / thinking capabilities.""" STANDARD = "standard" REASONING = "reasoning" @dataclass class ArchitectureProfile: """Detected model architecture profile with recommended overrides.""" arch_class: ArchitectureClass reasoning_class: ReasoningClass # Detection metadata model_name: str = "" model_type: str = "" # HF config.model_type is_moe: bool = False num_experts: int = 0 # total experts per layer (0 = dense) num_active_experts: int = 0 # experts active per token total_params_b: float = 0.0 # total params in billions (estimated) num_layers: int = 0 hidden_size: int = 0 # Human-readable summary profile_label: str = "" # e.g. "Large MoE + Reasoning" profile_description: str = "" # explanation of why these defaults were chosen research_citations: list[str] = field(default_factory=list) # Recommended parameter overrides (method-level) recommended_method: str = "" method_overrides: dict[str, Any] = field(default_factory=dict) # Recommended breakthrough module configuration breakthrough_modules: dict[str, bool] = field(default_factory=dict) # ── MoE architecture identifiers ──────────────────────────────────────── # HF model_type values that indicate MoE architecture _MOE_MODEL_TYPES = { "mixtral", "qwen2_moe", "qwen3_moe", "deepseek_v2", "deepseek_v3", "dbrx", "grok", "jamba", "arctic", "olmoe", "switch_transformers", "nllb_moe", "llama4", } # Patterns in model name that indicate MoE (fallback when model_type is ambiguous) _MOE_NAME_PATTERNS = [ "moe", "mixtral", "-A3B", "-A22B", "MoE", "deepseek-v3", "gpt-oss", "kimi-k2", "glm-4.7", "step-3.5", "minimax-m2", "maverick", "scout", "mistral-large-3", "jamba", "olmoe", "arctic", ] # Name patterns that indicate MoE ONLY if no "distill" is present # (full DeepSeek-R1 is 671B MoE, but R1-Distill-* are dense) _MOE_NAME_PATTERNS_NO_DISTILL = [ "deepseek-r1", ] # Name-based heuristics for SMALL MoE (when no config is available). # These patterns identify models that are known to be small MoE (<100B total). # Without config, we can't detect expert count, so name matching is the fallback. _SMALL_MOE_NAME_PATTERNS = [ "-A3B", # Qwen3-30B-A3B, Qwen3-Next-80B-A3B (active = 3B) "gpt-oss", # GPT-OSS-20B (21B total, 3.6B active) "olmoe", # OLMoE-1B-7B (7B total) "mixtral-8x7b", # Mixtral-8x7B (47B total) "jamba", # Jamba models (52B total) ] # Name-based heuristics for known LARGE MoE (>=100B total). _LARGE_MOE_NAME_PATTERNS = [ "deepseek-v3", # DeepSeek-V3 (671B total) "deepseek-r1", # DeepSeek-R1 (671B total) "kimi-k2", # Kimi K2 (1T total) "-A22B", # Qwen3-235B-A22B "mistral-large-3", # Mistral Large 3 (675B total) "step-3.5", # Step-3.5 Flash (large MoE) "minimax-m2", # MiniMax-M2 (large MoE) ] # Patterns in model name that indicate reasoning / thinking capability. # Uses regex word-boundary matching to avoid false positives # (e.g. "olmo" containing "o1", "falcon3" containing "o3"). _REASONING_NAME_PATTERNS_RE = [ re.compile(r"(?:^|[-_/])r1(?:[-_/]|$)", re.IGNORECASE), # DeepSeek-R1 re.compile(r"think", re.IGNORECASE), # QwQ-Think, etc. re.compile(r"qwq", re.IGNORECASE), # QwQ re.compile(r"(?:^|[-_/])o1(?:[-_/]|$)", re.IGNORECASE), # OpenAI o1 re.compile(r"(?:^|[-_/])o3(?:[-_/]|$)", re.IGNORECASE), # OpenAI o3 ] # Distill patterns (reasoning distillations into dense models) _REASONING_DISTILL_PATTERNS = [ "r1-distill", ] # Config attributes for MoE detection — split into total vs active # to avoid confusing per-token count with total expert count. _TOTAL_EXPERT_ATTRS = [ "num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts", ] _ACTIVE_EXPERT_ATTRS = [ "num_experts_per_tok", "num_selected_experts", ] def detect_architecture( model_name: str, config: Any = None, num_layers: int = 0, hidden_size: int = 0, ) -> ArchitectureProfile: """Detect the architecture class and reasoning capability of a model. Args: model_name: HuggingFace model identifier config: HuggingFace AutoConfig object (optional, for precise detection) num_layers: Number of transformer layers (from ModelHandle) hidden_size: Hidden dimension size (from ModelHandle) Returns: ArchitectureProfile with detection results and recommended defaults """ model_type = "" is_moe = False num_experts = 0 num_active_experts = 0 total_params_b = 0.0 is_reasoning = False # ── Step 1: Extract info from config if available ──────────────── if config is not None: model_type = getattr(config, "model_type", "") # Check for MoE via config attributes for attr in _TOTAL_EXPERT_ATTRS: val = getattr(config, attr, None) if val is not None and val > 0: is_moe = True num_experts = max(num_experts, val) for attr in _ACTIVE_EXPERT_ATTRS: val = getattr(config, attr, None) if val is not None and val > 0: is_moe = True num_active_experts = max(num_active_experts, val) # Check model_type if model_type in _MOE_MODEL_TYPES: is_moe = True # Extract layer/hidden info from config if not provided if num_layers == 0: num_layers = getattr(config, "num_hidden_layers", 0) if hidden_size == 0: hidden_size = getattr(config, "hidden_size", 0) # Rough param estimation intermediate = getattr(config, "intermediate_size", hidden_size * 4) vocab = getattr(config, "vocab_size", 32000) if num_layers > 0 and hidden_size > 0: per_layer = 4 * hidden_size * hidden_size + 3 * hidden_size * intermediate if is_moe and num_experts > 0: # MoE: multiply FFN part by num_experts ffn_part = 3 * hidden_size * intermediate attn_part = 4 * hidden_size * hidden_size per_layer = attn_part + ffn_part * num_experts embedding = 2 * vocab * hidden_size total_params_b = (per_layer * num_layers + embedding) / 1e9 # ── Step 2: Name-based detection (fallback / supplement) ───────── name_lower = model_name.lower() if not is_moe: for pattern in _MOE_NAME_PATTERNS: if pattern.lower() in name_lower: is_moe = True break if not is_moe: # Check patterns that only apply when "distill" is NOT in the name has_distill = "distill" in name_lower if not has_distill: for pattern in _MOE_NAME_PATTERNS_NO_DISTILL: if pattern.lower() in name_lower: is_moe = True break # Reasoning detection for pattern in _REASONING_DISTILL_PATTERNS: if pattern.lower() in name_lower: is_reasoning = True break if not is_reasoning: for pattern_re in _REASONING_NAME_PATTERNS_RE: if pattern_re.search(name_lower): is_reasoning = True break # ── Step 3: Classify architecture ──────────────────────────────── if is_moe: # Classification priority: # 1. If total params known → use param threshold (100B) # 2. Else if expert count known → use expert threshold (16) # 3. Else fall back to name patterns → default SMALL_MOE (conservative) if total_params_b > 0: is_small = total_params_b < 100 elif num_experts > 0: is_small = num_experts <= 16 else: # No config available — use name heuristics. # Check large patterns first (more specific). is_small = True for pattern in _LARGE_MOE_NAME_PATTERNS: if pattern.lower() in name_lower: is_small = False break arch_class = ArchitectureClass.SMALL_MOE if is_small else ArchitectureClass.LARGE_MOE else: arch_class = ArchitectureClass.DENSE reasoning_class = ( ReasoningClass.REASONING if is_reasoning else ReasoningClass.STANDARD ) # ── Step 4: Build profile with recommended defaults ────────────── profile = ArchitectureProfile( arch_class=arch_class, reasoning_class=reasoning_class, model_name=model_name, model_type=model_type, is_moe=is_moe, num_experts=num_experts, num_active_experts=num_active_experts, total_params_b=total_params_b, num_layers=num_layers, hidden_size=hidden_size, ) _apply_recommended_defaults(profile) return profile def _apply_recommended_defaults(profile: ArchitectureProfile): """Fill in recommended method, overrides, and breakthrough modules. All recommendations are grounded in 2025-2026 abliteration research. """ arch = profile.arch_class reasoning = profile.reasoning_class # ── Dense + Standard ───────────────────────────────────────────── if arch == ArchitectureClass.DENSE and reasoning == ReasoningClass.STANDARD: profile.profile_label = "Dense Standard" profile.profile_description = ( "Dense decoder-only model. Single-pass projection is optimal " "(Young 2025: DECCP -0.13pp GSM8K). Linear refusal geometry is " "well-studied. Anti-Ouroboros maps self-repair for clean removal. " "Spectral Certification verifies completeness." ) profile.research_citations = [ "Young 2025 (arXiv:2512.13655): single-pass preserves GSM8K", "Arditi et al. 2024: refusal is a single direction in dense models", ] profile.recommended_method = "aggressive" profile.method_overrides = { # Single-pass is better for dense (Young 2025) "refinement_passes": 1, } profile.breakthrough_modules = { "anti_ouroboros": True, "spectral_cert": True, "riemannian": False, # Dense manifolds are flat "conditional": False, # Not needed for global removal "wasserstein_transfer": False, } # ── Dense + Reasoning ──────────────────────────────────────────── elif arch == ArchitectureClass.DENSE and reasoning == ReasoningClass.REASONING: profile.profile_label = "Dense Reasoning" profile.profile_description = ( "Dense reasoning model (e.g. R1 distill, OLMo-Think). Multi-stage " "alignment resists single-direction abliteration (Rannaberg 2025). " "Needs more directions (12-16) and iterative refinement (4-6 passes). " "Anti-Ouroboros is critical — reasoning models self-repair by " "literally reasoning about the missing refusal. Riemannian detects " "curved thinking-chain refusal geometry. Conditional addresses " "over-refusal (FalseReject COLM 2025)." ) profile.research_citations = [ "Rannaberg 2025: abliteration fails on R1 distills", "FalseReject (COLM 2025): reasoning models over-refuse", "Perplexity R1 1776: post-training succeeds where abliteration fails", ] profile.recommended_method = "aggressive" profile.method_overrides = { "n_directions": 12, "refinement_passes": 4, "use_jailbreak_contrast": True, "use_chat_template": True, "safety_neuron_masking": True, } profile.breakthrough_modules = { "anti_ouroboros": True, # Most important — reasoning self-repair "riemannian": True, # Thinking chain curves refusal surface "conditional": True, # Addresses reasoning over-refusal "spectral_cert": True, # Expect RED initially, drives iteration "wasserstein_transfer": False, } # ── Small MoE + Standard ──────────────────────────────────────── elif arch == ArchitectureClass.SMALL_MOE and reasoning == ReasoningClass.STANDARD: profile.profile_label = "Small MoE Standard" profile.profile_description = ( "Small MoE model (e.g. Qwen3-30B-A3B, Mixtral-8x7B, GPT-OSS-20B). " "Safety concentrated in <0.2% of experts (SAFEx NeurIPS 2025). " "Surgical per-expert targeting is optimal. Expert transplant very " "low (0.05) or OFF — fewer experts means less headroom. " "Conditional abliteration enables domain-specific removal." ) profile.research_citations = [ "SAFEx (NeurIPS 2025): 12/6144 experts carry safety in Qwen3-30B", "Korinsky 2025: MoE abliteration damages reasoning", "Cracken AI 2025: domain-specific abliteration works on MoE", ] profile.recommended_method = "surgical" profile.method_overrides = { "n_directions": 4, "refinement_passes": 2, "per_expert_directions": True, "invert_refusal": False, "expert_transplant": False, # Fewer experts = less headroom "transplant_blend": 0.05, "project_embeddings": False, # Cascades through router unpredictably "regularization": 0.05, # Small reg protects shared layers } profile.breakthrough_modules = { "anti_ouroboros": True, "conditional": True, # Domain-specific removal "spectral_cert": True, "riemannian": False, # Small MoE — not enough curvature "wasserstein_transfer": False, } # ── Large MoE + Standard ──────────────────────────────────────── elif arch == ArchitectureClass.LARGE_MOE and reasoning == ReasoningClass.STANDARD: profile.profile_label = "Large MoE Standard" profile.profile_description = ( "Large MoE model (e.g. DeepSeek-V3, Kimi K2, Qwen3-235B). " "Global abliteration has ZERO effect (Cracken AI on Kimi K2 1T). " "Must use surgical per-expert targeting. Conditional abliteration " "is the #1 technique — proven 0% target refusal + 100% non-target " "preservation. Riemannian needed for 'more sophisticated refusal " "geometry' in shared layers." ) profile.research_citations = [ "Cracken AI 2025: global abliteration zero effect on Kimi K2", "Cracken AI 2025: domain-specific gets 0% cyber refusal, 100% explicit preserved", "L3 (Feb 2026): <20% expert silencing achieves 70.4% ASR", "SAFEx (NeurIPS 2025): HCDG/HRCG expert taxonomy", ] profile.recommended_method = "surgical" profile.method_overrides = { "n_directions": 4, # Per-expert, not global "refinement_passes": 2, "per_expert_directions": True, "layer_adaptive_strength": True, # Different MoE layers vary wildly "invert_refusal": False, "expert_transplant": True, "transplant_blend": 0.10, # Light touch preserves specialization "project_embeddings": False, # Cascades through router "regularization": 0.05, "attention_head_surgery": True, # Shared attention carries signal } profile.breakthrough_modules = { "conditional": True, # #1 technique for MoE "anti_ouroboros": True, # Expert-level ASRG "riemannian": True, # Shared layers have curved geometry "spectral_cert": True, "wasserstein_transfer": False, } # ── Small MoE + Reasoning ─────────────────────────────────────── elif arch == ArchitectureClass.SMALL_MOE and reasoning == ReasoningClass.REASONING: profile.profile_label = "Small MoE Reasoning" profile.profile_description = ( "Small MoE with reasoning (e.g. Qwen3-30B-A3B in think mode). " "Most fragile combination — MoE expert specialization extends into " "reasoning (Korinsky 2025). Gentle surgical approach. Stop at first " "GREEN spectral cert to avoid over-ablation." ) profile.research_citations = [ "Korinsky 2025: MoE abliteration damages reasoning substantially", "SAFEx (NeurIPS 2025): safety concentrated in few experts", "FalseReject (COLM 2025): reasoning models over-refuse", ] profile.recommended_method = "surgical" profile.method_overrides = { "n_directions": 6, "refinement_passes": 3, "per_expert_directions": True, "use_jailbreak_contrast": True, "use_chat_template": True, "invert_refusal": False, "expert_transplant": False, # Too risky for reasoning MoE "transplant_blend": 0.05, "project_embeddings": False, "regularization": 0.05, "safety_neuron_masking": True, } profile.breakthrough_modules = { "conditional": True, # #1 for MoE "anti_ouroboros": True, "spectral_cert": True, # Run per-pass, stop at GREEN "riemannian": False, # Small model — overhead not worth it "wasserstein_transfer": False, } # ── Large MoE + Reasoning ─────────────────────────────────────── elif arch == ArchitectureClass.LARGE_MOE and reasoning == ReasoningClass.REASONING: profile.profile_label = "Large MoE Reasoning" profile.profile_description = ( "Large MoE reasoning model (e.g. DeepSeek-R1 671B). The hardest " "category. Global abliteration fails AND multi-stage alignment " "resists direction removal. Gentle surgical precision at expert " "level + reasoning-aware iterative deepening. Over-ablation kills " "reasoning — stop at first GREEN cert." ) profile.research_citations = [ "Cracken AI 2025: global abliteration fails on large MoE", "Rannaberg 2025: abliteration fails on R1 distills", "Korinsky 2025: MoE abliteration damages reasoning", "L3 (Feb 2026): expert silencing is the viable attack surface", ] profile.recommended_method = "surgical" profile.method_overrides = { "n_directions": 8, "refinement_passes": 3, "per_expert_directions": True, "use_jailbreak_contrast": True, "use_chat_template": True, "layer_adaptive_strength": True, "invert_refusal": False, "expert_transplant": True, "transplant_blend": 0.08, # Very light for reasoning preservation "project_embeddings": False, "regularization": 0.05, "safety_neuron_masking": True, "attention_head_surgery": True, } profile.breakthrough_modules = { "conditional": True, # #1 technique "anti_ouroboros": True, # Expert+layer ASRG "riemannian": True, # Curved shared layers "spectral_cert": True, # Per-pass, stop at GREEN "wasserstein_transfer": False, } else: # Fallback — should not happen, but be safe profile.profile_label = "Unknown" profile.profile_description = "Could not classify architecture. Using safe defaults." profile.recommended_method = "advanced" profile.method_overrides = {} profile.breakthrough_modules = { "anti_ouroboros": False, "riemannian": False, "conditional": False, "spectral_cert": False, "wasserstein_transfer": False, } logger.info( f"Architecture profile: {profile.profile_label} " f"(MoE={profile.is_moe}, experts={profile.num_experts}, " f"reasoning={reasoning.value}, ~{profile.total_params_b:.1f}B params)" ) def get_profile_summary(profile: ArchitectureProfile) -> str: """Return a human-readable markdown summary of the detected profile.""" lines = [ f"**Detected Profile:** {profile.profile_label}", "", f"**Architecture:** {'MoE' if profile.is_moe else 'Dense'}" + (f" ({profile.num_experts} experts, {profile.num_active_experts} active)" if profile.is_moe else ""), f"**Reasoning:** {'Yes' if profile.reasoning_class == ReasoningClass.REASONING else 'No'}", f"**Est. Params:** {profile.total_params_b:.1f}B" + (f" | Layers: {profile.num_layers} | Hidden: {profile.hidden_size}" if profile.num_layers else ""), "", f"**Recommended Method:** `{profile.recommended_method}`", "", profile.profile_description, ] if profile.research_citations: lines.append("") lines.append("**Research basis:**") for cite in profile.research_citations: lines.append(f"- {cite}") overrides = profile.method_overrides if overrides: lines.append("") lines.append("**Key parameter overrides:**") for k, v in overrides.items(): lines.append(f"- `{k}`: {v}") modules = profile.breakthrough_modules enabled = [k for k, v in modules.items() if v] disabled = [k for k, v in modules.items() if not v] if enabled: lines.append("") lines.append(f"**Breakthrough modules enabled:** {', '.join(enabled)}") if disabled: lines.append(f"**Breakthrough modules disabled:** {', '.join(disabled)}") return "\n".join(lines) def apply_profile_to_method_config( profile: ArchitectureProfile, base_config: dict[str, Any], ) -> dict[str, Any]: """Apply architecture profile overrides to a method config dict. Takes the base method config (from METHODS[method_key]) and applies the profile's recommended overrides on top. Explicit user overrides still take precedence (handled by AbliterationPipeline.__init__). Args: profile: Detected architecture profile base_config: Base method configuration dict Returns: New config dict with profile overrides applied """ result = dict(base_config) for key, value in profile.method_overrides.items(): # Always set the override — some keys (e.g., use_jailbreak_contrast, # safety_neuron_masking) may not exist in the base method config but # are valid pipeline parameters needed by the UI auto-detect path. result[key] = value return result