ATWU: Alternating Token-Weighted Unlearning

• ATWU is a machine unlearning framework that learns token forget-specificity by balancing forget and retain objectives in autoregressive language models.
• It employs a lightweight linear scorer over hidden states to modulate token-level loss, ensuring structural tokens remain intact during unlearning.
• The method alternates between model parameter updates and scorer training, achieving superior unlearning quality with minimal impact on overall language competence.

Alternating Token-Weighted Unlearning (ATWU) is a framework for machine unlearning in autoregressive LLMs that treats token importance as a latent quantity to be learned jointly with model parameters during unlearning. Its central premise is that forget examples are internally heterogeneous: some tokens are forget-specific and encode the content to be removed, while others are structural tokens that support general language behavior and should not be suppressed indiscriminately. ATWU formalizes this distinction through the interaction between forget and retain objectives, and implements it with a lightweight linear scorer over hidden states trained in alternation with the LLM itself, without external token-level supervision (Yüce et al., 4 Jun 2026).

1. Retain conflict as the basis of forget-specificity

ATWU is formulated for an autoregressive LLM $\mathcal{M}_\theta$ with a forget set $\mathcal{D}_F$ and a retain set $\mathcal{D}_R$. The motivating observation is that a forget example may contain both answer-bearing tokens—such as names, dates, quantities, identifiers, and attribute-bearing phrases—and structural tokens such as function words, punctuation, and generic scaffolding. Uniformly applying a forget objective to all tokens can therefore remove target content while also degrading retained behavior and general language competence (Yüce et al., 4 Jun 2026).

The framework makes this intuition precise by defining forget-specificity through retain conflict. Let the retain objective be

$R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$

For forget tokens, ATWU introduces weights $z_{x,t}\in[0,1]$, intended to approximate the oracle indicator that a token is forget-specific. The token-weighted objective is

$$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$

The key quantity is the retain conflict of a subset of forget tokens. Writing $\mathcal{I}_F$ for all forget-token positions and $\ell^i_{\mathrm{forget}}(\theta)$ for the forget loss on token $i$, the conflict of $A\subseteq\mathcal{I}_F$ is

$\mathcal{D}_F$0

where

$\mathcal{D}_F$1

A token with small singleton conflict can be forgotten while staying close to the retain optimum; a token with large conflict behaves structurally. This shifts token selection away from probability heuristics and toward a criterion defined by the retain–forget trade-off itself (Yüce et al., 4 Jun 2026).

2. Oracle support recovery and relaxed token selection

ATWU introduces a separation condition under which the true forget-specific support is identifiable. The forget-token index set $\mathcal{D}_F$2 is partitioned into $\mathcal{D}_F$3, with oracle labels

$\mathcal{D}_F$4

such that

$\mathcal{D}_F$5

for some $\mathcal{D}_F$6 and $\mathcal{D}_F$7. Under this retain-conflict separation condition, the exact support-learning problem is

$\mathcal{D}_F$8

where $\mathcal{D}_F$9 is a token budget. The hard budget prevents the trivial all-zero and all-one selectors (Yüce et al., 4 Jun 2026).

The recovery theorem states that, if $\mathcal{D}_R$0 and $\mathcal{D}_R$1 is an integer, then any global minimizer $\mathcal{D}_R$2 satisfies

$\mathcal{D}_R$3

If $\mathcal{D}_R$4, then

$\mathcal{D}_R$5

In the paper’s interpretation, the joint objective can therefore recover the oracle forget-specific token support under a natural separation condition (Yüce et al., 4 Jun 2026).

Because discrete optimization over $\mathcal{D}_R$6 is combinatorial, ATWU also uses a continuous surrogate: $\mathcal{D}_R$7 with binary entropy

$\mathcal{D}_R$8

The entropy term encourages near-binary token decisions, while the budget penalty keeps the average selected-token mass near $\mathcal{D}_R$9 (Yüce et al., 4 Jun 2026).

3. Scorer parameterization and the practical ATWU objective

The practical ATWU method replaces free token variables by a shared scorer over token representations. Let $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$0 be the hidden representation for token $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$1. ATWU uses a linear scoring head

$R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$2

where $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$3 is initialized at

$R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$4

This makes all initial token scores equal to $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$5, so training begins without token preference and becomes selective only through optimization (Yüce et al., 4 Jun 2026).

On the forget side, ATWU adopts saturated negative cross-entropy rather than raw gradient ascent. For $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$6,

$R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$7

where $R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$8 is negative cross-entropy. The saturation factor bounds the forget loss below and stabilizes updates. The paper refers to this as SatGA (Yüce et al., 4 Jun 2026).

The empirical ATWU objective further injects the learned score into both the multiplicative weight and the saturation exponent, producing SatGA$R(\theta) \;\defeq\; \sum_{x \in \mathcal{D}_R}\sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta), \qquad R^\star \;\defeq\; \min_{\theta \in \Theta} R(\theta).$9: $z_{x,t}\in[0,1]$0 Here $z_{x,t}\in[0,1]$1 abbreviates $z_{x,t}\in[0,1]$2. For binary scores, SatGA$z_{x,t}\in[0,1]$3 reduces to standard SatGA. The added exponent modulation attenuates uncertain token updates early in training, before the scorer has converged (Yüce et al., 4 Jun 2026).

4. Alternating optimization and training dynamics

ATWU alternates between updating the LLM parameters $z_{x,t}\in[0,1]$4 and the scorer parameters $z_{x,t}\in[0,1]$5. The rationale is that the scorer determines which tokens receive forget pressure, while the model determines the hidden-state geometry on which the scorer operates. Simultaneous lockstep updates can therefore create an unstable feedback loop (Yüce et al., 4 Jun 2026).

The procedure is as follows. The scorer is initialized at $z_{x,t}\in[0,1]$6. Hidden states are computed for forget tokens under the current model, and scores are assigned by

$z_{x,t}\in[0,1]$7

In the model step, $z_{x,t}\in[0,1]$8 is fixed and token scores are detached, so they act as fixed coefficients while $z_{x,t}\in[0,1]$9 is updated using the retain loss and current token-weighted forget loss. In the scorer step, $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$0 is frozen and gradients flow only through $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$1. The method then repeats these two updates rather than jointly co-updating all parameters (Yüce et al., 4 Jun 2026).

The update schedule is empirically important. The appendix reports that refreshing the scorer every step performs poorly; the headline setting refreshes it every $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$2 model steps, and $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$3 is also competitive. Joint lockstep updates are clearly worse. The same ablations report that the population or budget penalty is the dominant regularizer: without it, the scorer remains near-uniform and forgetting remains weak, while entropy and $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$4 regularization provide further gains (Yüce et al., 4 Jun 2026).

ATWU is deliberately lightweight. The scorer is a single linear layer over hidden states; the core method uses no auxiliary model, no external token annotations, and no hand-designed token heuristics. At inference time there is no additional machinery beyond the unlearned LLM itself. A converged scorer from a previous ATWU run can be frozen and reused successfully, but a scorer trained only on the original memorized model performs poorly, even if later unfrozen; the scorer must track the unlearning trajectory rather than the initial checkpoint alone (Yüce et al., 4 Jun 2026).

5. Empirical performance and learned token semantics

ATWU is evaluated on TOFU and RWKU. On TOFU, the headline setting is forget10 with Llama-3.1-8B-Instruct; on RWKU, the main setting is the canonical 10-subject batch with Phi-3-Mini-4k-Instruct. The main reported trade-off metrics are Forget Quality ($$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$5), Retain Degradation ($$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$6), and

$$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$7

Utility probes include MMLU, repetitiveness, and win rate against the original checkpoint (Yüce et al., 4 Jun 2026).

On TOFU forget10 with Llama-3.1-8B-Instruct, ATWU reports

$$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$8

The paper compares this with RMU at $$\mathcal{L}(\theta,\mathbf{z}) \;=\; \sum_{x \in \mathcal{D}_R} \sum_{t=1}^{|x|} \ell_{\mathrm{retain}}(x,t;\theta) \;+\; \sum_{x \in \mathcal{D}_F} \sum_{t=1}^{|x|} z_{x,t}\,\ell_{\mathrm{forget}}(x,t;\theta). \tag{1}$$9, JensUn at $\mathcal{I}_F$0 but with severe utility collapse, SatImp at $\mathcal{I}_F$1, WGA at $\mathcal{I}_F$2, and SimNPO at $\mathcal{I}_F$3. On RWKU, ATWU reports

$\mathcal{I}_F$4

exceeding WGA at $\mathcal{I}_F$5, JensUn at $\mathcal{I}_F$6, SatImp at $\mathcal{I}_F$7, and GradDiff at $\mathcal{I}_F$8 (Yüce et al., 4 Jun 2026).

The framework is not tied to a single forget loss. On RWKU, adding the ATWU scorer to existing objectives improves unlearning quality over the corresponding unweighted methods by $\mathcal{I}_F$9 for DPO, $\ell^i_{\mathrm{forget}}(\theta)$0 for NPO, and $\ell^i_{\mathrm{forget}}(\theta)$1 for SimNPO, although the best overall result remains the SatGA$\ell^i_{\mathrm{forget}}(\theta)$2-based ATWU configuration at $\ell^i_{\mathrm{forget}}(\theta)$3 (Yüce et al., 4 Jun 2026).

ATWU also evaluates the quality of its token scores directly. Using ground-truth forget-span labels on TOFU, it ranks tokens by score and computes per-sample AUROC. The reported mean AUROC is $\ell^i_{\mathrm{forget}}(\theta)$4, compared with roughly $\ell^i_{\mathrm{forget}}(\theta)$5–$\ell^i_{\mathrm{forget}}(\theta)$6 for auxiliary-model baselines and around $\ell^i_{\mathrm{forget}}(\theta)$7–$\ell^i_{\mathrm{forget}}(\theta)$8 for the strongest heuristic baselines. Qualitative examples show the scorer emphasizing names, professions, named entities, and factual attribute tokens while downweighting sentence-template material such as “The author’s name is …” (Yüce et al., 4 Jun 2026).

The ablations further separate token scoring from forget-loss design. On TOFU forget10 with Llama-3.2-1B, the paper reports $\ell^i_{\mathrm{forget}}(\theta)$9 for ground-truth labels with GA, $i$0 for ground-truth labels with SatGA, $i$1 for ATWU scores with SatGA, and $i$2 for ATWU scores with SatGA$i$3. The implication is that good token scores alone are insufficient if the forget loss is brittle, and that the score-modulated saturation design is materially important (Yüce et al., 4 Jun 2026).

6. Relation to adjacent methods and principal limitations

ATWU belongs to a broader family of token-aware unlearning methods, but it is distinguished by how it defines token importance. Direct Token Optimization partitions forget sequences into target and non-target tokens using a perturbation-based delta-score and alternates forget ascent with KL-based utility preservation, but its token weighting is a binary top-$i$4 mask rather than learned continuous weights (Lee et al., 30 Sep 2025). TokenUnlearn provides hard selection and soft weighting from masking and entropy signals and recomputes scores online, but it does not formalize an explicit alternating outer loop between score learning and model updates (Wu et al., 1 May 2026). Entropy-guided Token Weighting uses predictive entropy as a proxy for token informativeness in a single-stage weighted optimization, not an alternating one (Koh et al., 20 Apr 2026). CATNIP tokenizes negative preference alignment and scales token forgetting by model confidence, again without an explicit alternating schedule (Yang et al., 2 Feb 2026). In multimodal settings, ViKeR regularizes unlearning with visually estimated token distributions and induces token-level gradient reweighting, but it uses a unified loss rather than alternating optimization (Cai et al., 29 Jan 2026). ALTER uses a two-phase asymmetric LoRA architecture with entropy-guided routing and parameter isolation, yet it is staged rather than explicitly token-weighted and alternating in the ATWU sense (Chen et al., 2 Mar 2026).

Against this background, ATWU’s distinctive claim is that token importance should be identified by compatibility with retain optimality rather than by token probability, entropy, surprisal, external linguistic parsers, or auxiliary-model disagreement alone (Yüce et al., 4 Jun 2026). This makes retain conflict both the conceptual criterion and the supervisory signal for the scorer.

The paper also states clear limitations. Gains are weaker on very small forget sets, where the scorer may not receive enough signal. The budget hyperparameter $i$5 remains user-specified. The theory assumes an approximate structural-versus-forget-specific separation that real language may not satisfy cleanly. There is also a theory–practice gap: the formal recovery results are proved for the multiplicative token-weighted formulation, while the main empirical objective SatGA$i$6 additionally places the score in the saturation exponent. Finally, the evaluations cover two benchmarks and rely on judge-based semantic evaluation, although the paper reports high agreement with humans (Yüce et al., 4 Jun 2026).

A plausible implication is that ATWU defines a general design pattern rather than a terminal formulation. Its scorer–model alternation can be combined with other token signals, other forget losses, or multimodal token sources, but those extensions are not part of the method as published. Within the current literature, its main contribution is to turn token forget-specificity from a heuristic into a learned quantity grounded in the retain–forget trade-off itself (Yüce et al., 4 Jun 2026).