Barlow Twins - Redundancy Reduction for Self-Supervised Learning.

Implements Barlow Twins from "Barlow Twins: Self-Supervised Learning via Redundancy Reduction" (Zbontar et al., ICML 2021). Barlow Twins prevents representation collapse by pushing the cross-correlation matrix of two augmented views toward the identity matrix.

Key Innovation

The loss has two terms on the cross-correlation matrix C:

1. Invariance: Diagonal elements should be 1 (views agree per dimension)
2. Redundancy reduction: Off-diagonal elements should be 0 (dimensions should be independent)

L = SUM_i (1 - C_ii)^2 + lambda * SUM_{i!=j} C_ij^2

Architecture

Augmented View 1         Augmented View 2
      |                         |
      v                         v
+------------+           +------------+
|  Encoder   |           |  Encoder   |  (shared weights)
+------------+           +------------+
      |                         |
      v                         v
+------------+           +------------+
| Projector  |           | Projector  |  (shared weights)
+------------+           +------------+
      |                         |
      v                         v
     Z_A        Cross-Corr     Z_B
      |                         |
      +-------> Loss <---------+

Usage

model = BarlowTwins.build(encoder_dim: 287, projection_dim: 256)

loss = BarlowTwins.barlow_twins_loss(z_a, z_b, lambda_param: 0.005)

References

• Paper: https://arxiv.org/abs/2103.03230