ForgeryTTT: Zero-Shot Image Manipulation Localization with Test-Time Training

Early methods in image manipulation localization primarily rely on statistical analysis and handcrafted features. They mainly focus on identifying jpeg compression [5], noise inconsistencies [4], and color filter array [6], etc. While effective in detecting global inconsistencies, these techniques often struggle with complex manipulations and precise localization. Recent advancements in deep learning have significantly influenced the field. Deep neural networks now automatically learn discriminative forensic features from large-scale datasets, resulting in notable improvements [7, 20, 21, 22, 23, 24, 25]. Most of these methods employ an encoder-decoder architecture, where an image encoder extracts the features and a decoder predicts the forgery mask. To capture imperceptible clues beyond RGB images, some approaches transform the RGB image into the noise domain or frequency domain, providing complementary views for more effective forgery detection. ManTraNet [7] integrates the SRM filter [26] and Bayar filter [27] into the pre-trained backbone to extract features from both the noise map and the RGB image. Objectformer [28] extracts RGB features and high-frequency features via two encoders and combines them as multimodal features. Edge artifacts have also been utilized as crucial supervision signals to better capture details. MVSSNet [29] employs the Sobel filter to create a multi-level edge-supervised branch. TANet [30] refines the boundary by bilaterally exploring foreground and background cues. However, the proliferation of fake images and the development of forgery techniques present ongoing challenges, while no one has yet focused on the robustness when confronted with unseen scenes.

Traditional image manipulation classification methods initially extract hand-crafted features and then analyze statistical characteristics to differentiate the manipulated and authentic images. Fridrich et al. [31] proposed a copy-move forgery detection method based on approximate block matching. Pan et al. [32] identified forgery images using a SIFT-based matching method. Shi et al. [33] employed multi-size block discrete cosine transform and Markov transition probabilities to detect splicing. In the deep learning era, image manipulation classification methods utilize deep neural networks to capture manipulation traces [34, 35, 36]. Zhang et al. [37] input image patch features into a five-layer neural network to determine whether the image has been manipulated. Bayar et al. [34] proposed constrained convolutional neural networks that suppress the influence of image content on manipulation traces and adaptively extract manipulation features. Recent works [7, 38] on image manipulation localization simultaneously address both image manipulation classification and localization. Different from those supervised methods, we design a mask-based self-supervised image manipulation classifier.

The performance of deep neural networks may significantly drop in unseen test data that diverges from the training distribution. Test-time training (TTT) aims at adapting the pre-trained model to out-of-distribution test data through self-supervised/ unsupervised object functions. Since the statistics in the BatchNorm layers represent domain-specific knowledge, the statistic-based methods recalculate the statistics to adapt the model to the test domain [13, 14, 39]. Schneider et al. [13] adjust the BatchNorm statistics of the model by estimating from the test sample. Wang et al. [14] optimizes the affine parameters in the BatchNorm layers by minimizing the prediction entropy. Auxiliary head-based algorithms leverage the self-supervised learning task to adjust the features from the training domain to the test instance [15, 17, 16]. They typically consist of a shared image encoder, a primary decoder for the main task, and an auxiliary decoder for the self-supervised task. Sun et al. [15] develop an auxiliary head for rotation prediction and fine-tune the backbone using the auxiliary head during testing. Gandelsman et al. [17] employ the mask autoencoder which reconstructs the test sample to update the pre-trained model. Some works have explored effective TTT approaches in other areas [40, 41, 42]. It has not yet been explored in image manipulation localization and we take the first step.

Image Manipulation Localization. Image manipulation localization involves localizing the pixel-level forgery mask in an RGB image. This is typically accomplished using an encoder-decoder network architecture [7, 38, 25]. In the image encoder, the image is first split into patch tokens, which are then processed through hierarchical transformer blocks to produce multi-scale representations. These representations are subsequently fed into the localization head, which is responsible for predicting the forgery mask. The ground truth object mask is used to supervise the model via the binary cross-entropy loss.

Test-Time Training. Test-time training (TTT) is introduced to enhance the generalization capability of models to out-of-distribution test data [15, 14, 13]. A commonly employed TTT framework based on the auxiliary head [15] consists of a shared encoder $\mathcal{E}$, a main head $\mathcal{D}$ for the main task, and an auxiliary head $\mathcal{C}$ for self-supervised learning. The TTT framework typically involves a two-stage training process. In the first stage, the network is trained using both the main loss $\mathcal{L}_{\text{main}}$ and a self-supervised loss $\mathcal{L}_{\text{ssl}}$:

where $\lambda$ is the hyper-parameter to balance the two losses.

In the subsequent stage, which is known as test-time training (TTT), the encoder $\mathcal{E}$ is fine-tuned for each test sample based on the self-supervised objective function:

This process allows the model to adapt to unseen data through the self-supervised learning task during testing, thereby improving its ability to generalize beyond the distribution of the training data.

Image manipulation classification has been integrated as an auxiliary task in several image manipulation localization algorithms [38, 43]. The classification head is trained in a supervised manner using labeled forgery and authentic images. The classification score helps detect global inconsistencies since forgery images exhibit a high response in at least one region, while authentic images show no response. However, test-time training (TTT) necessitates a self-supervised auxiliary task to allow for model optimization without relying on labels during testing. To address this requirement, we develop a mask-based self-supervised image manipulation classification algorithm (Figure 4).

Given an input image $i\in\mathbb{R}^{h\times w\times 3}$, we extract multi-scale features $F=\left\{f_{s}:s\in(1,2,...,s_{max})\right\}$ via the image encoder $\mathcal{E}$, where $f_{s}\in\mathbb{R}^{h_{s}\times w_{s}\times c_{s}}$ represent features extracted at scale $s$, and $h_{s}$, $w_{s}$, and $c_{s}$ are the height, width, and dimension of the feature at scale $s$, respectively. The proposed classification head is shown in Figure 5. Specifically, we first merge the multi-scale features obtained by the image encoder into tokens. The multi-scale features are then unified in the channel dimension by the corresponding linear layer:

These features are upsampled to $h_{1}\times w_{1}$ and concatenated together as $f$. The fused features $f$ are split into tokens with $p\times p$ patch size (in contrast to the original resolution) via a patch embedding layer:

where $e\in\mathbb{R}^{\frac{h}{p}\times\frac{w}{p}\times d}$. Next, we select some tokens as query $q$. The query and the learnable class token $c$ are fed into a series of transformer blocks:

where $L_{j}$ is the j-th transformer block. Each transformer block is composed of multi-head self-attention and feed-forward networks, together with layer normalization and residual connection. Finally, we map the class token into the class probability via a linear layer:

where $y$ represents the probability of whether the given query is manipulated.

On top of the basic architecture, we further introduce several novel features, including label-free query construction and mask-based token dropout, which collectively make a powerful model.

Lable-Free Query Construction. To adapt image manipulation classification to test-time training, we model it as a self-supervised task rather than a supervised one. The pseudo label is derived from the given mask. During training-time training, we employ the ground truth mask to group the foreground tokens and background tokens. The mask is downsampled using max-pooling to match the token size. The foreground tokens and background tokens are used to construct the ”manipulated” query, while the background tokens are used to construct the ”authentic” query. These two types of tokens are concatenated with a class token to train the classifier in a self-supervised manner. During TTT, we construct only the manipulated query using the predicted mask to further train the image encoder with the pseudo-training samples. The manipulated query is reliable since the authentic query could be wrong when the predicted mask is inaccurate. Thus, we can use the constructed query to train a self-supervised classifier and fine-tune the pre-trained model with pseudo-training samples during testing.

Mask-based Token Dropout. Previous methods that incorporate image manipulation classification as an auxiliary task input all tokens for classification. We believe it is redundant. A small number of authentic and manipulated tokens is sufficient to train a classifier. Inspired by recent advances in masked image modeling, we boost the performance via token dropout. To achieve it, we first utilize a mask to differentiate between foreground-manipulated tokens and background-authentic tokens within a given set of image tokens. We then perform random dropout to select a subset of both foreground and background tokens, discarding the remaining tokens. This dropout process follows a uniform distribution, ensuring an equal representation of different types of tokens. The dropout ratio r is the same in both foreground and background. The benefit of this mask-based dropout strategy is straightforward. It boosts classification performance by eliminating redundant information since image manipulation classification does not rely on most of the tokens.

With the development of image editing technology, massive forgery images and advanced editing methods are challenging existing image manipulation localization algorithms. Mainstream methods often directly apply the models trained on large-scale datasets to the target test images, which makes them struggle to generalize to unseen images. To overcome these limitations, we propose ForgeryTTT. ForgeryTTT first trains on a large-scale dataset and then adapts the pre-trained model to test samples through self-supervision during testing.

As illustrated in Figure 3, ForgeryTTT is a multi-task framework composed of a shared image encoder $\mathcal{E}$, an image manipulation localization head $\mathcal{D}$, and an image manipulation classification head $\mathcal{C}$. The image encoder $\mathcal{E}$ is a hierarchical transformer [19] that produces a hierarchical representation for the input image $i$. The localization head $\mathcal{D}$ is a lightweight multi-layer perception [44] that predicts the forgery mask. The classification head $\mathcal{C}$ is composed of several transformer blocks and linear layers, which distinguish whether the given query is manipulated.

We utilize multiple datasets, including SynCOCO, CASIA [45], Coverage [46], Columbia [47], NIST16 [48], and CocoGlide [24], in our experiments.

Training Data. We create a synthetic dataset named SynCOCO to train our model. This dataset is built upon MSCOCO [54] and contains forgery images of several manipulation types, including splicing, copy-move, and removal. For splicing, we extract an annotated instance from a random image and paste it into another image. For copy-move, we extract an annotated instance from a random image and paste it into a different region within the same image. For removal, we delete an annotated instance from an image and employ a cutting-edge image inpainting algorithm [2] to restore the region. In total, the synthetic dataset comprises 150,000 labeled forgery images.

Testing Data. We evaluate our method on five publicly available benchmarks. CASIA [45] is composed of CASIA v1 (5,123 images) and CASIA v2 (921 images), including splicing and copy-move samples. We evaluate our method on CASIA v2. Coverage [46] is a copy-move forgery dataset containing 100 manipulated images. We use the test split (25 images) to evaluate the proposed method. Columbia [47] includes 180 forgery images designated for splicing detection. We evaluate our method on the entire dataset. NIST16 [48] contains 564 forgery images involving splicing, copy-move, and removal types. We use the test split (160 images) to evaluate the proposed method. CocoGlide [24] comprises 512 images generated from the COCO 2017 validation set using the GLIDE diffusion model [55]. We evaluate our method on the entire dataset.

We implement our algorithm using Pytorch framework and all the experiments are performed on a single NVIDIA A40 GPU. As for optimization, we use the Adam [56] optimizer. The initial learning rate is set to $2e^{-4}$, which is adjusted by the exponential decay strategy with a decay rate of 0.9. We employ a fixed learning rate of $2e^{-5}$ during TTT. The models are trained for 20 epochs on SynCOCO dataset and 10 epochs during TTT on each test sample. The mini-batch size is equal to 32. The input image is resized into a fixed resolution of $512\times 512$. Random horizontal flipping, resizing, and cropping are used for data augmentation.

As for the network, we adopt Swin-Tiny [19] as the image encoder, the lightweight ALL-MLP decoder [44] as the localization head, and the number of transformer blocks in the classification head is set to 5. The loss weight $\lambda$ in training-time training (equation 7) is set to 0.01. The dropout ratio for token dropout is 0.5 and the patch size in the classification head is 16 $\times$ 16 (section III-B).

We evaluate the proposed method in image manipulation localization using several metrics following [24, 43]. We report the F1 score using both the best threshold (F_best) and the default 0.5 threshold (F_fix) to measure pixel-level performance. For image-level analysis, we utilize the Area Under the Curve (AUC) and balanced accuracy metrics (ACC). All these metrics are the larger the better.

We compare the performance of our method with other state-of-the-art image manipulation methods [7, 9, 49, 50, 29, 51, 52, 53, 24]. Since our method is designed for zero-shot image manipulation localization, the model is only trained on a large-scale synthetic dataset and evaluated directly on the test split of the target datasets. Instead, some methods further fine-tune their models on the training split of the target dataset [9, 38] or include the training split into their training data [24, 43]. For a fair comparison, we divided them into two categories by different training protocols: (1) None zero-shot methods: SPAN [9], MVSSNet [29], CATNetv2 [52], PSCCNet [53], TruFor [24], and UnionFormer [43]. (2) Zero-shot methods: ManTraNet [7], EXIF-SC [49], Noiseprint [50], IFOSN [51], and our ForgeryTTT.

Table I describes the comparison of F_best, F_fix, AUC, and ACC between our method and other state-of-the-art methods. Specifically, our method is much better than previous zero-shot methods (ManTraNet [7], EXIF-SC [49], Noiseprint [50], and IFOSN [51]). Our approach outperforms them on all five benchmarks and achieves an F_fix improvement of 20.1$\%$ on average at least. On the other hand, our approach also achieves the top performance when compared to the methods that include the training split of the target dataset into the training dataset (CATNetv2 [52], TruFor [24], and UnionFormer [43]) or fine-tuning the pre-trained model on the training split (SPAN [9], MVSSNet [29], and PSCCNet [53]). Our method gets the highest average  F_fix , even though our model has not seen these images before. It is worth noting that our approach performs best on the CocoGlide dataset, which is built with a prompt-based generative artificial intelligence technique [55]. None of these methods have included this kind of data during training. This shows that our approach has strong potential to be generalized to unseen scenarios. We show the visual comparison of our method and other state-of-the-art methods in Figure 8. The proposed method performs better than others on five benchmarks, which proves the superiority of the proposed method.

We further analyze the efficiency of the proposed method. The parameters of each component in our model are 27.5M for the image encoder, 0.6M for the localization head, and 5.1M for the classification head. Consequently, we introduce only a small number of additional parameters to the baseline model. As for comparison, our model (33.2M) is smaller than other top competitors, such as TruFor (68.7M), CATNetv2 (114.3M), IFOSN (128.8M), and MVSSNet (146.9M). Our model takes about 30 hours on one A40 GPU for training, while TruFor takes more than 14 days for training on one A6000 GPU. Our full model takes 12 milliseconds per frame for inference and 260 milliseconds per frame for TTT. As an inherent drawback of TTT, it does bring additional computational time. Therefore we also try to mitigate its impact. Equipped with the proposed strategy, the proposed method becomes more effective and faster than the basic strategy (please refer to the next section).

We also compare the proposed method with state-of-the-art TTT methods [13, 14, 15, 17]. For statistics-based TTT methods [13, 14], we directly apply TTT to the baseline model, and we follow a two-stage training pipeline for the TTT methods with an auxiliary head [15, 17]. All these methods are pre-trained on SynCOCO and evaluated on the test split of five benchmarks. As shown in Table II, the proposed method outperforms all other TTT methods significantly. On the one hand, although other TTT methods work well in some corruptions such as fog, snow, and rain, they may be less effective against natural domain shifts. On the other hand, designing an auxiliary task that is closely related to the main task has a significant impact on the performance of TTT. The proposed ForgeryTTT is proved to be advantageous to image manipulation localization, which brings significant improvements when applied to unseen forgery images.

We evaluate the robustness under several distortion settings. Table VIII shows the results on CAISA datasets of MVSSNet, IFOSN, TruFor, and our method, on which Gaussian Blur, Gaussian Noise, JPEG compression, and online social network transmission process the images. Specifically, the proposed method shows more stable results than others in most cases, which gains the least performance drop under various attacks. Our method performs poorly when faced with highly compressed and Gaussian blurred images, and we notice that the performance of TruFor also drops significantly in these cases. Instead, MVSSNet and IFOSN perform better on Gaussian blurred images. We think this is because of the different kinds of backbones. Our method and TruFor use a transformer-based backbone, while MVSSNet and IFOSN use a CNN-based backbone. When image quality is severely degraded, the non-semantic discrepancies between manipulated and authentic regions become smaller, while local inconsistencies in manipulated boundaries are more important. In addition, we believe that including these types of data augmentation in training could help improve the robustness of these attacks.