Diff-Feat: Diffusion-Based Cross-Modal Feature Extraction for Multi-Label Classification

This is the official implementation for our paper "[Diff-Feat: Diffusion-Based Cross-Modal Feature Extraction for Multi-Label Classification]".

Image-only representation extraction

This module extracts intermediate representations from pretrained DiT-XL/2(256$\times$256) models and trains a lightweight classifier on datasets like MSCOCO, VG500, or Tiny-ImageNet, using diffusion features across different timesteps and Transformer blocks.

---

Overview

This pipeline includes:

1. Loading raw images
2. Extracting transformer block features at specific diffusion timesteps
3. Saving extracted features as .npz files for reproducibility and fusion
4. Training a linear classifier on extracted features
5. Evaluating Top-1 / Top-5 accuracy (for single-label) or mAP, F1, Precision, Recall (for multi-label)

---

Requirements

The environment setup is identical to that required for training and evaluating the DiT model. Please refer to the DiT repository for installation and dependency instructions.

---

Usage

1. Prepare Latent Vectors

If starting from raw images:

# Use provided script to encode raw images into VAE latents(vae_latent.py)
torchrun --nproc_per_node=4 vae_latent.py
# Output: {dataset_name}/latent/*.npz

2. Extract Features and Train Classifier

torchrun --nproc_per_node=4 linear_image.py \
    --name layer-12 \
    --time 20 \
    --epoch 40 \
    --lr 1e-3 \
    --batch_size 128 \
    --use_amp

---

Output

• Extracted features are saved to:

  {dataset_name}/features/train_features_t20_layer-12.npz
  {dataset_name}/features/val_features_t20_layer-12.npz
  

• Final classification metrics are appended to:

  {dataset_name}/result/final_results.csv
  

Metrics Reported:

• Single-label classification: Top-1 Accuracy, Top-5 Accuracy;
• Multi-label classification: mAP, CP, CR, CF1, OP, OR, OF1, Top-3 versions of the above metrics;

---

Text-only representation extraction

This module provides the code for extracting intermediate representations from a pretrained language diffusion model Plaid 1B and training a simple linear classifier for multi-label classification. It is designed to support experiments on datasets such as VG500 with minimal modification.

---

Requirements

Please refer to the Plaid repository for installation and dependency instructions.

---

Usage

1. Prepare Pretrained Weights

Download the Plaid 1B weights and place them in:

../plaid1b_weights/
├── model.pt
├── embedding_matrix.pt
├── noise_schedule.pt
└── gamma_bounds.pt

2. Prepare the Dataset

Place your training and testing CSVs in dataset/{dataset_name}/.

Each row in the CSV should include:

• image_id: unique identifier
• caption: the input text
• label: a list of binary values (as stringified list), e.g., "[0, 1, 0, ..., 1]"

3. Run the Script

torchrun --nproc_per_node=4 linear_language.py --timestep 30 --blockname layer-12

Optional arguments:

• --timestep: index in the diffusion schedule (default: 0)
• --blockname: transformer block name to extract activations (e.g., layer-20)

---

Output

• Extracted features are saved as .npz files in {dataset_name}/features/
• After training, the script prints:
  • mAP, CP, CR, CF1, OP, OR, OF1
  • Top-3 versions of the above metrics

---

Notes

• Different datasets may require different seq_len values, which can significantly impact experimental results. Please adjust seq_len accordingly in the configuration, and analyze the token length distribution using token_distribution.py to make an informed choice.
• Feature extraction supports deterministic ODE-style sampling by default. For stochastic DDPM/Appendix experiments, you can modify the sampler accordingly.

---

Cross-modal fusion representation

This module fuses latent features extracted from pretrained diffusion models in two modalities — image and language — for downstream multi-label classification. It supports flexible fusion strategies, distributed training, and standard evaluation metrics.

---

Key Components

Inputs

• Image features: From diffusion models (e.g., DiT), saved as .npz
• Language features: From language diffusion models(e.g., Plaid), saved as .npz
• Labels: Multi-label vectors (e.g., VG500 object categories), saved as .csv

Fusion Methods

• Linear Addition
• Linear Concat
• Simple Concat
• Cross Attention

---

Usage

1. Prepare Input Data

Ensure the following files exist:

• Image latent features:

  ../DiT/VG500/features/train_t{image_time}_{image_name}.npz
  ../DiT/VG500/features/test_t{image_time}_{image_name}.npz
  

• Text latent features:

  ../plaid/VG500/features/train_t{lang_name}_{lang_time}.npz
  ../plaid/VG500/features/test_t{lang_name}_{lang_time}.npz
  

• Label files:

  ../plaid/dataset/VG500/train_vg.csv
  ../plaid/dataset/VG500/test_vg.csv
  

2. Run Training

torchrun --nproc_per_node=4 train_multimodal_fusion.py \
    --image_time 100 --image_name layer-8 \
    --lang_time 30 --lang_name layer-12 \
    --batch_size 128 --epoch 40 --lr 1e-3 --use_amp

---

Output & Metrics

At the end of training, the following evaluation metrics are printed:

• mAP, CP, CR, CF1, OP, OR, OF1
• Top-3 versions of the above metrics

---

Acknowledgments

This project is based on the DiT, Plaid and DDAE. Thanks for their great work!