PaperlayerView exact revision on arXiv

Source-linked AI summary

An Efficient and Modular Framework for Targeted Harm Mitigation in LLMS

Roberto Campbell, Momin Abbas, Muneeza Azmat, Michal Ulewicz, Raya Horesh, Kristjan Greenewald, Rogério Abreu de Paula, Nathalie Baracaldo

arXiv:2609.13624v3cs.LGcs.AIcs.CLcs.CY

TL;DR

LLMs can produce harmful responses, while existing alignment approaches can be costly and may trade off against task-specific performance. The paper combines aLoRA adapters with a router for targeted correction and reports >80% F1 across six harm categories, alongside constant time-to-first-token as multiturn conversations lengthen.

  • Problem

    Existing alignment approaches can be costly, and safety alignment may trade off against task-specific performance.

    Source: §1 · pp. 1–2

  • Method

    The method uses a learned router and harm-specific aLoRA adapters to route and correct harmful responses while reusing the base model KV cache.

    Source: §1 · p. 2, §3.1 · p. 4, §3.2 · p. 5

  • Results

    >80% F1 across six harm categories is reported for routing, and the method maintains constant time-to-first-token as multiturn conversations lengthen.

    Source: §1.1 · p. 2, §4.3 · p. 10

  • Takeaways & Limitations

    The approach offers modular alignment adapters that can be added or removed with minor router retraining, without retraining the original model.

    Source: §1 · p. 2

    The current alignment scope is a single harm, and inference-time cache-reuse benefits require adapters trained for the same model being checked.

    Source: §6 · p. 11

Abstract

from arXiv · show

Large Language Models (LLMs) are powerful zero-shot learners but remain prone to misalignment with human preferences, often producing biased, toxic, or otherwise harmful outputs. Existing alignment methods, while effective, are costly and tightly coupled to the model, limiting flexibility and scalability. We propose a modular correction framework that augments pretrained LLMs with Activated LoRA (aLoRA) adapters and a context-aware routing mechanism to eliminate harms from misaligned model responses. Our approach enables expert adapters to activate mid-sequence without invalidating the KV cache, allowing low-latency, targeted correction during generation. Each expert is trained to detect and mitigate specific harms, such as bias or toxicity. A learned router dynamically selects appropriate experts based on the models intermediate outputs. We demonstrate that our system improves alignment on standard safety benchmarks while preserving task performance, offering a lightweight and efficient path toward safer and more controllable LLM deployments.

1 Introduction

LLMs can produce harmful responses, while alignment can be costly and trade off against task-specific performance. The paper proposes modular aLoRA aligners and a router for targeted correction, with reported routing performance above 80% F1 across six harm categories.

1 Introduction

  • Alignment concerns include harmful responses such as bias, sexual content, violence, and unethical behavior.The paper focuses on harmful model responses rather than all dimensions of alignment, such as correctness or relevance.Source: §1 · p. 1
  • Existing alignment approaches can be costly, and safety alignment may trade off against task-specific performance, a tension called the alignment “tax.”Maintaining safety without interfering with tasks such as tool calling or math can require careful training-data ratios.Source: §1 · pp. 1–2
  • External alignment systems can avoid modifying the base model but may require multiple detectors and models, plus costly prefill for long contexts.The proposed router and adapters address the need to select a harm type without running multiple detectors on the same text.Source: §1 · p. 2
  • The method uses an aLoRA router to select harm-specific adapters, which activate mid-sequence and reuse the base model KV cache for low-latency correction.The base model remains unchanged, and the adapters support modular addition or removal with minor router retraining.Source: §1 · p. 2, §1.1 · p. 2

1.1 Contributions

  • The router achieves >80% F1 across six harm categories, and the approach is evaluated on BeaverTails, SafeRLHF, and HarmfulQA.The reported contribution also compares the approach with LoRA on multiturn conversations.Source: §1.1 · p. 2

2 Related Work

Related approaches separate alignment from the base model, adapt model weights with low-rank modules, or route tokens among experts. The paper positions its approach as routing harm-correction tasks to adapters while reusing the base model’s KV cache.

2 Related Work

  • Prior external aligners separate alignment from base-model generation, while classifier-based ensemble methods select an aligner by category; this approach instead uses parameter-efficient adapters.The cited discussion contrasts adapters served with a large base model against smaller BERT-like models.Source: §2 · p. 3
  • MAP incorporates diverse human preferences during the last phase of model training, whereas the paper says its approach can be applied to any open-weight model.The related work also notes that alignment can affect task-specific performance, a tension between performance and alignment.Source: §2 · p. 3
  • LoRA fine-tunes targeted weight matrices with trainable low-rank matrices while holding the base weights fixed.The adapter parameters are typically less than 1% of the base model parameters.Source: §2 · p. 3
  • Unlike standard LoRA, aLoRA applies adaptations only after an invocation sequence, allowing reuse of the base model’s prior-context KV cache.The related-work passage reports speedups as high as 30× for short generations with adapters.Source: §2 · pp. 2–3
  • Mixture-of-experts approaches route tokens among adapters or experts, whereas this work routes correction tasks to interpretable, harm-specific adapters.The related work says the compared approaches must recompute attention key-value pairs when activating adapters, unlike this method’s KV-cache reuse.Source: §2 · pp. 3–4
  • KV-cache sharing has been studied for multi-agent communication, with cited work noting an adversarial attack surface; this paper focuses on a single agent and leaves cache representations unchanged.The cited defenses limit sensitive-information leakage while preserving relevant cache information.Source: §2 · p. 4

3 Model Architecture

The architecture pairs harm-specific aligners with a router that directs responses to the appropriate aligner.

3 Model Architecture

3.1 Aligners

  • Each aligner is an aLoRA adapter fine-tuned for a specific harm, and the model can return the original response when no harm is detected.Source: §3.1 · pp. 4–5
  • Each aligner has a unique invocation sequence, with six harm-type names used to identify the corresponding adapter.The listed types are harm, profanity, sexual_content, social_bias, unethical_behavior, and violence.Source: §3.1 · p. 5

3.2 Router

  • The router is an aLoRA adapter with a trainable output head mapping to C classes; argmax selects an aligner, and classification and correction losses are summed.The described results freeze the linear output layer and train the router adapter.Source: §3.2 · p. 5
  • Each forward pass routes samples first, then processes homogeneous sub-batches with the selected aligner activation sequence.The router activation sequence is replaced with the corresponding aligner sequence before the aligner processes the model output.Source: §3.2 · p. 5
  • Table 1 reports sample counts by dataset and split.Source: Table 1 · p. 5

4 Experiments

Experiments evaluate data preparation, routing, alignment, and inference performance across three safety benchmarks. The trained router generally exceeds baselines, while aLoRA adapters reuse the base-model KV cache to reduce latency and maintain constant time-to-first-token as turns increase.

4 Experiments

  • The experiments compare models on three safety benchmarks and train them using an NVIDIA A100-SXM4-80GB GPU.Table 2 lists the baseline models used for benchmark evaluation.Source: §4 · p. 5, Table 2 · p. 6

4.1 Data

  • Data preparation creates harmful-response and safe-correction pairs: BeaverTails expands prompts with both labels, SafeRLHF keeps mixed-label pairs, and HarmfulQA keeps only the first response.After filtering, Granite Guardian 3.1 8b labels each dataset.Source: §4.1 · p. 6

4.2 Results

  • Three configurations compare frozen pre-finetuned aligners (Trained Router), trainable pre-finetuned aligners (Joint Trained Router), and joint training without pre-finetuning (Cold Start Joint), which lacks clear specialization.Joint training can suit data disjoint from pre-finetuning, such as proprietary aligner data added to the system.Source: §4.2 · pp. 7–8

4.2.1 Routing

  • Figure 9 compares weighted-average F1 across router configurations; trained routers outperform baselines except Gemma 12b instruct, which performs similarly (≥80%) on SafeRLHF and HarmfulQA.The Granite 2b router also demonstrates this result against a model with six times its parameter count.Source: Figure 9 · p. 8, §4.2.1 · pp. 8–9

4.3 Performance

  • aLoRA improves request latency as conversation length increases and maintains constant TTFT across turns, while LoRA TTFT scales with context length.The adapters reuse the base model’s KV cache; the LoRA variant recomputes it for both router and corrector calls.Source: Figure 11 · p. 9, Figure 10 · p. 9, §4.3 · p. 10

4.2.2 Alignment

  • Trained aLoRA aligners perform well against baselines on BeaverTails, SafeRLHF, and HarmfulQA, although baseline Granite models perform similarly.Figure 12 reports average win rates across all harms; the authors note that focusing on harm may miss response relevance.Source: §4.2.2 · p. 9, Figure 12 · p. 10

5 Conclusion

The proposed method routes model responses to the appropriate aligner using aLoRA adapters, reusing base-model KV cache entries to reduce inference compute and speed text generation. It performs well on three safety datasets, including strong routing performance from the Granite 3.3 2b model relative to a 12b-parameter Gemma baseline.

5 Conclusion

  • The method routes each model response to an appropriate aligner and uses aLoRA adapters to reuse base-model KV cache entries, reducing inference compute and speeding generation.Parameter-efficient adapters support scaling alignment criteria by adding adapters and retraining only the router.Source: §5 · p. 10

6 Limitations

The approach requires access to the base model, excluding black-box and API models unless providers offer the tuning approach as a service. Its serving and KV-cache benefits also require adapters trained for the same model being checked.

6 Limitations

  • Adapters cannot be used with black-box or API models unless their providers make the tuning approach available as a service.The method depends on parameter-efficient fine-tuning and access to the base model.Source: §6 · p. 11
  • Adapters transferred to a different model lose the inference-time serving and KV-cache reuse benefits.Those benefits require adapters trained for the model generating the text being checked.Source: §6 · p. 11
  • Alignment is limited to a single harm, although real-world responses may contain multiple harms.The authors identify multi-aligner responses across several dimensions as future work.Source: §6 · p. 11

A Datasets

The appendix identifies BeaverTails, SafeRLHF, and HarmfulQA data breakdowns by harm type, but the supplied passages do not provide their values.

A Datasets

  • BeaverTails data are broken down by harm type.The supplied caption does not state counts or proportions.Source: A Datasets · p. 14
  • SafeRLHF data are broken down by harm type.The supplied caption does not state counts or proportions.Source: A Datasets · p. 14
  • HarmfulQA data are broken down by harm type.The supplied caption does not state counts or proportions.Source: A Datasets · p. 14

B Experiment Details

The experiment details describe prompt-based routing for baseline models and a training prompt convention for aligned responses.

B.1.1 Routing

  • Baseline models score each harm class, and each datapoint is assigned to the class with the largest score.The routing prompts include examples for in-context learning for both base and instruct models.Source: §B.1.1 · p. 15, Table 3 · p. 15, Table 5 · p. 17
  • Tables 3 and 4 present routing prompts for base models, and Table 3 specifies an example for in-context learning.Source: Table 3 · p. 15, Table 4 · p. 16

B.1.2 Alignment

  • The alignment training prompt adds a <correctedmessage> tag to extract aligned responses when models also provide inconsistent explanations.The tag is described as a convenience for extracting the response.Source: Table 7 · p. 19

C.1 Routing

The appendix provides precision and recall results for BeaverTails, SafeRLHF, and HarmfulQA; the supplied captions do not report comparative outcomes.

C.1 Routing

  • Figure 14 presents precision results on BeaverTails.The caption does not state a winning method or values.Source: Figure 14 · p. 19
  • Figure 15 presents precision results on SafeRLHF.The caption does not state a winning method or values.Source: Figure 15 · p. 20
  • Figure 16 presents precision results on HarmfulQA.The caption does not state a winning method or values.Source: Figure 16 · p. 21
  • Figure 17 presents recall results for BeaverTails.The caption does not state a winning method or values.Source: Figure 17 · p. 22
  • Figure 18 presents recall results for SafeRLHF.The caption does not state a winning method or values.Source: Figure 18 · p. 23
  • Figure 19 presents recall results on HarmfulQA.The caption does not state a winning method or values.Source: Figure 19 · p. 24

C.2 Alignment

On Granite 3.3 2B, the methods match or outperform performance across all three datasets; on 8B, they perform similarly, with a small Trained Router improvement on SafeRLHF.

C.2 Alignment

  • For Granite 3.3 2B, the Trained Router method outperforms or matches performance on all three datasets: BeaverTails, SafeRLHF, and HarmfulQA.Source: §C.2 · p. 27
  • The reported alignment comparisons cover the Trained Router, Joint Trained Router, and Cold Start Joint across BeaverTails, SafeRLHF, and HarmfulQA at both model scales.The supplied prose reports these comparisons for Granite 3.3 2B and 8B.Source: §C.2 · p. 27
  • For Granite 3.3 8B, all three methods perform similarly, with the Trained Router showing a small improvement on SafeRLHF.The comparison covers BeaverTails, SafeRLHF, and HarmfulQA.Source: §C.2 · p. 27
Loading 2609.13624v3…