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

[surajyadav-research]

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

[natinew77-creator]

This is a great suggestion! Adding these popular self-supervised losses would be a valuable addition to Optax.

I can start working on implementing byol_loss, simsiam_loss, dino_loss, and barlow_twins_loss and submit a PR.

[surajyadav-research]

@natinew77-creator I also had code prepared for this feature and was hoping to contribute it myself. It would’ve helped if I was asked before the implementation was taken over. No worries just sharing this so future contributions can be better coordinated.

[natinew77-creator]

Hi @surajyadav-research,

I sincerely apologize for that! I saw the feature request and got excited to contribute, so I jumped straight into the implementation. I didn't realize you intended to work on it yourself since the issue was unassigned.

Thank you for being gracious about it. Since the PR is already up, I would strictly value your feedback if you'd be willing to review it? I want to make sure the implementation aligns with what you had in mind.

I'll definitely be more mindful of coordinating on future issues. Thanks again for the great suggestion!

[surajyadav-research]

Hi @natinew77-creator ,
No worries at all. I just wanted to mention that I was already working on this feature on my end. I’ve made several improvements and addressed issues, so I’ll be opening a new PR shortly with the updated implementation.

[natinew77-creator]

Thanks for the update, @surajyadav-research! That sounds great.

Since you have a more comprehensive implementation ready, I will close my PR now to avoid confusion for the maintainers. I look forward to seeing your implementation!

Best of luck with the PR!

[surajyadav-research]

Hi @natinew77-creator
I just went through the related papers, the authors’ implementations, and your code as well. I’ve committed an updated version with a few improvements—maybe we can collaborate on this and refine it further together?

[natinew77-creator]

Hi @surajyadav-research, that sounds like a great plan! I'd be happy to collaborate.

Thanks for taking the time to review the papers and my code. I see you've opened #1531. I can head over there to review your changes, and we can discuss any further refinements together.

To keep the history clean, I'll go ahead and close my PR (#1530) so we can focus our efforts on yours.

Looking forward to working with you on this!

[surajyadav-research]

Hi @natinew77-creator happy to collaborate!

I’m not entirely sure how to give collaborators access on a PR. Could you guide me on how to set that up?

Also, the CI is currently failing on the “check links in markdown” step. The links themselves seem valid when opened manually, so I’m not sure how to resolve that.
If you have suggestions on how to handle this (or whether we need to update the link-check config), please let me know.

Looking forward to working together!

[natinew77-creator]

Hi @surajyadav-research! Great to collaborate with you on this.

Regarding collaborator access:

There are a few ways we can work together:

Review Comments (easiest): I can review your PR and leave detailed suggestions using GitHub's "Suggest changes" feature. You can then batch-commit those directly.
Enable "Allow edits from maintainers": If you edit your PR settings and check "Allow edits from maintainers", fork collaborators can push directly. Though I believe this only works for repository maintainers, not other fork contributors.
Add me to your fork: Go to your fork's Settings → Collaborators → Add people, and add my GitHub username. This would let me push commits to your branch directly.
For now, I'll head over to the PR and leave detailed review comments!

Regarding the CI "check links in markdown" failure:

Good news - this isn't caused by your code changes! I investigated the failure logs and found:

The failing links are in README.md and
docs/development.md
(existing repo files, not files you touched)
Links to
docs/getting_started.ipynb
and CONTRIBUTING.md are returning 503 (Service Unavailable)
This is a transient issue - GitHub sometimes rate-limits or blocks automated link checkers
How to resolve it:

Simply re-run the failed job - it will likely pass on retry since it's not a real broken link
If it fails repeatedly, a maintainer may need to temporarily ignore these links in the CI config, or you can ask for a maintainer to override
The actual tests for your code (pytest, linting, etc.) should all pass since your implementation looks solid.

I'll head over to PR #1531 now to review the code!

[surajyadav-research]

Hi @natinew77-creator ,
I’ve given you access—so you can now commit directly.

Regarding the “check links in markdown” failure: I reran the workflow but it’s still failing.
Feel free to take a look and leave your review and make changes; we can improve this together.

Happy to collaborate on this!

[natinew77-creator]

Awesome, I've accepted the invite! Thanks for adding me.

I'll head over to #1531 right now. Since the link checker is still being stubborn, I'll see if I can tweak the config to ignore those specific flaky URLs or just double-check why they are timing out.

I'll also start reviewing the implementation and see where I can chip in with tests or docs. Let's get this merged!