Aligned Text Image Encoders

Introduction

Aligned Text-Image Encoders is a PyTorch framework for learning a shared embedding space between text and images using frozen pretrained encoders, trainable adaptors, and the SigLIP contrastive objective [1].

The framework is designed to be simple to understand, modify, and extend for education, research and exploration on multimodal representation learning.

Features:

• Supports learning text-image alignment using frozen pretrained vision and text encoders with trainable adaptor networks.
• Uses the SigLIP sigmoid contrastive loss with learnable temperature and bias.
• Adapted the SigLIP chunking mechanism for single-GPU training, enabling training on a single GPU.
• Fully configurable training pipeline, including datasets, model architectures, optimizer, learning-rate scheduler, loss, and training hyperparameters.
• Implements checkpointing, training resumption, TensorBoard logging, and configuration management.
• Provides utilities for precomputing image and text backbone features for faster training on large datasets.
• Includes inference tools for bidirectional image-to-text and text-to-image retrieval, with both standard and streaming modes for small and large image collections.
• Scripts for curating DataComp metadata by filtering samples, selecting high-quality image-text pairs using CLIP similarity scores, and exporting train/validation TSV files.
• Designed to be easy to read, modify, and extend for experimentation with multimodal retrieval models.

The current implementation has been successfully trained on approximately 3 million curated DataComp image-text pairs using a single GPU. Evaluated in a zero-shot setting on the ImageNet-1K validation set (50,000 images and 1000 classes) with CLIP-style class prompts, the model achieves:

• Top-1 Accuracy: 31.12%
• Top-5 Accuracy: 51.40%

These results were obtained without extensive hyperparameter tuning and should be considered a baseline. The framework is fully configurable, and performance may be further improved by tuning the training hyperparameters and model architecture (e.g., the text adaptor depth).

Quick Installation

1. Create a virtual environment

make venv

2. Activate it

source .venv/bin/activate

3. Install Dependencies and CUDA-Enabled PyTorch

make install-gpu
make install-dev

If you want to change the version of cuda enabled torch (currently CUDA 12.8), you can modify install-gpu section in Makefile.

4. Verify GPU support

make check-gpu

How To Use

Dataset Preparation

Before training, an image-text dataset can be prepared from DataComp metadata [2]. The provided scripts download DataComp metadata, filter and curate high-quality image-text pairs using CLIP similarity scores, create train/validation TSV files, and then download the corresponding images into WebDataset format.

The dataset preparation pipeline performs the following steps:

• Downloads DataComp Small metadata.
• Cleans captions and image URLs.
• Filters invalid samples and likely non-English captions.
• Selects the highest-scoring image-text pairs using CLIP similarity scores.
• Creates train/validation TSV files.
• Uses img2dataset to download images and save them as WebDataset shards.

To download, filter, and curate an image-text dataset from the DataComp Small metadata, run:

bash scripts/download_datacomp_random_subset.sh

By default, the script creates a dataset containing 5.7 million image-text pairs, which due to expired or inaccessible image URLs, the final downloaded dataset typically contains around 3.1 million samples (approximately a 53% download success rate). You can change the target dataset size by modifying TARGET_TOTAL_SAMPLES in scripts/create_random_datacomp_subset.py.

Training

Training consists of two steps.

• First, you precompute frozen vision and text backbone features to reduce training time.
• Then, train the text adaptor using the SigLIP objective.

Precomputing Backbone Features

The framework provides a utility for augmenting an existing image-text dataset with precomputed vision and text backbone features. The resulting dataset can significantly reduce training time since the frozen backbone encoders no longer need to be executed during training.

aligned_text_image_encoders_augment_dataset \
  --root /path/to/dataset/datacomp-3M \
  --augmented-root /save/path/dataset/datacomp-3M-with-Feature \
  --batch-size 1024

Training

Once the augmented dataset has been created, start training by pointing the configuration to the augmented dataset:

aligned_text_image_encoders train \
  --set data.augmented_root=/path/to/augmented-dataset

Training can be resumed from any saved checkpoint:

aligned_text_image_encoders train \
  --set training.resume_path=/path/to/checkpoints/last.pt

The training CLI is configuration-driven. Any field in the default experiment configuration can be overridden from the command line using the --set section.field=value syntax. Multiple --set arguments may be supplied to customize datasets, model architecture, optimizer, learning-rate scheduler, loss function, and training hyperparameters without modifying the source code.

For all available configuration options, use:

aligned_text_image_encoders train --help

To inspect the resolved experiment configuration before starting training:

aligned_text_image_encoders train --print-config

For practical examples, please refer to commands_example/train.sh.

Training progress, including losses, learning rate, temperature, bias, and other metrics, is logged to TensorBoard. To monitor training:

tensorboard --logdir runs

• The default configuration is a reasonable starting point. No extensive hyperparameter sweep was performed, so you are encouraged to tune the configuration for your own use case.

Inference

After training, image-text retrieval can be performed using one of three inference commands, depending on the size of your dataset.

For all available options and configurable parameters, use the built-in help:

aligned_text_image_encoders infer --help
aligned_text_image_encoders infer-files --help
aligned_text_image_encoders infer-stream --help

For practical examples, please refer to commands_example/infer.sh.

• infer: Runs bidirectional image-to-text and text-to-image retrieval for a small collection of images and texts provided directly on the command line. This mode loads all images and texts into memory and is intended for quick experiments and debugging.

• infer-files: Runs bidirectional retrieval for all images in a folder against all candidate texts stored in a text file (one text per line). Both images and texts are loaded into memory, making this mode suitable for small to medium-sized datasets.

• infer-stream: Runs the same retrieval pipeline as infer-files, but loads all texts only once while loading and processing images in configurable chunks. This greatly reduces memory usage and is recommended for large image collections such as ImageNet or larger datasets.

The inference commands automatically restore the model configuration from the checkpoint, including the pretrained backbone configuration, trainable adaptor weights, and the learned SigLIP temperature and bias parameters.

For ImageNet-1K evaluation, CLIP-style class prompts are provided in src/aligned_text_image_encoders/evaluation/imagenet_class_prompts/. To see how to evaluate on ImageNet-1K validation set, please refer to commands in commands_example/evaluate_on_imagenet.sh.

Results

The default configuration (aligned_text_image_encoders train --print-config) was used without hyperparameter tuning. The model consists of a frozen DINOv3 ViT-L/16 vision encoder, a frozen BERT-Base text encoder, and a 2-layer Transformer text adaptor, trained using the Single-GPU Chunked SigLIP Sigmoid Loss.

Training was performed on approximately 3 million curated DataComp image-text pairs using a single GPU. To reduce training time, image and text backbone features were precomputed beforehand. For training commands, please see commands_example/train_on_datacomp_3M.sh.

Model performance was evaluated using zero-shot image-to-text retrieval on the ImageNet-1K validation set (50,000 images and 1000 classes) with CLIP-style class prompts.

Metric	Result
Top-1 Accuracy	31.12%
Top-5 Accuracy	51.40%

Training was configured for 30,000 optimization steps. However, it was stopped manually after approximately 14,000 steps, since the ImageNet-1K evaluation accuracy had already reached the level typically obtained after training on roughly 3 million image-text pairs. The remaining training was therefore omitted.

The following TensorBoard plots show the training dynamics, including the training and validation losses, learning rate schedule, and the learned SigLIP temperature and bias parameters. As shown by the TensorBoard curves, the validation loss increases after approximately 5,000 training steps, while the training loss continues to decrease. Training was nevertheless continued because the validation loss, which depends heavily on the sampled negative pairs and a relatively small validation set, was not considered a reliable indicator of downstream retrieval performance.

The reported results should be considered a baseline, as no extensive hyperparameter search was performed. The framework exposes all training, optimization, and model parameters through the CLI, making it straightforward to explore alternative configurations, such as deeper text adaptors, different learning-rate schedules, optimizer settings, or larger training datasets.

License

Released under the MIT License. For used datasets or pre-trained models, please check their respective licenses.

References

- Zhai, Xiaohua, et al. "Sigmoid loss for language image pre-training." Proceedings of the IEEE/CVF international conference on computer vision. 2023.
- [DataComp](https://github.com/mlfoundations/datacomp)
- [ImageNet Dataset](https://www.image-net.org/update-mar-11-2021.php)