google-deepmind / google-deepmind/optax

Feature request: add BYOL, SimSiam, DINO and Barlow Twins losses to `optax.losses`

Open
#1,528 20 comments 0 reactions 0 assignees View on GitHub
type:feature
Dominant language
Python
Stars
2.3k
Forks
369
Avg merge
10h 15m
Merged PRs (30d)
7

Description

I’d like to propose adding a few popular self-supervised losses to
`optax.losses._self_supervised`:

- `byol_loss`
- `simsiam_loss`
- `dino_loss`
- `barlow_twins_loss`

---

### Motivation

These objectives are widely used in modern self-supervised representation
learning pipelines, especially for vision, and having them in Optax would:

- Make it easier to prototype and compare self-supervised methods on top of
JAX/Flax.
- Provide a single, well-tested implementation instead of many slightly
different copies.
- Nicely complement the existing `ntxent` and triplet margin losses already
present in `_self_supervised.py`.

---

### Proposed API (high-level)

All functions follow the same style as existing Optax losses:

- Pure JAX functions, compatible with `jit`/`vmap`.
- `jax.typing.ArrayLike` arguments and `jax.Array` return types.
- Shape checks and `utils.check_subdtype` for float inputs.
- Docstrings with examples and references.

Rough signatures:

def byol_loss(
online_projection_1: jax.typing.ArrayLike,
target_projection_2: jax.typing.ArrayLike,
online_projection_2: jax.typing.ArrayLike,
target_projection_1: jax.typing.ArrayLike,
eps: jax.typing.ArrayLike = 1e-6,
) -> jax.Array:

def simsiam_loss(
predictor_projection_1: jax.typing.ArrayLike,
target_projection_2: jax.typing.ArrayLike,
predictor_projection_2: jax.typing.ArrayLike,
target_projection_1: jax.typing.ArrayLike,
eps: jax.typing.ArrayLike = 1e-6,
) -> jax.Array:

def dino_loss(
student_logits: jax.typing.ArrayLike,
teacher_logits: jax.typing.ArrayLike,
student_temperature: jax.typing.ArrayLike = 0.1,
teacher_temperature: jax.typing.ArrayLike = 0.04,
teacher_center: jax.typing.ArrayLike = 0.0,
) -> jax.Array:

def barlow_twins_loss(
projection_1: jax.typing.ArrayLike,
projection_2: jax.typing.ArrayLike,
off_diagonal_scale: jax.typing.ArrayLike = 5e-3,
eps: jax.typing.ArrayLike = 1e-12,
) -> jax.Array:

---

### References

- BYOL – *Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning*
https://arxiv.org/abs/2006.07733

- SimSiam – *Exploring Simple Siamese Representation Learning*
https://arxiv.org/abs/2011.10566

- DINO – *Emerging Properties in Self-Supervised Vision Transformers*
https://arxiv.org/abs/2104.14294

- Barlow Twins – *Barlow Twins: Self-Supervised Learning via Redundancy Reduction*
https://arxiv.org/abs/2103.03230

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.