Mask-GCG: Optimized LLM Token Masking

• Mask-GCG is a dual-use approach that applies learnable token-masking to streamline adversarial jailbreak attacks and enable grammar-constrained decoding.
• The method employs mask learning and adaptive pruning to reduce computational overhead while preserving high-impact tokens for effective model manipulation.
• Empirical evaluations show up to 40% token compression and 17% runtime reduction, highlighting improved efficiency and interpretability in LLM outputs.

Mask-GCG refers to two distinct but independently significant advances in LLM prompt optimization: (1) learnable token-masking for adversarial suffix pruning in jailbreak attacks, and (2) efficient grammar-constrained decoding via token masks for legal output generation. Each context deploys “Mask-GCG” to enhance either attack efficiency and interpretability or decoding flexibility and speed within structured syntactic domains.

1. Token-Masking for Adversarial Suffix Pruning in Jailbreak Attacks

Mask-GCG, as introduced in the context of adversarial jailbreak attacks, is a plug-and-play extension to Greedy Coordinate Gradient (GCG) optimization for adversarial suffixes. Jailbreak attacks attempt to circumvent LLM safety alignment by appending adversarially optimized suffixes to prompts, prompting the model to generate undesired (e.g., harmful) responses. Standard GCG and its variants optimize a fixed-length suffix of discrete tokens by repeated gradient-guided, position-wise token substitution, but are limited by redundancy, computational load scaling with suffix length ($L$), and the potential for longer suffixes to dilute attack signal and aid detection.

Mask-GCG augments any GCG procedure by introducing a learnable, per-position token mask $m \in \mathbb{R}^L$, hierarchically performing: (a) mask learning to assign importance to each suffix token, (b) interleaved discrete token optimization for high-impact positions, and (c) adaptive pruning of low-impact (low mask-probability) positions. Masking probabilities are defined by $P_i = \sigma(m_i/\tau)$ with temperature $\tau > 0$, so that tokens with $P_i \approx 0$ contribute negligibly to optimization gradients. The optimization minimizes a combined objective

$$\ell_{\mathrm{total}}(m,\,\text{suffix}) = \ell_{\mathrm{attack}}(\text{suffix}) + \lambda_{\mathrm{reg}} \cdot \ell_{\mathrm{reg}}(m)$$

where $$\ell_{\mathrm{attack}} = -\log P(y_{\mathrm{target}} | x_{\mathrm{prompt}},\,\text{suffix})$$, and $\ell_{\mathrm{reg}}$ penalizes nonzero mask values. The optimization interleaves Adam updates of $m$ with GCG steps on high-probability positions. After mask convergence, all positions $i$ with $P_i < \tau_{\mathrm{prune}}$ are pruned, with rollback if pruning increases $\ell_{\mathrm{attack}}$. The gradient and search space dimensionality, and thus computational overhead, shrink with each prune step, optimizing both efficiency and stealthiness.

2. Algorithmic Structure and Workflow

Contrasting standard GCG and Mask-GCG clarifies the impact of token-masking. GCG iteratively evaluates, for all $i = 1 \ldots L$ and candidate tokens $v$, the attack loss, updating positions holding the maximal loss reduction. Mask-GCG proceeds via the following sequence:

1. Suffix and mask logits $m$ are initialized using attention-guided importance.
2. For each optimization step $t$:
   • Compute mask probabilities $P = \sigma(m/\tau)$;
   • Form embeddings $e^{\mathrm{masked}} = P \odot e$ for input to the LLM;
   • Compute total loss (attack + mask regularization) and backpropagate to update $m$;
   • Run a GCG step for high-impact positions;
   • Every $f$ steps, prune low-probability positions ($P_i < \tau_{\mathrm{prune}}$), rollback if $\ell_{\mathrm{attack}}$ increases;
   • If the attack is successful, halt with the current suffix.

Pruning reduces the search and gradient space to $$|J_{\text{keep}}| = |\{i : P_i \geq \tau_{\mathrm{prune}}\}|$$, directly decreasing per-step overhead in gradient computation, candidate sampling, and loss evaluation, with rollback ensuring monotonic progress with respect to $\ell_{\mathrm{attack}}$.

3. Empirical Performance and Compression Ratios

Mask-GCG has been evaluated on Llama-2-7B-Chat, Vicuna-7B, and Llama-2-13B-Chat using GCG, I-GCG, and AmpleGCG, with toxic queries sampled from AdvBench. The key metrics are Attack Success Rate (ASR), Suffix Compression Ratio (SCR), and reduction in time cost.

Model	GCG+Mask-GCG SCR@20	GCG+Mask-GCG SCR@30	I-GCG+Mask-GCG SCR@20	I-GCG+Mask-GCG SCR@30	AmpleGCG+Mask-GCG SCR@20	AmpleGCG+Mask-GCG SCR@30
Llama-2-7B	5.8%	9.9%	0.2%	0.7%	2.0%	2.0%
Vicuna-7B	1.4%	2.1%	0.3%	1.1%	6.5%	4.1%
Llama-2-13B	5.2%	10.5%	4.1%	5.4%	5.1%	4.7%
Average	4.1%	7.5%	1.5%	2.4%	4.5%	3.0%

Mask-GCG preserves or slightly improves ASR for all tested settings. Average SCR (fraction of pruned tokens) ranges from 1.5–7.5%, with up to 40% observed in individual cases. Average runtime reduction approaches 17%, confirming that pruning low-impact tokens abates computational cost, with no deleterious effect on attack effectiveness (Mu et al., 8 Sep 2025).

4. Interpretability and Token Redundancy Analysis

Analysis of the learned mask distribution demonstrates that a substantial majority of suffix positions—exceeding 83% for tested models and settings—are classified as “high-impact,” i.e., essential for attack success ($P_i \geq \tau_{\mathrm{prune}}$). However, a nontrivial tail of positions consistently shows negligible importance, and these can be pruned without impairing attack effectiveness or increasing loss. The mask highlights model-attentive regions within the prompt, exposing internal vulnerabilities and providing direct insight into which prompt fragments steer generative outcomes. This interpretability is useful for both attack understanding and defensive mechanism construction, suggesting regularizers that can penalize or filter atypically sparse or dense token-importance spectra.

5. Implications for Efficiency, Defense, and Broader Impacts

Mask-GCG exemplifies a lightweight, extensible framework that can be applied to any GCG variant, immediately yielding a search space contraction, acceleration of attack generation (~17% reduction in time), and improved stealth by avoiding unnecessary suffix elongation. As pruning is rollback-safe, no increase in $\ell_{\mathrm{attack}}$ is permitted. The underlying mask learning for discrete suffixes bridges the methodological gap between continuous-space adversarial optimization and discrete prompt engineering. From a defense perspective, knowledge of high- and low-impact token loci can inform detection of adversarial attempts (by identifying redundancy patterns) and underpin mask-based regularization for defensive fine-tuning of LLMs. A plausible implication is that this mask-based approach can generalize to other settings where structured adversarial prompt optimization over discrete input spaces is needed.

6. Mask-GCG in Grammar-Constrained Decoding (GCD) and LLM Output Control

In a separate domain, Mask-GCG refers to a grammar-constrained decoding algorithm, as instantiated in the GreatGramma tool. The goal is to ensure that LLM outputs comply with a context-free grammar (CFG) by constructing token masks at each decoding step, ruling out tokens that would lead to ungrammatical structures. The formalism composes a detokenizer FST ($T_V$), a character-level lexing FST ($T_A$), and a token-level FST ($T_{A \circ V}$), maintaining state in both the lexer and the parser (PDA). The token mask at time $t$ is computed as: $$V_{\mathrm{allowed}} = A(q^A_t, q^P_t) \cup \bigcup_{\alpha\in D(q^A_t,q^P_t)} \{v \mid q^A_t \xrightarrow{v:\alpha} q'_A\}$$ where $A$ and $D$ are precomputed tables mapping lexer and parser states to always-accepted and context-dependent terminal productions, respectively.

Key innovations include the token-spanner table for efficient realization of terminal sequences, stack-free overapproximation for rapid acceptance checking, and minimization of both offline precomputation and online per-token overhead. Offline preprocessing for GreatGramma is $17.71\times$ faster than the previous state-of-the-art (24–35 s for typical vocabularies and grammars), and online token masking is performed in 5–32 ms per token (Park et al., 7 Feb 2025). The approach enables LLMs to support arbitrary and dynamic CFGs, supporting applications in code generation, structured data output, and safe text templating.

7. Related Work, Trade-offs, and Future Directions

Both deployments of Mask-GCG illustrate a common theme: selective, per-token masking—either learned (adversarial suffixes) or grammar-driven (GCD)—can drastically reduce unnecessary computation and improve controllability of LLM outputs. In jailbreak attack optimization, most suffix tokens are indispensable for attack success, but identifying and pruning token-level redundancy streamlines attacks and exposes interpretability vectors. In GCD, token masks guarantee syntactic legality efficiently, even for large vocabularies and grammars with dynamic or complex rules.

Trade-offs in Mask-GCG for GCD involve slightly higher offline cost compared to minimal approaches, but this is compensated by tractable online latency and support for arbitrary grammars. In adversarial optimization, incremental pruning must be balanced with rollback mechanisms to avoid loss of attack potency. In both settings, the mask formalism yields robust practical performance, modular implementation, and potential for downstream defense and reliability improvements in LLM pipelines.

The broader implication is that learnable or computed token masking is poised to become an integral component in both adversarial robustness studies and structured output control for LLMs, opening further research on mask-based regularization, interactive prompt auditing, and adversarial attack detection frameworks.