README.md

---
license: cc-by-nc-4.0
library_name: pytorch
pipeline_tag: image-feature-extraction
tags:
  - chest-xray
  - medical-imaging
  - self-supervised-learning
  - ijepa
  - jepa
  - vision-transformer
  - pytorch
  - mimic-cxr
datasets:
  - MIMIC-CXR-JPG
  - VinBigData-CXR
---

# TexJEPA Model Weights

This repository hosts the trained model weights for **TexJEPA** from:

**When Texture Becomes the World: Texture-aware JEPA for Chest X-ray Representation Learning**

Authors: **Yi Zhao\***, **Ruilang Wang\***, Bowen Liu, Donglong Chen†<br>
\* Equal contribution. † Corresponding author.

Code repository: [YiZhao-Jasper/TexJEPA](https://github.com/YiZhao-Jasper/TexJEPA)

## Summary

TexJEPA is a texture-aware extension of Image Joint-Embedding Predictive Architecture (I-JEPA) for chest X-ray representation learning. The models are trained on MIMIC-CXR-JPG with a ViT-H/14 backbone and an I-JEPA-style context-target predictive objective. The released weights are intended to reproduce the paper's pre-training lineage and downstream evaluation using the public TexJEPA codebase.

The repository includes four checkpoints:

| Model | File | Initialization | Training stage | Main purpose |
| --- | --- | --- | --- | --- |
| **I-JEPA-300** | `weights/I-JEPA-300/jepa-latest.pth.tar` | I-JEPA ViT-H/14 pre-training run | 300-epoch final/latest checkpoint | Baseline chest X-ray I-JEPA representation |
| **TexJEPA-N** | `weights/TexJEPA-N/jepa-latest.pth.tar` | I-JEPA-300 lineage | +50 epoch texture-noise specialization | Learn invariance to nuisance texture corruption |
| **TexJEPA-R** | `weights/TexJEPA-R/jepa-latest.pth.tar` | TexJEPA-N | +50 epoch register-token branch | Route high-norm texture artifacts through register tokens |
| **TexJEPA-C** | `weights/TexJEPA-C/jepa-latest.pth.tar` | TexJEPA-N | +50 epoch covariance-regularized branch | Preserve local patch diversity and lesion-scale texture sensitivity |

The historical internal experiment names are not used in this release. Please cite and report the public names above.

## Files

```text
weights/
  I-JEPA-300/jepa-latest.pth.tar
  TexJEPA-N/jepa-latest.pth.tar
  TexJEPA-R/jepa-latest.pth.tar
  TexJEPA-C/jepa-latest.pth.tar
configs/pretrain/
  ijepa_300.yaml
  texjepa_n.yaml
  texjepa_r.yaml
  texjepa_c.yaml
checksums.sha256
requirements.txt
LICENSE
```

`checksums.sha256` contains SHA-256 hashes for the uploaded checkpoint files.

## Model Differences

### I-JEPA-300

I-JEPA-300 is the chest X-ray baseline. It uses the original I-JEPA multi-block masking objective with a ViT-H/14 context encoder, EMA target encoder, and predictor. It is the reference model for clean representation quality and texture-sensitivity analysis.

### TexJEPA-N

TexJEPA-N adds **context-target asymmetric texture corruption**. The target branch receives clean chest radiographs, while the context branch receives mild stochastic texture perturbations such as Gaussian noise, Poisson noise, and JPEG-like compression. This trains the model to predict clean latent targets from corrupted context views, directly targeting texture robustness.

### TexJEPA-R

TexJEPA-R starts from TexJEPA-N and adds **register tokens** to the encoder. Register tokens participate in self-attention but are stripped before patch outputs are returned, preserving compatibility with the I-JEPA predictor. The branch also uses tighter local masking to encourage stronger local, lesion-scale structure modeling.

### TexJEPA-C

TexJEPA-C starts from TexJEPA-N and adds a **patch-token variance/covariance auxiliary loss**. The regularizer discourages patch-token collapse and helps preserve local texture diversity. The implementation computes the auxiliary loss per rank and avoids cross-rank all-gather, which is important because I-JEPA masks can produce variable visible-token counts across distributed workers.

## Intended Use

These weights are intended for:

- Reproducing TexJEPA pre-training and downstream experiments with the GitHub repository.
- Chest X-ray representation learning research.
- Linear probing and fine-tuning on VinBigData/VinDr-style multi-label chest X-ray classification.
- Texture robustness, lesion sensitivity, and model comparison experiments.

They are **not** intended for clinical deployment or diagnostic use.

## Reproduction

Clone the code:

```bash
git clone https://github.com/YiZhao-Jasper/TexJEPA.git
cd TexJEPA
python -m pip install -r requirements.txt
```

Download this model repository and place or symlink the checkpoints under the paths expected by the configs:

```text
logs/ijepa_300/jepa-latest.pth.tar             <- weights/I-JEPA-300/jepa-latest.pth.tar
logs/texjepa_n/jepa-latest.pth.tar             <- weights/TexJEPA-N/jepa-latest.pth.tar
logs/texjepa_r/jepa-latest.pth.tar             <- weights/TexJEPA-R/jepa-latest.pth.tar
logs/texjepa_c/jepa-latest.pth.tar             <- weights/TexJEPA-C/jepa-latest.pth.tar
```

The pre-training configs in this model repository mirror the GitHub release:

```bash
python scripts/sanity_check.py
```

For downstream evaluation:

```bash
export CHECKPOINT=logs/texjepa_n/jepa-latest.pth.tar
export VINBIG_IMAGE_DIR=data/vinbig/images_1024/train
export VINBIG_CSV=data/vinbig/annotations/train.csv
bash downstream/run_downstream.sh all
```

## Data

Pre-training uses **MIMIC-CXR-JPG**. This model repository does not redistribute MIMIC-CXR data. Researchers must obtain MIMIC-CXR-JPG through the official PhysioNet credentialing and data-use process.

Downstream evaluation code supports VinBigData/VinDr-style multi-label chest X-ray annotations; those datasets are not redistributed here.

## Checkpoint Format

The checkpoints are PyTorch `.pth.tar` training checkpoints containing encoder, predictor, target encoder, optimizer metadata, and training metadata where available. Load only checkpoints from trusted sources because PyTorch checkpoint loading relies on Python pickle serialization.

For representation extraction, the downstream code in the GitHub repository loads the `target_encoder` weights and supports ViT-H/14, ViT-L/14, and register-token variants by reading checkpoint metadata and tensor shapes.

## Citation

```bibtex
@misc{zhao2026texjepa,
  title  = {When Texture Becomes the World: Texture-aware JEPA for Chest X-ray Representation Learning},
  author = {Zhao, Yi and Wang, Ruilang and Liu, Bowen and Chen, Donglong},
  year   = {2026},
  note   = {TexJEPA model weights}
}
```

## License

The released code and model weights are provided for non-commercial research use under the license included in this repository. MIMIC-CXR-JPG and downstream datasets remain governed by their own licenses and data-use agreements.