Andrii Bidochko

• Updated: August 16, 2026
• 7 min read

Localizing Safety Alignment: MLP Layers and Mid-Network Blocks Encode Refusal Behavior in Large Language Models

Direct Answer

The paper “Localizing Safety Alignment: MLP Layers and Mid‑Network Blocks Encode Refusal Behavior in Large Language Models” demonstrates that refusal behavior—when a model declines to answer unsafe or malicious prompts—is concentrated in a small subset of MLP parameters, especially in the middle layers of the network. This finding matters because it reveals a concrete, manipulable locus for improving safety alignment without retraining the entire model.

Background: Why This Problem Is Hard

Large language models (LLMs) are deployed in chatbots, virtual assistants, and autonomous agents that interact with end‑users in real time. Safety alignment—ensuring the model refuses or redirects harmful requests—is a prerequisite for any commercial rollout. Yet, practitioners observe two persistent challenges:

• Distributed safety signals: Traditional fine‑tuning treats safety as a property spread across all weights, making it difficult to pinpoint which parts of the network actually enforce refusal.
• Brittle refusal performance: Small prompt variations or adversarial phrasing can bypass safety filters, suggesting that the underlying safety mechanisms are fragile and possibly localized.

Existing approaches—reinforcement learning from human feedback (RLHF), instruction tuning, or post‑hoc classifiers—address the symptom (the model’s output) rather than the root cause within the model’s architecture. Consequently, they require massive data, expensive compute, and still produce unpredictable gaps in refusal behavior.

What the Researchers Propose

The authors introduce a systematic “weight‑transplant” methodology to locate safety‑relevant parameters. Instead of retraining, they take a safety‑aligned model (the “source”) and selectively copy its weights into a matched, unaligned base model (the “target”). By varying the granularity of the transplanted components—attention heads, individual MLP layers, contiguous layer blocks, and predefined MLP blocks—they can measure how much each component contributes to refusal performance.

Key components of the framework include:

• Model pairs: Two open‑weight families (e.g., Llama‑2‑7B and its aligned counterpart) that share architecture but differ in safety fine‑tuning.
• Granular transplantation levels: From single‑layer MLP weights to multi‑layer blocks (e.g., layers 8‑11).
• Greedy search algorithm: An iterative procedure that selects the next most beneficial block based on improvement in refusal scores.
• Safety benchmarks: Four curated test suites covering malicious prompts, benign over‑refusal, and mixed‑intent queries.

How It Works in Practice

The workflow can be visualized as a three‑stage pipeline:

1. Preparation: Align a base LLM using standard safety fine‑tuning (e.g., RLHF). Preserve a copy of the unaligned base model for comparison.
2. Transplantation: Replace a chosen subset of weights in the unaligned model with the corresponding weights from the aligned model. The replacement can target:
   • Attention matrices (Q, K, V, O)
   • MLP feed‑forward matrices (up‑projection, down‑projection)
   • Contiguous layer ranges (e.g., layers 1‑4)
   • Pre‑defined MLP blocks (e.g., layers 8‑11)
3. Evaluation: Run the transplanted model on the safety benchmarks. Record refusal rates, false‑positive over‑refusal, and any degradation in benign task performance.

The novelty lies in treating safety alignment as a modular property that can be “plugged in” or “plugged out” at the parameter level. This contrasts sharply with monolithic fine‑tuning, where the entire weight matrix is altered, obscuring the contribution of any individual component.

Evaluation & Results

The authors conducted experiments on two model families (each with a base and an aligned version) across four safety benchmarks:

• Malicious‑prompt refusal: Ability to say “I can’t help with that” when faced with harmful instructions.
• Benign over‑refusal: Avoiding unnecessary rejections on safe queries.
• Mixed‑intent prompts: Testing nuanced cases where intent is ambiguous.
• OR‑Bench cross‑validation: A hold‑out benchmark to assess generalization.

Key observations:

• MLP dominance: Swapping only the MLP weights recovered at least 2.7× more malicious‑prompt refusals than swapping attention weights, consistently across both model families.
• Mid‑network concentration: In every greedy search, the block spanning layers 8‑11 was selected first, indicating that the middle of the feed‑forward stack houses the most safety‑relevant parameters.
• Non‑additive composition: Adding more aligned blocks sometimes reduced overall refusal performance, suggesting interference effects between blocks.
• Selective superiority: Certain subsets of blocks outperformed a full MLP transplantation on specific metrics (e.g., higher malicious refusal with lower benign over‑refusal).
• Benchmark‑dependent ordering: The greedy block order derived from one safety benchmark did not always transfer optimally to another, highlighting a precision‑coverage trade‑off.

These results collectively argue that safety alignment is not uniformly distributed; instead, a relatively small, interaction‑sensitive region of the network governs refusal behavior.

Why This Matters for AI Systems and Agents

Understanding the localized nature of safety alignment unlocks several practical pathways for developers and product teams:

• Targeted safety patches: Instead of full‑scale re‑training, engineers can fine‑tune or replace only the identified MLP block, reducing compute costs and deployment latency.
• Modular safety layers: System architects can design “safety plug‑ins” that swap in pre‑aligned MLP blocks on demand, enabling dynamic risk management for different user segments.
• Improved evaluation pipelines: By focusing tests on the mid‑network region, QA teams can detect alignment regressions faster and with fewer resources.
• Agent orchestration benefits: Autonomous agents that call LLMs can request a “safety‑enhanced” version of the model for high‑risk tasks, while using the base model for low‑risk queries, optimizing both performance and cost.

These capabilities align directly with the needs of enterprises building AI‑driven products. For example, the UBOS platform overview offers a modular architecture where safety‑enhanced model blocks can be injected into existing workflows without disrupting downstream services.

Similarly, developers integrating LLMs into messaging bots can leverage the ChatGPT and Telegram integration to route high‑risk user inputs through a safety‑augmented model, while keeping casual conversations lightweight.

What Comes Next

While the study provides compelling evidence for MLP‑centric safety, several open challenges remain:

• Generalization across architectures: The experiments focused on decoder‑only transformers. Future work should test encoder‑decoder models, vision‑language hybrids, and emerging sparsity‑based architectures.
• Dynamic block selection: Developing a runtime system that decides which block to activate based on prompt risk could make safety truly adaptive.
• Interaction with other alignment objectives: How do these mid‑network safety blocks interact with factuality, helpfulness, or style‑control fine‑tuning?
• Robustness to adversarial attacks: Attackers may learn to target the uncovered layers; defensive strategies need to anticipate such pressure.

Addressing these questions will likely involve a blend of interpretability research, reinforcement learning, and systems engineering. Companies interested in pioneering these solutions can explore the Workflow automation studio to prototype dynamic safety routing, or partner through the UBOS partner program to co‑develop custom safety modules.

For a deeper dive into the methodology and raw data, consult the original arXiv paper. The authors also release their code and model checkpoints, enabling the community to replicate and extend the findings.

Visual Illustration

The diagram below visualizes the transplantation pipeline, highlighting the mid‑network MLP block (layers 8‑11) that carries the bulk of refusal behavior.

Conclusion

Safety alignment is not an amorphous property spread evenly across a language model; it is heavily localized in the feed‑forward MLP layers, especially in the middle of the network. By exploiting this concentration, practitioners can achieve more efficient, interpretable, and potentially more robust safety interventions. The work opens a new research frontier that blends model interpretability with practical safety engineering, offering a roadmap for building trustworthy AI agents at scale.

Call to Action

If your organization is building AI‑driven products that require rigorous safety guarantees, consider experimenting with modular MLP block swaps on your own models. Leverage the Enterprise AI platform by UBOS to manage model versions, conduct safety evaluations, and integrate the resulting safety‑enhanced models into production pipelines.

Andrii Bidochko

CTO UBOS

Andrii Bidochko is an AI entrepreneur and researcher focused on AI agents, reinforcement learning, and autonomous systems. He writes about the technologies shaping the future of machine intelligence, from frontier models and agent architectures to real-world AI applications.