Files
OBLITERATUS/obliteratus/architecture_profiles.py
T
2026-03-04 12:38:18 -08:00

585 lines
25 KiB
Python

"""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