Reconstruction and colorization
One frozen visual backbone, two jobs: put a masked photo back together, and paint colour, without training two networks or the whole encoder.
- DINOv2
- LoRA
- PyTorch
- COCO
- 2 headsRGB repair and LAB colour
- ~90%fewer parameters trained
- LoRA 32rank that fits a 4 GB GPU
One backbone, two judgments
Reconstruction and colour are related, but they fail in different ways. Sharing the encoder and splitting the heads means each can be looked at on its own.
- Encoder: DINOv2 ViT-S/14, weights frozen.
- Shared convolutional decoder, then an RGB head and a LAB a* b* head.
- Course project with Prof. Biplab Banerjee.
What was allowed to train
Full fine-tuning of the backbone does not fit a 4 GB card. LoRA trains a low-rank difference instead of every weight.
- LoRA rank 32, plus the last transformer block. Trainable parameters drop by about 90%.
- COCO 2017, 224×224, 25% patch masking. Reconstruction loss is L1 with a little SSIM, so structure is kept. Colour is L1 on a* and b*.
- On a held-out photo the pose survives the mask. Colour is the head learning the scene, grass and kit, from the same features.