Pivotal Tuning for Latent-based Editing of Real Images

In plain words

StyleGAN can change a face's age, pose or smile, but only for faces the generator could already draw. To edit a real photo, you first need a code that makes the generator draw that photo. The usual codes either match the photo closely and are hard to edit, or are easy to edit and do not match the photo.

This paper keeps the easy-to-edit code and changes the generator a little instead, so that it draws this one photo from that code. A short extra training run does this, and a penalty stops nearby faces from changing too. The result is a close copy of the photo that still responds well to edits, at a cost of a few minutes per face.

TL;DR

• PTI inverts a real face into StyleGAN2's native W space, then fine-tunes the whole generator for up to 350 steps so the face is reproduced from that fixed code §3.1 §3.2 §4.5.
• A locality term keeps the tuned generator's outputs for random codes near the pivot close to the original's; the code ships it switched off by default §3.3 hyperparameters.py.
• Author-reported on CelebA-HQ: LPIPS 0.09 against 0.34 for the best baseline (SG2 W+), and identity similarity 0.9 against 0.85 tab. 1.
• Demonstrated only with StyleGAN2 on faces; other generators are named as future work §5.

Problem

StyleGAN's native space W gives the most editable codes, but inverting a real image into W gives large distortion, such as identity loss or unnatural appearance §1. Extended spaces such as W+ reduce distortion but make edits weaker, so there is a trade-off between distortion and editability §1 §2. Out-of-domain faces (heavy makeup, headwear, elaborate hairstyles) are the hardest case §1 §4.

Method

The method has two steps, plus a regularizer.

Inversion (§3.1). Find a pivot code w_p in W with the generator frozen. Both w and the noise maps n are optimized against the input image x using an LPIPS perceptual loss plus a noise regularizer eq. 1. The paper uses the W space, not W+, because distortion is later removed by tuning §3.1. The paper says it does not use the mapping network from Z to W in this step §3.1. It uses the off-the-shelf inversion of Karras et al. §3.1.

w_p, n = \underset{w, n}{\arg\min}\; \mathcal{L}_{\text{LPIPS}}(x, G(w, n; \theta)) + \lambda_{n}\,\mathcal{L}_{n}(n) \qquad [\text{eq. 1}]

Pivotal tuning (§3.2). Unfreeze the generator and tune it, starting from the pretrained weights θ, so that the fixed pivot w_p reproduces the input image. The loss is LPIPS plus a weighted L2 term eq. 2. The paper extends this to N images by averaging the per-image losses eq. 3. The pivot is held fixed. The paper states that random or mean codes as the pivot fail to converge, so the inversion step matters §3.2 §4.4.

\mathcal{L}_{pt} = \mathcal{L}_{\text{LPIPS}}(x, x^{p}) + \lambda_{L2}\,\mathcal{L}_{L2}(x, x^{p}), \quad x^{p} = G(w_p; \theta^{*}) \qquad [\text{eq. 2}]

\mathcal{L}_{pt} = \frac{1}{N}\sum_{i=1}^{N}\Big(\mathcal{L}_{\text{LPIPS}}(x_i, x_i^{p}) + \lambda_{L2}\,\mathcal{L}_{L2}(x_i, x_i^{p})\Big), \quad x_i^{p} = G(w_i; \theta^{*}) \qquad [\text{eq. 3}]

Locality regularization (§3.3). Each iteration samples z from a normal distribution, maps it to w_z with the mapping network, and places w_r at distance α from the pivot in the direction of w_z eq. 4. The tuned generator must then match the original generator at w_r. The regularizer is an LPIPS plus a weighted L2 term between the original and tuned outputs at w_r eq. 5, averaged over N_r samples eq. 6. The final objective weights it by λ_R eq. 7.

w_r = w_p + \alpha\,\frac{w_z - w_p}{\lVert w_z - w_p \rVert_2} \qquad [\text{eq. 4}]

\mathcal{L}_{R} = \mathcal{L}_{\text{LPIPS}}(x_r, x_r^{*}) + \lambda^{R}_{L2}\,\mathcal{L}_{L2}(x_r, x_r^{*}), \quad x_r = G(w_r; \theta),\; x_r^{*} = G(w_r; \theta^{*}) \qquad [\text{eq. 5}]

\mathcal{L}_{R} = \frac{1}{N_r}\sum_{i=1}^{N_r}\Big(\mathcal{L}_{\text{LPIPS}}(x_{r,i}, x_{r,i}^{*}) + \lambda^{R}_{L2}\,\mathcal{L}_{L2}(x_{r,i}, x_{r,i}^{*})\Big) \qquad [\text{eq. 6}]

\theta^{*} = \underset{\theta^{*}}{\arg\min}\; \mathcal{L}_{pt} + \lambda_{R}\,\mathcal{L}_{R} \qquad [\text{eq. 7}]

The paper does not give a single algorithm box. The following summarizes the steps above, with the settings from §4.5:

# Pivotal Tuning Inversion (PTI). No algorithm box in the paper; this
# summarizes [§3] with the settings in [§4.5].
# Input:  aligned face image (3, H, W); pretrained weights theta
# Output: pivot code w_pivot (512,); tuned weights theta_star

def pivotal_tuning_inversion(image, theta):
    # Invert with the generator frozen; optimize w and the noise maps,
    # then keep only w as the pivot. [§3.1] [eq. 1]
    w_pivot = invert_to_w(image, theta, steps=450)

    # Start tuning from the pretrained weights. The pivot stays fixed. [§3.2]
    theta_star = copy(theta)
    for step in range(350):  # [§4.5]
        # Reconstruct the input from the pivot. [§3.2] [eq. 2]
        recon = generator(w_pivot, theta_star)
        pivot_loss = lpips(image, recon) + 1.0 * l2(image, recon)

        # Sample z and map it to W. [§3.3]
        w_z = mapping(normal_noise(512))
        # Move from the pivot toward w_z by distance alpha = 30. [§3.3] [eq. 4]
        direction = w_z - w_pivot
        w_r = w_pivot + 30.0 * direction / norm(direction)

        # The original generator is a fixed target. [§3.3] [eq. 5]
        target = generator(w_r, theta)
        tuned_at_r = generator(w_r, theta_star)
        locality_loss = lpips(target, tuned_at_r)
        locality_loss = locality_loss + 1.0 * l2(target, tuned_at_r)

        # Weight the locality term by lambda_R = 0.1. [eq. 7] [§4.5]
        loss = pivot_loss + 0.1 * locality_loss
        theta_star = adam_step(theta_star, loss, learning_rate=3e-4)
    return w_pivot, theta_star

Notes on the pseudocode: the paper defines the regularizer as an average over N_r samples and sets N_r = 1 §4.5, so the sum reduces to one sample.

Hyperparameters

Values from the paper are cited to the paper. Code values are shown only where they are cited separately under Paper vs code.

Results

Author-reported. Setup: StyleGAN2 pretrained on FFHQ; the first 1000 CelebA-HQ test samples are inverted §4.5; four methods: PTI ("Ours"), e4e, SG2 (W), and SG2 W+ §4 §4.1.

Reconstruction tab. 1. Arrows as printed in the paper; see Pitfalls for the MS-SSIM row.

Source: tab. 1. The caption says "all metrics indicate superior reconstruction" tab. 1; see Pitfalls for the MS-SSIM arrow.

Editing tab. 2. Edit magnitude is measured as the yaw change under a pose edit, using Microsoft Face API tab. 2. The paper's text and caption do not state its units. ID similarity uses a pretrained face recognition network §4.1.

Source: tab. 2. Single edits are smile, pose and age; sequential means all three applied one after another tab. 2 §4.2.

Qualitative. Figures show reconstructions of out-of-domain faces fig. 3, CelebA-HQ reconstructions fig. 4, editing comparisons fig. 5, StyleClip edits with and without PTI fig. 6 fig. 7, famous figures fig. 8, out-of-domain edits fig. 9, multi-identity inversion fig. 10, and sequential edits fig. 11 fig. 12. The paper says StyleClip uses e4e pivots when combined with PTI, because StyleClip was trained on e4e codes §4.2.

Ablation §4.4 fig. 15. Qualitative only, with no table. Using the W+ pivot (B), the mean code in place of the pivot (C), a random code in place of the pivot (D), or optimizing the pivot from the mean (E) or a random code (F) gives weaker edits or more distortion. The one number given: tuning with the pivot left free gives an output nearly identical to the fixed pivot, LPIPS 0.015 ± 5e-6 and MSE 0.0012 ± 1e-6 §4.4.

Locality regularization §4.3 fig. 13 fig. 14. Tuned on 12 identities. The paper reports that the regularizer almost removes artifacts on random codes and keeps target reconstruction similar fig. 13. The plotted curves fig. 14 give no numbers in the text, so values are not stated.

α sweep app. Locality Regularization. α = 2 gives little regularization; α = 60 behaves like a random code; α = 120 gives unrealistic images. The paper picks α = 30 app. Locality Regularization.

Paper vs code

Official code: github.com/danielroich/PTI at commit da94d59d, owned by author Daniel Roich, and linked in the paper's abstract. Code paths are relative to the repository root.

• First inversion learning rate. The paper uses 5e-3 §4.5. The config sets first_inv_lr = 5e-3 hyperparameters.py, but calc_inversions does not pass it on base_coach.py. The projector's initial_learning_rate defaults to 0.01 w_projector.py, multiplied by a ramp w_projector.py and written into the optimizer at every step w_projector.py. The effective peak rate is therefore 0.01, not 5e-3.
• First inversion loss. The paper uses LPIPS eq. 1 §3.1. The projector uses a squared distance between VGG16 features w_projector.py w_projector.py, not the AlexNet LPIPS that the tuning uses hyperparameters.py.
• Mapping network in inversion. The paper says the mapping network is not used in inversion §3.1. The projector starts from the mean of 600 mapped samples w_projector.py base_coach.py. Only the starting point comes from the mapping network, which the paper does not describe.
• Locality regularizer off by default. The paper's regularized results need it §4.3. The config sets use_locality_regularization = False hyperparameters.py. Reproducing the paper requires turning it on.
• Locality weights. The paper sets λ^R_L2 = 1 and λ_R = 0.1 §4.5. The code uses 0.1 for both the L2 and LPIPS parts of the regularizer hyperparameters.py and has no λ_R. The effective weights match: 1 × 0.1 = 0.1 for L2 and 0.1 for LPIPS.
• Locality sampling. The paper samples z from a normal distribution and maps it without truncation §3.3. The code calls the mapping with truncation_psi=0.5 localitly_regulizer.py.
• Locality distance. The code computes get_morphed_w_code as fixed_w + alpha * direction / norm localitly_regulizer.py, which matches eq. 4. Line 20 computes an expression that is never used localitly_regulizer.py. It has no effect.
• Multi-identity update. The paper averages the loss over N images into one update eq. 3. The code takes a separate Adam step for each image in every pass multi_id_coach.py.
• Early stop. The single-identity code stops tuning when the LPIPS term is at most 0.06, before the update of that step single_id_coach.py hyperparameters.py. The paper does not mention this; tuning may run fewer than 350 steps.
• Multi-identity has no early stop multi_id_coach.py.
• Pivot shape. The code repeats the 512-d pivot across 18 layers before synthesis w_projector.py. The paper describes a single W code.
• Noise mode. Reconstruction uses noise_mode='const' base_coach.py. The locality regularizer uses noise_mode='none' localitly_regulizer.py localitly_regulizer.py. The paper does not state a noise mode.
• Face alignment. The code aligns faces with dlib before inversion align_data.py README.md. The paper does not describe any alignment step.
• Inversion weight. The code uses weight 1e5 on the noise regularizer w_projector.py. The paper gives no value for λ_n eq. 1.
• Generator weights. The code loads the G_ema weights from a StyleGAN2-ADA FFHQ pickle models_utils.py paths_config.py. The paper says StyleGAN2 pretrained on FFHQ §4.

Pitfalls

• Inversion learning rate. The paper's 5e-3 and the code's effective 0.01 differ (see above). Reproducing the paper's inversion needs a deliberate choice.
• Locality is off by default. Enable use_locality_regularization hyperparameters.py; without it the code gives the unregularized method.
• Scale of α. Eq. 4 normalizes by the L2 norm of the difference. The code's codes are the 512-d pivot broadcast to 18 layers w_projector.py. If the mapping output is also broadcast, the norm covers all 18 copies, and α = 30 is about 30/√18 ≈ 7.1 in a single 512-d copy. We did not check the mapping network's output shape, because StyleGAN2-ADA's network code is not in the repository. Verify this before reusing α.
• Per-identity cost. Each identity needs its own full copy of the generator, which the code saves to disk after each image single_id_coach.py. Storage per identity is the size of the generator, not the code.
• Time. About 3 minutes per identity on an RTX 2080 §4.5 §5. Multi-identity tuning time grows linearly with the number of identities §4.5.
• Multi-identity sharing. One tuned generator serves several identities. The regularizer ablation uses 12 identities §4.3, and the largest multi-identity figure shows 6 at once fig. 10. Scaling beyond this is not tested.
• Pivot. The pivot must come from inversion; random and mean codes fail §4.4. Optimizing the pivot does not help, which the paper reports as negligible change §4.4.
• MS-SSIM arrow. The table marks MS-SSIM as "↓" and says ours is best at 0.21 tab. 1. Standard MS-SSIM rises for more similar images, so the arrow or the values may be inverted. The paper does not define the quantity; treat this row as unclear.
• Units. The paper does not state the units of the edit-magnitude row tab. 2.
• Edits beyond the regularized region. Edits run on the tuned copy. The locality term only constrains random codes near the pivot eq. 4 eq. 5. The paper's sequential-edit results tab. 2 are its only test of stacked edits.
• Deterministic seed. The projector's mean samples use a fixed seed of 123 w_projector.py. The paper does not state a seed.

Lineage

• Builds on: the StyleGAN2 generator and its direct-optimization inversion, both from Karras et al. (CVPR 2020) §3.1 §4.
• Builds on: the W+ inversion of Abdal et al. (2019), which PTI argues against for editability §1 §4.1.
• Compared against: e4e, the encoder of Tov et al., arXiv:2102.02766 §4.1 §4.2 tab. 1 tab. 2.
• Editing methods used: GANSpace §4.2, InterfaceGAN §4.2, and StyleCLIP, which runs on e4e pivots §4.2.
• Related generator tuning: Bau et al. (tuning the generator for interactive editing) and Pan et al. (optimizing noise and the generator for BigGAN) §2.
• Future work named by the paper: a trainable mapper that approximates PTI in one forward pass, using several photos per identity, and other generators such as BigGAN §5.
• Superseded: not stated.