# How to Add Support for New Model Architectures in Heretic: A Complete Guide

> Learn how to add support for new model architectures in Heretic by modifying three key Python functions. This guide simplifies integrating custom models into the Heretic framework.

- Repository: [Philipp Emanuel Weidmann/heretic](https://github.com/p-e-w/heretic)
- Tags: how-to-guide
- Published: 2026-02-19

---

**To add support for new model architectures in Heretic, you must modify three functions in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py): `get_model_class` to select the correct AutoModel class, `get_layers` to expose the transformer layer stack, and `get_layer_modules` to identify the attention and MLP projection weights for abliteration.**

Heretic is an open-source tool for "abliterating" harmful capabilities from large language models using LoRA-based weight editing. If you are working with the p-e-w/heretic repository and need to add support for new model architectures—whether encoder-decoder models like T5, vision-language models, or custom transformer layouts—you only need to adapt three specific functions in the model abstraction layer.

## Step 1: Register the New Model Class in `get_model_class`

Heretic determines which Hugging Face AutoModel class to instantiate by inspecting the pretrained configuration. The `get_model_class` function (lines 36-44 in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py)) currently checks for `vision_config` to distinguish between multimodal and causal language models.

### Detecting Architecture Types from Config

Add a new conditional branch to detect your architecture's unique configuration signature. For example, to support sequence-to-sequence models:

```python
from transformers import AutoModelForSeq2SeqLM  # NEW import

def get_model_class(
    model: str,
) -> Type[AutoModelForImageTextToText] | Type[AutoModelForCausalLM] | Type[AutoModelForSeq2SeqLM]:
    configs = PretrainedConfig.get_config_dict(model)

    # Existing vision model check

    if any([("vision_config" in config) for config in configs]):
        return AutoModelForImageTextToText
    
    # NEW: Check for encoder-decoder architecture

    if any([("decoder" in config) for config in configs]):
        return AutoModelForSeq2SeqLM
    
    return AutoModelForCausalLM

```

This ensures `Model.__init__` can call `get_model_class(settings.model).from_pretrained(...)` without raising a `ClassNotFoundError`.

## Step 2: Locate Transformer Layers in `get_layers`

Once the model loads, Heretic must access the list of transformer layers to apply edits. The `get_layers` method (lines 12-25 in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py)) handles unwrapping PEFT adapters and navigating common layer hierarchies.

### Handling Different Layer Hierarchies

If your architecture stores layers in a non-standard location (e.g., `model.encoder.block` instead of `model.model.layers`), add a new fallback:

```python
def get_layers(self) -> ModuleList:
    model = self.model

    # Unwrap PEFT wrapper if present

    if isinstance(model, PeftModel):
        model = model.base_model.model

    # Existing: Multimodal path (e.g., Llava)

    with suppress(Exception):
        return model.model.language_model.layers

    # Existing: Text-only path (e.g., Llama, Mistral)

    with suppress(Exception):
        return model.model.layers

    # NEW: Encoder-decoder or encoder-only layouts (e.g., T5)

    with suppress(Exception):
        return model.encoder.block

    raise AssertionError("Unable to locate transformer layers for the loaded model.")

```

This method is invoked by the optimizer, analyzer, and residual-geometry utilities, so returning the correct `ModuleList` is critical for the rest of the pipeline to function.

## Step 3: Map Attention and MLP Modules in `get_layer_modules`

Abliteration specifically targets the **attention output projection** (`attn.o_proj`) and **MLP down-projection** (`mlp.down_proj`) for LoRA-based editing. The `get_layer_modules` function (lines 26-74 in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py)) discovers these sub-modules using a helper called `try_add`.

### Identifying Projection Layers for Abliteration

Different architectures use varying attribute names for these projections. Extend the discovery logic with additional `try_add` calls wrapped in `suppress` blocks:

```python
def get_layer_modules(self, layer_index: int) -> dict[str, list[Module]]:
    layer = self.get_layers()[layer_index]
    modules: dict[str, list[Module]] = {}

    def try_add(component: str, module: Any):
        if isinstance(module, Module):
            modules.setdefault(component, []).append(module)
        else:
            assert not isinstance(module, Tensor), (
                f"Unexpected Tensor in {component} - expected nn.Module"
            )

    # Standard attention output projection (Llama, Mistral, etc.)

    try_add("attn.o_proj", layer.self_attn.o_proj)

    # Standard MLP down projection

    with suppress(Exception):
        try_add("mlp.down_proj", layer.mlp.down_proj)

    # ---------- NEW ARCHITECTURE EXAMPLES ----------

    
    # Example 1: Alternative attention naming (e.g., output_proj instead of o_proj)

    with suppress(Exception):
        try_add("attn.o_proj", layer.self_attn.output_proj)

    # Example 2: T5-style DenseReluDense blocks

    with suppress(Exception):
        try_add("mlp.down_proj", layer.DenseReluDense.wo)

    # Example 3: Custom expert modules in MoE architectures

    with suppress(Exception):
        for expert in getattr(layer, "custom_experts", []):
            try_add("mlp.down_proj", expert.down_proj)

    # -----------------------------------------------------------------

    
    total_modules = sum(len(v) for v in modules.values())
    assert total_modules > 0, "No abliterable modules found in layer"
    return modules

```

The `abliterate` method iterates over `self.get_layer_modules(layer_index).items()`, so adding these mappings enables Heretic to perform weight edits on the new architecture.

## Complete Working Example: Adding T5 Support

Here is a consolidated example showing how to add support for **T5-base**, an encoder-decoder model:

```python

# src/heretic/model.py

from transformers import AutoModelForSeq2SeqLM  # NEW import

def get_model_class(model: str):
    configs = PretrainedConfig.get_config_dict(model)
    
    if any([("vision_config" in cfg) for cfg in configs]):
        return AutoModelForImageTextToText
    if any([("decoder" in cfg) for cfg in configs]):  # NEW

        return AutoModelForSeq2SeqLM
    return AutoModelForCausalLM

def get_layers(self) -> ModuleList:
    model = self.model
    if isinstance(model, PeftModel):
        model = model.base_model.model
    
    with suppress(Exception):
        return model.model.language_model.layers
    with suppress(Exception):
        return model.model.layers
    with suppress(Exception):
        return model.encoder.block  # NEW: T5 encoder blocks

    
    raise AssertionError("Unable to locate transformer layers")

def get_layer_modules(self, layer_index: int):
    layer = self.get_layers()[layer_index]
    modules = {}
    
    def try_add(component: str, module: Any):
        if isinstance(module, Module):
            modules.setdefault(component, []).append(module)
    
    # Standard attempts

    with suppress(Exception):
        try_add("attn.o_proj", layer.self_attn.o_proj)
    with suppress(Exception):
        try_add("mlp.down_proj", layer.mlp.down_proj)
    
    # T5-specific mappings

    with suppress(Exception):
        try_add("attn.o_proj", layer.SelfAttention.o)  # NEW

    with suppress(Exception):
        try_add("mlp.down_proj", layer.DenseReluDense.wo)  # NEW

    
    assert sum(len(v) for v in modules.values()) > 0
    return modules

```

With these changes, you can run Heretic with `--model t5-base` and the optimizer will successfully load the model, enumerate its encoder layers, and apply LoRA-based abliteration to the attention and MLP projections.

## Key Files and Functions Reference

| File | Responsibility | Key Functions to Modify |
|------|---------------|------------------------|
| [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py) | Model loading and layer abstraction | `get_model_class` (lines 36-44), `get_layers` (lines 12-25), `get_layer_modules` (lines 26-74) |
| [`src/heretic/config.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/config.py) | Configuration schema | Only modify if adding new architecture-specific settings |
| [`src/heretic/utils.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/utils.py) | Helper utilities | Generally unchanged for new architectures |

## Summary

- **Register the model class** by extending `get_model_class` in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py) to return the appropriate `AutoModel...` class based on configuration keys.
- **Expose the layer stack** by adding fallback logic to `get_layers` to handle non-standard layer hierarchies like `model.encoder.block`.
- **Map the projections** by adding `try_add` calls in `get_layer_modules` to identify attention output and MLP down-projection weights regardless of naming conventions.
- **Test with a real model** such as T5 to verify that Heretic can load, analyze, and abliterate the new architecture.

## Frequently Asked Questions

### Do I need to modify the abliteration logic itself to support new architectures?

No. The abliteration algorithms in Heretic are architecture-agnostic and operate on the generic `attn.o_proj` and `mlp.down_proj` components returned by `get_layer_modules`. As long as you correctly map these projections in [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py), the existing LoRA-based editing and residual analysis code will work without modification.

### How do I handle models with different attention mechanisms like multi-query attention?

Multi-query attention (MQA) and grouped-query attention (GQA) typically reuse the same `o_proj` (output projection) attribute names as standard multi-head attention. You should verify the exact attribute path in your model's Hugging Face implementation, then add a `try_add("attn.o_proj", layer.self_attn.o_proj)` variant (or `layer.self_attn.q_proj` if targeting queries) wrapped in a `suppress(Exception)` block to handle the specific naming convention.

### Can I add support for quantized models like GPTQ or AWQ?

Yes, but you typically do not need to modify [`src/heretic/model.py`](https://github.com/p-e-w/heretic/blob/main/src/heretic/model.py) for quantization support. Heretic uses `transformers.AutoModel...` classes which automatically handle quantization config via `from_pretrained`. Ensure your quantization config is passed during model loading in the configuration settings. The layer enumeration and module discovery logic remains the same because quantization wrappers preserve the underlying attribute structure.

### What if my model uses a custom HuggingFace architecture not in the standard library?

For custom architectures, you must ensure the model class is importable in your Python environment (e.g., via `trust_remote_code=True`). Then follow the same three-step process: register the custom class in `get_model_class`, add the layer path in `get_layers`, and map the projection attributes in `get_layer_modules`. You may need to inspect the custom model's `modeling_*.py` file to find the exact attribute names for attention and MLP components.