google / google/tunix

Feature request: Add Contrastive Representation Distillation (CRD)

Open
#887 0 comments 0 reactions 0 assignees View on GitHub
type:feature/enhancement
Dominant language
Python
Stars
2.5k
Forks
345
Avg merge
1d 7h
Merged PRs (30d)
240

Description

Requesting a built-in **Contrastive Representation Distillation (CRD)** strategy (Tian et al., 2019) for `tunix.distillation.DistillationTrainer`.

## What CRD does
Distills at the **representation level** using an InfoNCE/contrastive loss:
- positive pair: (student rep, teacher rep) from the **same sample**
- negatives: mismatched pairs (e.g., in-batch negatives)

Minimal form:
- `logits = (z_s @ z_t.T) / tau`
- `labels = arange(B)`
- `loss = CE(logits, labels)` (optional symmetric term on `logits.T`)

## Why this helps
- Complements / improves over logit-KD in many setups

## Reference
Tian et al., *Contrastive Representation Distillation*, arXiv:1910.10699 (2019): https://arxiv.org/pdf/1910.10699

I can contribute a PR + tests if you’re open to it.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.