[ICML 2026] SOTAlign: Semi-Supervised Alignment of Unimodal Vision and Language Models via Optimal Transport

Simon Roschmann*, Paul Krzakala*, Sonia Mazelet, Quentin Bouniot, Zeynep Akata

Abstract

The Platonic Representation Hypothesis posits that neural networks trained on different modalities converge toward a shared statistical model of the world. Recent work exploits this convergence by aligning frozen pretrained vision and language models with lightweight alignment layers, but typically relies on contrastive losses and millions of paired samples. In this work, we ask whether meaningful alignment can be achieved with substantially less supervision. We introduce a semi-supervised setting in which pretrained unimodal encoders are aligned using a small number of image–text pairs together with large amounts of unpaired data. To address this challenge, we propose SOTAlign, a two-stage framework that first recovers a coarse shared geometry from limited paired data using a linear teacher, and then refines the alignment on unpaired samples via an optimal-transport-based divergence that transfers relational structure without overconstraining the target space. Unlike existing semi-supervised methods, SOTAlign effectively leverages unpaired images and text, learning robust joint embeddings across datasets and encoder pairs, and significantly outperforming supervised and semi-supervised baselines.

Methodology

SOTAlign is a two-step method for the alignment of pretrained unimodal image and text encoders. First, we fit a linear alignment model only using the limited amount of available image-text pairs. Then, we use this linear model as a teacher to regularize the training of alignment layers $f$ and $g$ for a joint embedding space leveraging unimodal (unpaired) data.

KLOT loss

For each batch of unpaired data, the linear teacher produces a pseudo-label affinity matrix $K^\ast$. Because the teacher is trained on only a small paired dataset, these predictions can be noisy. We therefore use entropy-regularized optimal transport to retain the most plausible globally consistent matches as training signal for the student model. Formally, the optimal transport plan is defined as

$$ \mathrm{OT}_{\epsilon}(K) = \underset{P \in \Pi_n}{\mathrm{arg,min}} \left[-\langle P, K \rangle + \epsilon H(P)\right], $$

where $\Pi_n$ is the set of $n \times n$ doubly stochastic transport plans and $H$ is the entropy regularizer. The transport plan is computed efficiently with the Sinkhorn algorithm.

Given the student's affinity matrix $K$, our KLOT loss compares its transport plan with the filtered teacher plan:

$$ \mathrm{KLOT}(K \Vert K^\ast) = D_{\mathrm{KL}}!\left( \mathrm{OT}_{\epsilon^\ast}(K^\ast) \Vert \mathrm{OT}_{\epsilon}(K) \right). $$

KLOT is fully differentiable, with the efficient gradient

$$ \nabla_K \mathrm{KLOT}(K \Vert K^\ast) = \frac{ \mathrm{OT}_{\epsilon}(K) - \mathrm{OT}_{\epsilon^\ast}(K^\ast) }{\epsilon}. $$

The implementation is available in src/optimal_transport/klot.py.

Installation

This repository works with Python 3.11 and PyTorch 2.7. Please create a conda environment and install the dependencies specified in requirements.txt.

conda create -n sotalign_env python=3.11 -y
conda activate sotalign_env
python -m pip install torch==2.7.1 torchvision==0.22.1
python -m pip install -r requirements.txt

The experiment scripts activate a conda environment themselves. Point them to the environment created above:

export CONDA_SH="$(conda info --base)/etc/profile.d/conda.sh"
export CONDA_ENV=sotalign_env

Data Preparation

The following commands will download CC3M images with multiple caption variants and extract the corresponding features. Ensure that the selected storage locations have sufficient capacity, as the CC3M images and extracted representations require several hundred gigabytes.

export PROJECT_DIR="$PWD"
export DATA_ROOT=/path/to/datasets
export ARTIFACT_ROOT=/path/to/artifacts
mkdir -p "$DATA_ROOT" "$ARTIFACT_ROOT"

Download CC3M

Download the PixParse WebDataset shards:

python tools/data_preparation/download_cc.py \
  --dataset cc3m \
  --output-dir "$DATA_ROOT"

Download the DreamLIP captions and match them to the CC3M PixParse images contained in the downloaded shards:

wget -P "$DATA_ROOT" \
  https://huggingface.co/datasets/qidouxiong619/dreamlip_long_captions/resolve/main/cc3m_3long_3short_1raw_captions_url.csv

python tools/data_preparation/match_image_caption.py \
  --dataset cc3m \
  --base_dir "$DATA_ROOT"

The recaptioned shards are saved to: $DATA_ROOT/cc3m_recaptioned/cc3m-train-*.tar.

Feature Extraction

The default encoders in this repository are facebook/dinov3-vitl16-pretrain-lvd1689m and nvidia/NV-Embed-v2. NV-Embed-v2 requires accepting the model terms on Hugging Face and authenticating with huggingface-cli login. Furthermore, you need to request access to the DINOv3 weights, download the ViT-L/16 checkpoint, and set DINOV3_CHECKPOINT_PATH=/path/to/checkpoint.pth.

For image tokens, we concatenate the class token and mean patch token (AGG_MODE=concat). For text features, we use model-specific pooling to produce one embedding per caption. Encoding stores each image or text embedding bank in a single HDF5 (.h5) file. Batches are appended incrementally to a gzip-compressed FP16 dataset named embeddings.

Extract the training features from the recaptioned CC3M shards:

PROJECT_DIR="$PROJECT_DIR" OUTPUT_DIR="$ARTIFACT_ROOT" \
  DATA=cc3m USE_CONDA=1 \
  DATA_ROOT="$DATA_ROOT/cc3m_recaptioned" \
  DOMAIN=image SPLIT=train \
  VISION_MODEL=facebook/dinov3-vitl16-pretrain-lvd1689m \
  AGG_MODE=concat BATCH_SIZE=32 \
  bash scripts/encode.sh

PROJECT_DIR="$PROJECT_DIR" OUTPUT_DIR="$ARTIFACT_ROOT" \
  DATA=cc3m USE_CONDA=1 \
  DATA_ROOT="$DATA_ROOT/cc3m_recaptioned" \
  DOMAIN=text SPLIT=train \
  TEXT_MODEL=nvidia/NV-Embed-v2 \
  SOURCE_CAPTION=raw_caption BATCH_SIZE=32 \
  bash scripts/encode.sh

Extract the validation features:

PROJECT_DIR="$PROJECT_DIR" OUTPUT_DIR="$ARTIFACT_ROOT" \
  DATA=cc3m USE_CONDA=1 \
  DATA_ROOT="$DATA_ROOT/cc3m" \
  DOMAIN=image SPLIT=validation \
  VISION_MODEL=facebook/dinov3-vitl16-pretrain-lvd1689m \
  AGG_MODE=concat BATCH_SIZE=32 \
  bash scripts/encode.sh

PROJECT_DIR="$PROJECT_DIR" OUTPUT_DIR="$ARTIFACT_ROOT" \
  DATA=cc3m USE_CONDA=1 \
  DATA_ROOT="$DATA_ROOT/cc3m" \
  DOMAIN=text SPLIT=validation \
  TEXT_MODEL=nvidia/NV-Embed-v2 \
  SOURCE_CAPTION=txt BATCH_SIZE=32 \
  bash scripts/encode.sh

The extracted features are organized as follows:

$ARTIFACT_ROOT/tensor_data/
├── image_embedding/dinov3-vitl16-pretrain-lvd1689m/
│   ├── cc3m_concat.h5
│   └── cc3m_concat_validation.h5
└── text_embedding/NV-Embed-v2/
    ├── cc3m_raw_caption.h5
    └── cc3m_raw_caption_validation.h5

Multimodal Alignment

The training scripts load the extracted image and text features and only train lightweight alignment layers for a joint embedding space. The best checkpoint is selected based on the mean image-to-text/text-to-image Recall@1 on the CC3M validation split. Training can be logged using wandb after authentication via wandb login.

All experiments can be run on one A100 GPU with 80GB memory. When scheduling a cluster job, request sufficient CPU memory. The embedding datasets load all HDF5 features into host RAM during initialization such that subsequent __getitem__ calls require only tensor indexing. This enables fast training with minimal data loading overhead.

Set the following paths once:

export BASE_EMBEDDING_DIR="$ARTIFACT_ROOT/tensor_data"
export OUTPUT_DIR="$ARTIFACT_ROOT/logs"
export BACKBONE_FEATURES_DIR="$ARTIFACT_ROOT/backbone_features"
export DATASET_ROOT_DIR="$DATA_ROOT"

export IMAGE_MODEL=facebook/dinov3-vitl16-pretrain-lvd1689m
export DINOV3_CHECKPOINT_PATH=/path/to/dinov3_vitl16_checkpoint.pth
export TEXT_MODEL=nvidia/NV-Embed-v2

export WANDB_PROJECT=semisupervised_alignment

By default, training targets 2,000 optimization steps, uses 10% of the scheduled steps for warmup, and chooses a validation interval such that there are at most 50 validation runs. These values can be overridden through the launchers with TARGET_STEPS, EPOCH_NUM, WARMUP, VAL_FREQUENCY, and SAVE_FREQUENCY. Set only one of TARGET_STEPS and EPOCH_NUM; explicitly provided warmup and frequency values are always respected.

After training, every launcher evaluates its best checkpoint by default. Prepare the required downstream datasets using the instructions in the Evaluation section below.

(Semi-)Supervised Baselines

The baseline launchers default to the best DINOv3 ViT-L/NV-Embed-v2 configuration found in our hyperparameter sweeps:

bash scripts/linear_baseline.sh
bash scripts/sail_supervised.sh
bash scripts/structure_supervised.sh
bash scripts/ot_clip.sh
bash scripts/sail_semisupervised.sh
bash scripts/structure_semisupervised.sh
bash scripts/sclip.sh
bash scripts/nnclr.sh

NNCLR Nearest Neighbors

NNCLR requires a neighbor index matrix per modality. Use generate_neighbor_indices.py with the commands below to compute these files. Run the index generator once for the text data and once for the image data.

python tools/data_preparation/generate_neighbor_indices.py \
  --query "$BASE_EMBEDDING_DIR/text_embedding/NV-Embed-v2/cc3m_raw_caption.h5" \
  --output "$BASE_EMBEDDING_DIR/text_embedding/NV-Embed-v2/cc3m_raw_caption_neighbors.npy" \
  --k 64 --batch_size 65536

python tools/data_preparation/generate_neighbor_indices.py \
  --query "$BASE_EMBEDDING_DIR/image_embedding/dinov3-vitl16-pretrain-lvd1689m/cc3m_concat.h5" \
  --output "$BASE_EMBEDDING_DIR/image_embedding/dinov3-vitl16-pretrain-lvd1689m/cc3m_concat_neighbors.npy" \
  --k 64 --batch_size 65536

SOTAlign

The default SOTAlign configuration utilizes 10k pairs and 1M unpaired images and texts. SOTAlign first fits a linear CCA teacher on the paired examples to recover a coarse cross-modal geometry. The linear projection layers are then trained with a supervised SAIL/SigLIP loss on the pairs and the KLOT regularizer on unpaired batches using the teacher's OT plan as the training signal.

SUPERVISED_TEXT_EMBEDDING="$BASE_EMBEDDING_DIR/text_embedding/NV-Embed-v2/cc3m_raw_caption.h5" \
SUPERVISED_IMAGE_EMBEDDING="$BASE_EMBEDDING_DIR/image_embedding/dinov3-vitl16-pretrain-lvd1689m/cc3m_concat.h5" \
UNSUPERVISED_TEXT_EMBEDDING="$BASE_EMBEDDING_DIR/text_embedding/NV-Embed-v2/cc3m_raw_caption.h5" \
UNSUPERVISED_IMAGE_EMBEDDING="$BASE_EMBEDDING_DIR/image_embedding/dinov3-vitl16-pretrain-lvd1689m/cc3m_concat.h5" \
UNSUPERVISED_INDEX_MODE=disjoint \
N_SUPERVISED_PAIRS=10000 \
N_UNSUPERVISED_IMAGE=1000000 \
N_UNSUPERVISED_TEXT=1000000 \
BATCH_SIZE=32768 \
BATCH_SIZE_SUPERVISED=10000 \
AFFINITY=cca \
CCA_LAM_X=0.1 \
CCA_LAM_Y=0.1 \
EIG_EPS=1e-6 \
ALPHA_SEMISUPERVISED_SAIL=1.0 \
ALPHA_SEMISUPERVISED_KLOT=0.001 \
EPSILON_SINKHORN_STUDENT=0.05 \
N_ITERS_SINKHORN_STUDENT=100 \
EPSILON_SINKHORN_TEACHER=0.01 \
N_ITERS_SINKHORN_TEACHER=100 \
TARGET_DIM=1024 \
TARGET_STEPS=2000 \
P_NOISE=0.0 \
P_NOISE_UNSUPERVISED=0.0 \
SIGMA_K=0.0 \
bash scripts/sotalign.sh

The corresponding launcher parameters are:

Environment variable	Default	Description
UNSUPERVISED_INDEX_MODE	disjoint	How indices are selected when the unpaired image and text pools are compatible.
N_SUPERVISED_PAIRS	10000	Number of image-text pairs.
N_UNSUPERVISED_IMAGE	1000000	Maximum number of unpaired images.
N_UNSUPERVISED_TEXT	1000000	Maximum number of unpaired texts.
BATCH_SIZE	32768	Total number of paired and unpaired samples in each batch.
BATCH_SIZE_SUPERVISED	10000	Number of pairs included in each batch.
AFFINITY	cca	Linear teacher used to construct the target cross-modal geometry.
CCA_LAM_X	0.1	Covariance regularization for the teacher's first input space.
CCA_LAM_Y	0.1	Covariance regularization for the teacher's second input space.
EIG_EPS	1e-6	Minimum eigenvalue used for numerical stability when fitting CCA.
ALPHA_SEMISUPERVISED_SAIL	1.0	Weight of the SAIL/SigLIP loss.
ALPHA_SEMISUPERVISED_KLOT	0.001	Weight of the KLOT regularization.
EPSILON_SINKHORN_STUDENT	0.05	Entropic OT regularization in the learned joint space.
N_ITERS_SINKHORN_STUDENT	100	Number of student-space Sinkhorn iterations.
EPSILON_SINKHORN_TEACHER	0.01	Entropic OT regularization in the teacher reference space.
N_ITERS_SINKHORN_TEACHER	100	Number of teacher-space Sinkhorn iterations.
TARGET_DIM	1024	Dimension of the learned joint space.
TARGET_STEPS	2000	Number of training iterations.
P_NOISE	0.0	Fraction of paired examples corrupted when fitting the teacher; nonzero values for robustness experiments.
P_NOISE_UNSUPERVISED	0.0	Fraction of each unpaired pool replaced with random unit vectors; nonzero values for robustness experiments.
SIGMA_K	0.0	Standard deviation of Gaussian noise added to the teacher affinity matrix; nonzero values for robustness experiments.

Checkpoints are written to $OUTPUT_DIR/<run-name>/checkpoints/epoch_best.pt.

Evaluation

Representational Alignment

We measure the intrinsic similarity of frozen vision and language encoders before alignment using mutual k-nearest-neighbor overlap on paired CC3M samples. For each image-text pair, we find its 10 nearest examples independently in the image and text spaces and report the average fraction of shared neighbors. The probe reads the CC3M image and text embeddings extracted above from $BASE_EMBEDDING_DIR.

Compute the mutual k-NN score for DINOv3 ViT-L and NV-Embed-v2 on 1,024 CC3M pairs:

BASE_EMBEDDING_DIR="$ARTIFACT_ROOT/tensor_data" \
bash scripts/measure_alignment.sh

The resulting table is saved to $ARTIFACT_ROOT/alignment/cc3m_mutual_knn.csv.

Zero-Shot Retrieval and Classification

We evaluate the (semi-)supervised alignment methods on the following tasks:

Task	CLI name	Metrics
COCO 2017 image–text retrieval	COCO	Text-to-image (T2I) and image-to-text (I2T) Recall@1, @5, and @10
ImageNet-1K zero-shot classification	imagenetv1	Top-1 and top-5 accuracy

We additionally evaluate retrieval on Flickr30k and classification on Food-101, CIFAR-10, CIFAR-100, FGVC-Aircraft, DTD, and Oxford Flowers using the SAIL evaluation repository.

Download and extract the COCO 2017 validation split:

mkdir -p "$DATA_ROOT/coco/2017"

wget -P "$DATA_ROOT/coco/2017" http://images.cocodataset.org/annotations/annotations_trainval2017.zip
wget -P "$DATA_ROOT/coco/2017" http://images.cocodataset.org/zips/val2017.zip

unzip -q "$DATA_ROOT/coco/2017/annotations_trainval2017.zip" -d "$DATA_ROOT/coco/2017"
unzip -q "$DATA_ROOT/coco/2017/val2017.zip" -d "$DATA_ROOT/coco/2017"

Prepare the ImageNet-1k validation split:

1. Accept the access conditions on the ImageNet-1K huggingface page and authenticate with huggingface-cli login.
2. Download the validation split and prepare an ImageFolder-compatible layout at $DATA_ROOT/imagenet/val/<class-index>/*.JPEG:

python tools/data_preparation/download_imagenet_huggingface.py --output-dir "$DATA_ROOT"

Finally, evaluate any alignment checkpoint by specifying the checkpoint path and the image and text encoders used to generate the training features:

CHECKPOINT=/path/to/run/checkpoints/epoch_best.pt \
DATA_ROOT=/path/to/datasets \
IMAGE_MODEL=facebook/dinov3-vitl16-pretrain-lvd1689m \
TEXT_MODEL=nvidia/NV-Embed-v2 \
EVAL_TASKS="COCO imagenetv1" \
bash scripts/eval.sh

By default, evaluation results are stored under the corresponding run directory in results/<task>/. Extracted backbone features are cached in $ARTIFACT_ROOT/backbone_features/; set BACKBONE_FEATURES_DIR to override this location.

Acknowledgments

This project builds on the Platonic Representation Hypothesis, SAIL, OpenCLIP and DreamLIP. We appreciate their work.

Citation

If you find SOTAlign useful, please star the repository and cite:

@inproceedings{roschmann2026sotalign,
  title={SOTAlign: Semi-Supervised Alignment of Unimodal Vision and Language Models via Optimal Transport},
  author={Simon Roschmann and Paul Krzakala and Sonia Mazelet and Quentin Bouniot and Zeynep Akata},
  booktitle={International Conference on Machine Learning},
  year={2026}
}

Contact

Questions are welcome at simon.roschmann@tum.de.