InterpTok

Overview

Latent diffusion models generate images as codes of an image tokenizer, whose decoder turns them into the pixel image. Recent tokenizers improve how well these codes reconstruct an image and how easily a generator learns them. Yet a generator does not reproduce the codes of training images, and when a generated code lands outside what the decoder expects, the decoded image contains imperfections. Current methods make the decoder more robust to generated codes by adding noise during training. However, generated codes lie between the training codes and resemble their mixture more than one of them plus noise. To improve generated images, decoders should therefore be trained on codes in the space between training samples. We propose InterpTok, which fine-tunes MAETok to decode mixtures of image codes into natural images. Since no target image exists between two training samples, a discriminator judges whether each decoded mixture looks as natural as its two sources. A mixing ratio loss requires the decoded mixture to keep the mixing ratio of the two codes. Together, they prevent collapse to a trivial solution and yield a sharp and natural composition of both images. InterpTok keeps MAETok's reconstruction quality and improves interpolation quality more than threefold, the best of all tokenizers in a recent benchmark. Beyond interpolation, InterpTok improves the generation FID of a matched diffusion transformer from 11.61 to 8.48 without guidance, and even reaches 8.11 at an eighth of the sampling steps, where the generator on MAETok's codes worsens. The generator trained on our tokenizer also reaches the final result of the one trained on MAETok at half the training steps. Our findings show that the space between image codes can and should be explicitly supervised rather than ignored.

Between Two Images

Drag the slider to mix the codes of two ImageNet validation images, (1 − α) a + α b, and decode the mixture. MAETok superimposes the two photographs. InterpTok decodes one natural image that changes from the first photograph to the second. Most of the pairs are those of the paper's figures.

MAETok

InterpTok (ours)

Play

Method

InterpTok trains the decoder on the codes between two images. The encoder maps two similar training images A and B to their codes a and b, and the decoder maps the mixture of the codes to the decoded mixture. No image exists between two training images to compare the decoded mixture with. Instead, an interpolation discriminator judges each decoded mixture beside its two source images. A mixing ratio loss requires the decoded mixture to keep the ratio of the mixture, a fraction α of the way from A to B. This prevents a collapse onto one of the two sources and yields a sharp and natural composition of both images. Reconstruction keeps its losses, so the codes of training images still decode to the images.

Results

On ImageNet at 256 × 256, InterpTok lowers MAETok's interpolation FID from 14.42 to 4.30, lower than that of every tokenizer in the comparison released by Xu et al., and keeps its reconstruction FID, 0.597 against 0.582.

A Smaller Code

An earlier experiment trained the codes between images with a code four times smaller, 32 tokens of 32 values each, on a frozen REPA-E autoencoder with a transformer encoder and a flow decoder in between, trained on ImageNet and the image-editing pairs of ImgEdit. The smaller code is more semantic: its mixtures change the content of the scene rather than superimpose two images. Reconstruction is the cost, with a PSNR of 19.31 against InterpTok's 24.02. The endpoint images were generated with Z-Image-Turbo.

Two images

rabbit to applesrabbit to doberman

Play

Four images

Code and Models

The code trains InterpTok from MAETok's released weights and evaluates reconstruction, interpolation and generation. The weights of InterpTok, of its decoder for MAETok's generators, and of the paper's two SiT-B generators are on Hugging Face.

Acknowledgements

Funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's Excellence Strategy – EXC number 2064/1 – Project number 390727645. This work was supported by the German Research Foundation (DFG): SFB 1233, Robust Vision: Inference Principles and Neural Mechanisms, TP 02, project number: 276693517. This work was supported by the Tübingen AI Center. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Jan-Niklas Dihlmann.

Citation

More Information