Coverage for cafl4ds/ssl/simsiam.py: 100%

1"""SimSiam SSL method — the joint-embedding ``C`` backbone (Chen & He 2021). 

2 

3SimSiam is the simplest of the joint-embedding family: a projector + predictor with a 

4**stop-gradient**, and *no* negatives and *no* momentum/EMA target (that is BYOL, a later 

5drop-in). Its learning signal is pulling together two augmented views, which makes it 

6susceptible to representational **collapse** — the degradation mode the Phase-1 positive 

7control will elicit and the health instruments (rank, alignment/uniformity) will catch. The 

8health monitor reads the *encoder* embedding, never the projector/predictor outputs. 

9""" 

10 

11from __future__ import annotations 

12 

13import torch 

14from torchvision.transforms import v2 

15 

16from cafl4ds.models.heads import MLPHead 

17from cafl4ds.models.vit import TinyViTEncoder 

18from cafl4ds.ssl.augment import TwoView, make_ssl_augment 

19from cafl4ds.ssl.base import SSLMethod 

20 

21 

22def _neg_cosine(p: torch.Tensor, z: torch.Tensor, stop_grad: bool = True) -> torch.Tensor: 

23 """Negative cosine similarity, with the target branch optionally stop-gradient'd. 

24 

25 Args: 

26 p: Predictor output of one branch ``[B, d]`` (gradient flows). 

27 z: Projector output of the other branch ``[B, d]``. 

28 stop_grad: When ``True`` (SimSiam as published) the target ``z`` is detached — the 

29 stop-gradient that prevents collapse. When ``False`` the gradient flows through 

30 both branches, removing that anti-collapse mechanism (the positive-control 

31 ablation, see :class:`SimSiam`). 

32 

33 Returns: 

34 The mean negative cosine similarity (scalar); minimized when ``p`` aligns with ``z``. 

35 """ 

36 return _neg_cosine_per_sample(p, z, stop_grad=stop_grad).mean() 

37 

38 

39def _neg_cosine_per_sample(p: torch.Tensor, z: torch.Tensor, stop_grad: bool = True) -> torch.Tensor: 

40 """Per-sample negative cosine similarity ``[B]`` (the un-reduced :func:`_neg_cosine`). 

41 

42 Args: 

43 p: Predictor output of one branch ``[B, d]``. 

44 z: Projector output of the other branch ``[B, d]``. 

45 stop_grad: Whether to detach the target ``z`` (see :func:`_neg_cosine`). 

46 

47 Returns: 

48 The per-sample negative cosine similarity ``[B]`` (lower is better). 

49 """ 

50 p = torch.nn.functional.normalize(p, dim=1) 

51 z = torch.nn.functional.normalize(z.detach() if stop_grad else z, dim=1) 

52 return -(p * z).sum(dim=1) 

53 

54 

55class SimSiam(SSLMethod): 

56 """SimSiam over the shared :class:`~cafl4ds.models.vit.TinyViTEncoder`. 

57 

58 SimSiam avoids collapse with two coupled mechanisms: a **predictor** on one branch and a 

59 **stop-gradient** on the target branch (Chen & He 2021). The ``anti_collapse`` flag toggles 

60 *both* off together, giving the documented collapse ablation used as the Phase-0 positive 

61 control (P0.2): with the predictor bypassed (``p = z``) and no stop-gradient, the objective 

62 reduces to ``-cosine(z1, z2)`` with gradients flowing through both branches, whose trivial 

63 global optimum maps every input to one constant vector (cosine → 1, loss → −1). Collapse is 

64 then mathematically forced and scale-independent — which is exactly why the toy CPU/HPU 

65 regime suffices to calibrate the collapse instruments (RankMe, alignment/uniformity). 

66 """ 

67 

68 def __init__( 

69 self, 

70 encoder: TinyViTEncoder, 

71 projector: MLPHead, 

72 predictor: MLPHead, 

73 augment: v2.Transform | None = None, 

74 anti_collapse: bool = True, 

75 ) -> None: 

76 """Build the SimSiam method. 

77 

78 Args: 

79 encoder: The shared backbone encoder. 

80 projector: The projection MLP (3-layer, last BatchNorm) mapping the pooled 

81 embedding to the latent space. 

82 predictor: The prediction MLP (2-layer, no last BatchNorm) applied to one branch. 

83 augment: Per-view augmentation; defaults to 

84 :func:`~cafl4ds.ssl.augment.make_ssl_augment` sized to the encoder. 

85 anti_collapse: When ``True`` (default) run SimSiam as published — predictor + stop- 

86 gradient. When ``False`` disable **both** anti-collapse mechanisms (predictor 

87 bypassed, stop-gradient off): the forced-collapse positive control (P0.2). 

88 """ 

89 super().__init__(encoder) 

90 self.projector = projector 

91 self.predictor = predictor 

92 self.anti_collapse = anti_collapse 

93 img_size = int(round(encoder.num_patches**0.5)) * encoder.patch_size 

94 base_augment = augment if augment is not None else make_ssl_augment(img_size) 

95 self.two_view = TwoView(base_augment) 

96 

97 @property 

98 def name(self) -> str: 

99 """Return the method identifier (``"simsiam_collapse"`` for the PC ablation).""" 

100 return "simsiam" if self.anti_collapse else "simsiam_collapse" 

101 

102 def _branch(self, view: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 

103 """Encode → project (→ predict, unless collapsing) one augmented view. 

104 

105 Args: 

106 view: An augmented image batch ``[B, C, H, W]``. 

107 

108 Returns: 

109 A ``(prediction, projection)`` pair, each ``[B, latent_dim]``. With 

110 ``anti_collapse=False`` the predictor is removed from the path, so ``p = z``. 

111 """ 

112 z = self.projector(self.encoder.embed(view)) 

113 p = self.predictor(z) if self.anti_collapse else z 

114 return p, z 

115 

116 def training_step(self, imgs: torch.Tensor) -> torch.Tensor: 

117 """Compute the symmetric negative-cosine loss on two views. 

118 

119 With ``anti_collapse=True`` this is the published stop-gradient objective; with it 

120 ``False`` the stop-gradient is dropped (and the predictor already bypassed in 

121 :meth:`_branch`), yielding the forced-collapse positive control. 

122 

123 Args: 

124 imgs: A batch of raw images ``[B, C, H, W]``. Labels never enter this path. 

125 

126 Returns: 

127 The symmetrized SimSiam loss (scalar), in ``[-1, 1]`` (lower is better). 

128 """ 

129 view_1, view_2 = self.two_view(imgs) 

130 p1, z1 = self._branch(view_1) 

131 p2, z2 = self._branch(view_2) 

132 sg = self.anti_collapse 

133 return 0.5 * (_neg_cosine(p1, z2, stop_grad=sg) + _neg_cosine(p2, z1, stop_grad=sg)) 

134 

135 def per_sample_loss(self, imgs: torch.Tensor) -> torch.Tensor: 

136 """Per-image symmetric negative-cosine loss ``[B]`` (no gradient). 

137 

138 The per-frame version of :meth:`training_step`, used as an informativeness signal by 

139 the loss-gate knob. Two fresh views are drawn (as in a training step) and the symmetric 

140 negative cosine is kept per image instead of averaged over the batch. 

141 

142 Args: 

143 imgs: A batch of raw images ``[B, C, H, W]``. 

144 

145 Returns: 

146 A detached ``[B]`` tensor of per-image losses in ``[-1, 1]`` (lower is better). 

147 """ 

148 with torch.no_grad(): 

149 view_1, view_2 = self.two_view(imgs) 

150 p1, z1 = self._branch(view_1) 

151 p2, z2 = self._branch(view_2) 

152 sg = self.anti_collapse 

153 per_sample = 0.5 * ( 

154 _neg_cosine_per_sample(p1, z2, stop_grad=sg) + _neg_cosine_per_sample(p2, z1, stop_grad=sg) 

155 ) 

156 return per_sample