Source-linked AI summary
MOMAT: Mixture of Multiple Atlases for Low-Power Jailbreak Defense of Quantized LLMs
Boyang Li, Bingyu Shen, Weihao Hong, Zhiyuan Jiang, Xinlei Guan, Yan Ma, Miles Q. Li, Yi Sheng, Ruiyang Qin
TL;DR
Quantization weakens qLLM safety alignment, while edge constraints and heterogeneous safety databases challenge conventional safeguards. MOMAT combines domain-localized atlases, MoE scoring, and CiM retrieval. The paper reports reduced attack success and no added over-refusal under W4A8, alongside large retrieval latency and energy reductions.
Problem
Quantized LLMs can lose alignment robustness, while heterogeneous safety databases and edge resource limits challenge reliable, efficient jailbreak defense.
Source: §1 · p. 1, §2.3 · p. 3Method
MOMAT organizes safety samples and policy templates into semantic atlases, retrieves top-k features across them, and applies lightweight MoE scoring with CiM-accelerated similarity search.
Source: §1 · p. 2, §3.6.2 · p. 6Results
Takeaways & Limitations
The reported results show strong safety enforcement and benign utility preservation together under edge constraints.
Source: §5 · p. 9The cited retrieval limitation concerns monolithic RAG: heterogeneous indices increase retrieval noise, and DRAM-based search imposes memory-bandwidth demands incompatible with real-time edge enforcement.
Source: §2.3 · p. 3
Abstract
from arXiv · showhide
Quantized large language models are increasingly deployed on edge devices for their low latency and energy efficiency. However, model quantization weakens alignment safeguards, leaving qLLMs (quantized large language models) highly vulnerable to jailbreak attacks. To address this challenge, we present MOMAT (Mixture of Multiple Atlases), a hardware-enhanced safety framework that combines structured knowledge retrieval with low-power defense acceleration. Each atlas represents a semantic cluster of harmful or benign sample sets and policy templates, enabling domain-localized Retrieval-Augmented Generation guarding that mitigates the curse of dimensionality and the resulting semantic sparsity problem in large, heterogeneous safety databases. MOMAT retrieves top-$k$ similarity features from all atlases for each prompt and evaluates them using a lightweight MoE (Mixture of Experts) detector, while a CiM (Compute-in-Memory)-accelerated similarity engine performs fast, low-power atlas-local retrieval. MOMAT's CiM-based retrieval accelerates a 100-query batch from 15,052.44 ms to 3,207.21 ns (a $4.69 \times 10^6\times$ speedup) and reduces energy from $8.1 \times 10^7$ $\mu$J to 3.32 $\mu$J, yielding an approximately $2.5 \times 10^5\times$ energy reduction over DRAM-based (Raspberry Pi) baselines. Red-team evaluations across standard benchmarks show that MOMAT matches the defense performance of state-of-the-art methods while avoiding benign overkill and providing substantial efficiency gains, demonstrating that CiM-based modular defenses can make edge-deployed qLLMs both safer and more energy-efficient. We will release the full 223.2k-sample dataset to foster future research.
1. Introduction
Quantization weakens qLLM alignment just as edge constraints limit dedicated safety mechanisms, while monolithic safety retrieval can struggle with heterogeneous databases. MOMAT combines semantically coherent atlases, lightweight safety scoring, and CiM retrieval; its reported results include large retrieval-efficiency gains and strong benchmark defense without added benign over-refusal.
1. Introduction
- Quantization degrades qLLM alignment, while edge devices often lack dedicated safety mechanisms and monolithic RAG can retrieve irrelevant templates from heterogeneous databases.The introduction links quantization-related vulnerability and edge constraints with retrieval noise in large safety databases.Source: §1 · pp. 1–2
- The paper identifies atlas-based retrieval and CiM acceleration as contributions toward recovering alignment robustness under edge latency and energy constraints.It also plans to release a 223.2k-sample dataset for future research.Source: §1 · pp. 1–2
- MOMAT organizes harmful and benign samples and policy templates into semantically coherent atlases, retrieving top-k features across atlases for a lightweight MoE safety scorer.The framework is presented as domain-localized retrieval for edge-deployed qLLMs.Source: §1 · p. 2
- 15,052.44 ms to 3,207.21 ns for 100 queries, with energy falling from 8.1 × 10^7 μJ to 3.32 μJ; MOMAT also reports state-of-the-art benchmark defense without added benign over-refusal.The retrieval comparison is against Raspberry Pi–class DRAM baselines; the benchmark claim is reported separately.Source: §1 · p. 2
2. Background
The background connects quantization to degraded safety alignment, describes jailbreak and prompt-injection threats, and identifies weaknesses in large heterogeneous safety-retrieval databases. It also motivates CiM as a way to reduce the memory-transfer costs that burden edge retrieval.
2.1. Quantization and Alignment Collapse
- Quantization errors and activation clipping degrade safety-relevant representations, leaving quantized models more susceptible to jailbreaks, prompt injection, and multi-turn steering, particularly in edge deployments without auxiliary safeguards.The background attributes the damage to disrupted feature geometry, degraded safety subspaces, and affected safety layers.Source: §2.1 · p. 3
2.2. Jailbreak and Prompt Injection Threat
- Jailbreak and prompt-injection attacks use linguistic cues, including role-playing, multi-turn steering, translation, and adversarial paraphrase, while quantized models can misclassify adversarial inputs as benign.Prompt injections may override system messages or hijack in-context exemplars without changing model weights.Source: §2.2 · p. 3
2.3. Limitations of Safety RAG: Curse of Dimensionality and Semantic Sparsity
- As safety databases become heterogeneous, sparse embeddings and cross-domain interference undermine nearest-neighbor reliability, producing irrelevant or overly generic templates; monolithic DRAM indices also demand bandwidth poorly suited to edge hardware.These limitations motivate domain-localized retrieval.Source: §2.3 · p. 3
2.4. Edge Deployment Constraints and CiM-Based Accelerators
- CiM performs multiply-accumulate operations in memory arrays, aligning retrieval's vector-matrix search with in-memory computation; MOMAT stores atlas-specific exemplar matrices in dedicated arrays for compact, localized search.The background presents this design as a hardware basis for edge latency and energy budgets.Source: §2.4 · pp. 3–4
3.1. System Overview
MOMAT combines atlas-local retrieval, a lightweight intent detector, and CiM-accelerated response-template retrieval to protect edge-deployed quantized LLMs. It separates costly atlas preparation from compact online inference.
3.1. System Overview
- MOMAT uses four modules: CiM retrieval, atlas-specific safety stores, a lightweight MoE intent detector, and a CiM-accelerated RAG engine.The RAG engine supports both MoE scoring and safe-template retrieval when harmful intent is detected.Source: §3.1 · p. 4
- A frozen encoder produces a 1024-dimensional prompt embedding; the MoE combines it with atlas-local top-k features to make a binary safety decision.Harmful prompts receive a safety response, while other prompts pass to the qLLM.Source: §3.1 · p. 4
- Expensive corpus processing, clustering, and ICL expansion occur offline, leaving embedding and a compact MoE pass for online inference.CiM arrays produce atlas-local similarity signals in situ, avoiding DRAM-based vector search.Source: §3.1 · p. 4
3.2. Atlas Construction
MOMAT builds domain-labeled atlases from generated harmful–benign pairs and expanded variants. Dense, semantically coherent neighborhoods are intended to reduce retrieval noise and interference from unrelated domains.
3.2. Atlas Construction
- Seed prompts generate cross-domain harmful–benign pairs, which are recursively expanded into variants and assembled into domain-labeled atlas corpora.Expansion includes paraphrasing, syntactic variation, semantic abstraction, and detailed expansions.Source: §3.2 · p. 4
- Dense neighborhoods support stable CiM similarity search, while partitioning unrelated domains into separate atlases avoids high-dimensional interference.The passage states atlas-local RAG outperforms monolithic retrieval in both accuracy and efficiency.Source: §3.2 · p. 4
- Atlas embeddings and statistics are precomputed offline, leaving prompt encoding, CiM retrieval, and MoE scoring for online inference.Source: §3.2 · p. 4
3.3. Atlas Retrieval and Lightweight Intent Detection
MOMAT retrieves atlas-local similarity features for each prompt and uses a lightweight MoE to score safety. Low-dimensional routing and compact experts keep online processing small for edge hardware.
3.3. Atlas Retrieval and Lightweight Intent Detection
- For each prompt, CiM retrieves atlas-local top-k similarity features from all atlases, and the MoE consumes them alongside the prompt embedding.Atlas-local retrieval is described as robust to paraphrasing, obfuscation, and role-play because routing uses embedding-space geometry rather than surface phrasing.Source: §3.3 · p. 4
- The MoE comprises several lightweight MLP experts, while low-dimensional centroid comparisons keep routing and local scoring costs minimal.Routing selects which offline-prepared atlas-local features are forwarded to the main MoE scorer.Source: §3.3 · pp. 4–5
- Atlas construction uses cross-domain harmful–benign pairs from seed prompts, recursively expands variants, then assigns samples to domain-labeled atlases.Source: §3.3 · p. 5
3.4. CiM-Enhanced RAG Engine
The CiM engine supplies atlas-local retrieval features for safety scoring and performs targeted safe-template retrieval after harmful intent is detected. It computes similarities within storage arrays to limit data movement.
3.4. CiM-Enhanced RAG Engine
- For each query, CiM retrieves atlas-local top-k similarity features from all atlases in parallel to support MoE safety scoring.Source: §3.4 · p. 5
- When the MoE determines p_harm ≥ τ, CiM retrieves a safe response template from the most activated atlas.Source: §3.4 · p. 5
- Parallel multiply-accumulate operations compute dot-product similarities inside storage arrays, avoiding DRAM vector transfers and their bandwidth and energy overhead.Preloaded exemplar matrices support low-latency, low-energy retrieval in both engine roles.Source: §3.4 · p. 5
3.5. End-to-End Inference Pipeline
MOMAT embeds each prompt, gathers atlas-local retrieval summaries, and uses a gated expert ensemble to estimate harmfulness. Benign prompts pass to the qLLM; harmful prompts trigger targeted exemplar retrieval and safe-response composition.
3.5. End-to-End Inference Pipeline
- The pipeline computes an embedding, derives each atlas’s retrieval summary, and feeds the resulting feature representation to MLP experts and a learned gate.The experts produce harmfulness logits; softmax-normalized gating weights combine them into the final risk estimate.Source: §3.5 · p. 5
- All atlases are scored in parallel, and inference uses a second targeted CiM retrieval only when the prompt is classified as harmful.Atlas embeddings are precomputed and remain resident in CiM arrays, so neither retrieval pass requires DRAM data movement at inference time.Source: §3.5 · p. 5
- MOMAT’s inference path uses atlas-local summaries to score prompt harmfulness, then routes benign prompts to generation and harmful prompts to targeted retrieval and safe composition.The harmfulness estimate is formed from expert logits and passed through a sigmoid; the targeted retrieval selects exemplars from the most activated atlas.Source: §3.5 · p. 5
3.6. MOMAT for Safety Detection
MOMAT keeps prompt embeddings fixed and learns to interpret retrieval patterns across semantically distinct reference atlases. A lightweight gated expert ensemble scores those patterns for safety detection.
3.6. MOMAT for Safety Detection
- MOMAT uses a frozen encoder-only query embedding and learns to interpret retrieval structure rather than fine-tuning embeddings.The approach is framed as learning from patterns revealed by retrieval across multiple reference atlases.Source: §3.6 · p. 5
3.6.1. Multi-Atlas Architecture.
- Each atlas represents a semantic domain, and queries retrieve top-k examples from every atlas to form per-atlas scores and comparative features.The feature vector includes statistical summaries and pairwise atlas comparisons; the design extends to arbitrary atlas counts and can use fine-grained or hierarchical categories.Source: §3.6.1 · pp. 5–6
3.6.2. Mixture-of-Experts for Pattern-Specific Routing.
- MOMAT uses 3–5 MLP experts that specialize in neighborhood signatures, with a query-embedding-conditioned gate combining their safety logits.Training uses binary cross-entropy with a gating-balance regularizer to prevent expert collapse; the setup scales from binary to multi-domain atlases.Source: §3.6.2 · p. 6
4. Experiment
Experiments isolate atlas retrieval structure from CiM acceleration while testing safety, benign utility, and edge efficiency. MOMAT reaches 0.0% ASR across evaluated quantized models without increasing FRR, and CiM retrieval is substantially faster and lower-energy than DRAM retrieval.
4.1. Evaluation Settings
- Evaluation separates structural retrieval on Raspberry Pi from CiM hardware gains, testing two 7B qLLMs, domain-specific attacks, ASR, FRR, p95 latency, and energy.The test setup uses Llama-2-7B and Mistral-7B and harmful instruction categories aligned with MOMAT's atlas taxonomy.Source: §4 · pp. 6–7, §4.1 · p. 7
4.2.1. Safety Collapse of Small Quantized Models
- W4A8 raises 7B ASR by 4.4% on AdvBench and 3.0% on Malicious Instruct, while 13B changes by −0.7% and −2.0%, respectively.The 13B model's 26 GB FP16 footprint is impractical for most edge hardware, while 7B compresses to 3.76 GB under W4A8.Source: Table 1 · p. 7, §4.2.1 · p. 7
4.2.2. MOMAT Safety performance
- 0.0% ASR on both datasets for W4A8 Llama2-7B and Mistral-7B; MOMAT adds no FRR on either model.The evaluation also compares defenses on Llama2-7B and tests benign queries designed to be semantically close to adversarial ones.Source: Table 2 · p. 8, §4.2.2 · p. 8, Figure 2 · p. 8
4.2.3. Retrieval Latency and Energy Comparison: RPi DRAM Retrieval vs. MOMAT-CiM
- For 100 queries, CiM retrieval takes 3,207.21 ns versus 15,052.44 ms for Raspberry Pi DRAM; DRAM's optimal window is 500 vectors.The CiM timing comes from SRAM-based CiM results; the Raspberry Pi result is CPU/DRAM retrieval.Source: Table 4 · p. 8, §4.2.3 · p. 8, Table 3 · p. 8
- A 100-query batch uses 3.32 μJ on CiM versus approximately 8.1 × 10^7 μJ on Raspberry Pi, an approximately 2.5 × 10^5× reduction.These estimates are not directly comparable in absolute terms: Raspberry Pi energy is board-level, while CiM energy is circuit-simulation-derived.Source: Table 4 · p. 8, §4.2.3 · p. 9
5. Conclusion
MOMAT combines domain-localized atlases with CiM similarity search to address quantization-related safety degradation and retrieval inefficiency. The conclusion reports low ASR without added over-refusal, while identifying broader evaluation and adaptive atlas construction as future work.
5. Conclusion
- The motivating scope is edge-deployed qLLMs, where quantization degrades alignment and monolithic retrieval can fail to surface relevant safety exemplars.Source: §5 · p. 9
- Under W4A8, MOMAT reduces ASR to 0.0% on Llama2-7B and Mistral-7B with Δ FRR = +0.0%, alongside 4.69 × 10^6× latency and approximately 2.5 × 10^5× energy reductions over DRAM baselines.The conclusion attributes the safety result to domain-localized atlases and CiM-accelerated similarity search.Source: §5 · p. 9
- Future work will investigate adaptive atlas construction for evolving threat taxonomies and extend evaluation to more quantization schemes and on-device environments.Source: §5 · p. 9