Location-aware Argmax Inversion (LAI)

• Location-aware Argmax Inversion (LAI) is a pseudo-inverse method that computes inverse noise for discrete token generation in visual autoregressive models.
• It employs label and truncated Gumbel sampling to ensure that the perturbed logits recover the original token assignment under a controlled margin (τ).
• LAI underpins VARIN by enabling precise source reconstruction and smooth prompt-guided image editing without additional model training.

Searching arXiv for the cited paper to ground the article in the latest record. Location-aware Argmax Inversion (LAI) is a pseudo-inverse function for argmax sampling introduced within VARIN, a noise inversion-based method for prompt-guided image editing in visual autoregressive models. It is designed for discrete token generation settings in which a model samples codebook indices by adding Gumbel noise to logits and then taking an argmax. Because argmax is many-to-one, the forward sampling step has no true inverse; LAI addresses this by constructing inverse noises that both recover the original token assignment under argmax and remain compatible with Gumbel-distributed sampling. In the formulation presented for next-scale autoregressive image generators such as VAR and HART, LAI is the mechanism that enables precise source reconstruction and controllable editing without additional training (Dao et al., 2 Sep 2025).

1. Problem setting and design objectives

Autoregressive visual models such as VAR or HART predict discrete token maps $r_1,\dots,r_K$ via a sequence of multinomial draws. At scale $t$, the model produces unnormalized log-probabilities

$$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$

draws independent Gumbel noises $g_t \sim \mathrm{Gumbel}(0,1)$, forms perturbed logits

$q_t = p_t + g_t,$

and selects tokens by

$r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$

The inversion problem arises when one wishes to recover a noise-like latent description from a source image so that the autoregressive generator can later be steered by a new prompt while preserving source content. Since the argmax map is many-to-one, exact recovery of the original Gumbel noise is impossible from $r_t$ and $p_t$ alone. The objective is therefore to compute inverse noises $n_t$ such that $p_t + n_t$ has argmax equal to the original $t$0, while also retaining compatibility with genuine Gumbel noise. The paper formulates two requirements for such a pseudo-inverse: first, the extracted noises should resemble draws from $t$1; second, they should retain meaningful location-specific bias derived from the original logits $t$2, so that editing can preserve unedited content (Dao et al., 2 Sep 2025).

2. Formal definition and notation

LAI is described at a single scale $t$3, with the scale index omitted for brevity. The basic objects are:

• $t$4: number of token positions in the map.
• $t$5: size of the discrete codebook.
• $t$6: model logit tensor, where $t$7 is the log-unnormalized probability of class $t$8 at position $t$9.
• $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$0: ground-truth discrete labels extracted by the VAR-VAE encoder from the source image.
• $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$1: one-hot encoding of $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$2, with $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$3.
• $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$4: inverse-Gumbel noise to be computed, satisfying $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$5.

A scalar hyperparameter $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$6 controls the enforced gap between the sampled label logit and all other logits. Larger $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$7 imposes a stronger bias toward the source label. This parameter is central to the trade-off between preservation and randomness in the inversion procedure (Dao et al., 2 Sep 2025).

3. Construction of the pseudo-inverse

LAI proceeds position-wise in two stages. For a fixed position $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$8, it first isolates the logit corresponding to the source label: $$p_t = \log p_\theta(\cdot \mid r_{<t}) \in \mathbb{R}^{\ell \times C},$$9 It then samples a “label Gumbel” by

$g_t \sim \mathrm{Gumbel}(0,1)$0

For each non-label class $g_t \sim \mathrm{Gumbel}(0,1)$1, LAI samples from a truncated Gumbel distribution: $g_t \sim \mathrm{Gumbel}(0,1)$2 The truncation enforces

$g_t \sim \mathrm{Gumbel}(0,1)$3

which guarantees that the argmax remains the source label. The full vector is then assembled by setting the label entry to $g_t \sim \mathrm{Gumbel}(0,1)$4 and the remaining entries to the truncated samples, and the inverse noise is defined as

$g_t \sim \mathrm{Gumbel}(0,1)$5

The operator $g_t \sim \mathrm{Gumbel}(0,1)$6 may be implemented by sampling $g_t \sim \mathrm{Gumbel}(0,1)$7 and returning

$g_t \sim \mathrm{Gumbel}(0,1)$8

By construction, $g_t \sim \mathrm{Gumbel}(0,1)$9, and the resulting inverse noise consists of Gumbel or truncated-Gumbel draws centered around the original logits. This is the sense in which LAI satisfies both the distributional objective and the location-specific bias objective (Dao et al., 2 Sep 2025).

4. Algorithmic form, hyperparameters, and computational profile

In vectorized form, LAI takes as input the ground-truth labels $q_t = p_t + g_t,$0, model logits $q_t = p_t + g_t,$1, and truncation gap $q_t = p_t + g_t,$2. The algorithm forms the one-hot mask $q_t = p_t + g_t,$3, computes $q_t = p_t + g_t,$4, samples $q_t = p_t + g_t,$5, samples truncated Gumbels for all classes using truncation level $q_t = p_t + g_t,$6, overwrites the label entry with $q_t = p_t + g_t,$7, and returns $q_t = p_t + g_t,$8. The inverse noise is then $q_t = p_t + g_t,$9.

Two implementation details are explicitly noted. First, numerical stability in $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$0 may require clamping the argument of $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$1. Second, the choice of $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$2 trades off strict preservation against increased randomness; empirically, $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$3 works well in HART. The computational cost is $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$4, since LAI samples one label Gumbel and $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$5 non-label Gumbels per position, and the entire procedure is fully parallelizable on GPU (Dao et al., 2 Sep 2025).

A toy example in the paper illustrates the mechanism. For $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$6, logits $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$7, source token $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$8, and $r_t[i] = \arg\max_{c \in \{1,\ldots,C\}} q_t[i,c].$9, one may sample $r_t$0, giving $r_t$1. If the truncated draws for the remaining classes are $r_t$2 and $r_t$3, then

$r_t$4

The argmax remains class $r_t$5, while the inverse noise remains a plausible Gumbel-like perturbation.

5. Role within VARIN for text-based image editing

LAI is the inversion component of VARIN, where it is applied during the source-image inversion phase. Given a source image $r_t$6, the VAR-VAE encoder produces token maps $r_t$7. For each scale $r_t$8, the autoregressive model produces logits $r_t$9, and LAI$p_t$0 generates $p_t$1, from which the inverse noise

$p_t$2

is obtained. Collecting $p_t$3 yields a latent “noise fingerprint” that perfectly reconstructs the source when reused with $p_t$4.

During editing, for each scale $p_t$5, fresh random Gumbel noise $p_t$6 is interpolated with the inverse noise: $p_t$7 where $p_t$8 denotes the logits under the edited prompt and $p_t$9 is a scale-dependent mixing coefficient. The reported schedule is usually $n_t$0 at the start scale $n_t$1, then linearly decreasing to $n_t$2 by scale $n_t$3. Taking the argmax produces $n_t$4, and the VAR-VAE decoder reconstructs the edited image. In this formulation, LAI-extracted inverse noises keep sampling close to the source in unedited regions while permitting controlled deviations where the edited prompt exerts pressure (Dao et al., 2 Sep 2025).

6. Relation to alternative inversions, reported behavior, and limitations

LAI is contrasted directly with One-Hot Argmax Inversion (OAI), which sets

$n_t$5

OAI guarantees that the source label is the argmax and recovers perfect reconstruction, but the implied noise $n_t$6 is not Gumbel distributed; the paper characterizes it as highly biased and degenerate, violating the requirement that inverse noise remain compatible with Gumbel sampling. LAI differs in two respects: it centers its sampling on the original logits $n_t$7, and it enforces a controllable margin $n_t$8 rather than a hard one-hot collapse.

The reported empirical behavior is that LAI yields inverse noises that both perfectly reconstruct the source, with zero structure distance at $n_t$9, and allow smooth editing when $p_t + n_t$0, thereby preserving background. In quantitative comparisons described for Table 1, VARIN + LAI outperforms OAI-based editing and discrete diffusion inversion baselines (DICE), while matching or exceeding continuous diffusion inversion methods in speed and background preservation. The paper also notes several limitations and open directions: LAI depends on choosing $p_t + n_t$1 per model; extreme $p_t + n_t$2 values can under-constrain or over-constrain inversion; truncated Gumbel sampling is more expensive than a single argmax; and a possible extension is to generalize the construction to top-$p_t + n_t$3 or other structured sampling schemes via the Gumbel-Top-$p_t + n_t$4 trick (Dao et al., 2 Sep 2025).

A common misconception is to regard LAI as an exact inverse of discrete sampling. The formulation does not support that interpretation: argmax remains non-invertible, and LAI is explicitly a pseudo-inverse that constructs surrogate perturbed logits $p_t + n_t$5 consistent with the observed label assignment. Its significance lies not in recovering the original noise realization, but in producing an invertible noise proxy that remains aligned with the model’s pre-softmax geometry and is therefore useful for training-free, prompt-guided editing in next-scale autoregressive image generators.