Coverage for cafl4ds/monitor.py: 100%

1"""The health monitor — the *dependent variable* (the thermometer). 

2 

3The monitor periodically reads the SSL backbone with the Phase-0 instruments 

4(:mod:`cafl4ds.measurements`) on a **fixed** held-out probe set, and returns a flat 

5metric dictionary the loop logs as the ``health`` series. In Phase 0 it is a pure readout; 

6in Phase 3 the same signal drives the monitor→filter controller (the slow edge), so the 

7measurement surface is deliberately kept separate from any action. 

8 

9Metrics reported: 

10 

11* ``rankme`` — effective rank of the probe embeddings (the core collapse readout). 

12* ``cka_drift`` / ``cosine_drift`` — representation drift of the fixed probe set vs. its 

13 first-checkpoint embeddings (content drift and coordinate-frame churn). 

14* ``uniformity`` / ``offdiag_cov`` / ``mean_feature_var`` — extra label-free geometry. 

15* ``knn_acc`` / ``linear_acc`` — downstream probe accuracy on frozen features (labels used 

16 HERE ONLY). 

17 

18Drift is undefined at the first checkpoint (no reference yet), so it is reported as ``0.0`` 

19there and against the stored ``t0`` embeddings thereafter. 

20""" 

21 

22from __future__ import annotations 

23 

24import torch 

25 

26from cafl4ds import measurements 

27from cafl4ds.data.streams import EvalSets 

28from cafl4ds.ssl.base import SSLMethod 

29 

30 

31class HealthMonitor: 

32 """Runs the representation-health instruments on a fixed held-out probe set.""" 

33 

34 def __init__( 

35 self, 

36 eval_sets: EvalSets, 

37 knn_k: int = 20, 

38 run_knn: bool = True, 

39 run_linear: bool = True, 

40 ) -> None: 

41 """Configure the monitor. 

42 

43 Args: 

44 eval_sets: The stream's held-out eval sets (probe support/query, per-era). 

45 knn_k: Number of neighbours for the kNN probe. 

46 run_knn: Whether to compute the kNN probe (labels used HERE ONLY). 

47 run_linear: Whether to compute the linear probe (labels used HERE ONLY). 

48 """ 

49 self.eval_sets = eval_sets 

50 self.knn_k = knn_k 

51 self.run_knn = run_knn 

52 self.run_linear = run_linear 

53 self._z_ref0: torch.Tensor | None = None 

54 

55 def measure(self, method: SSLMethod, step: int) -> dict[str, float]: 

56 """Compute the health metrics for the current model state. 

57 

58 Args: 

59 method: The live SSL method (its encoder supplies the frozen embedding). 

60 step: The current global step (recorded in the returned dict). 

61 

62 Returns: 

63 A flat ``metric -> value`` dictionary (all Python floats). 

64 """ 

65 was_training = method.training 

66 method.eval() 

67 try: 

68 query = self.eval_sets.probe_query 

69 z_query = method.encode(query.images) # [M, d], no grad 

70 metrics: dict[str, float] = { 

71 "step": float(step), 

72 "rankme": measurements.rankme(z_query), 

73 "uniformity": measurements.uniformity(z_query), 

74 "offdiag_cov": measurements.offdiag_covariance(z_query), 

75 "mean_feature_var": float(measurements.feature_variance(z_query).mean().item()), 

76 } 

77 metrics.update(self._drift(z_query)) 

78 if self.run_knn: 

79 metrics["knn_acc"] = measurements.knn_probe( 

80 method.encode, 

81 (self.eval_sets.probe_support.images, self.eval_sets.probe_support.labels), 

82 (query.images, query.labels), 

83 k=self.knn_k, 

84 ) 

85 if self.run_linear: 

86 metrics["linear_acc"] = measurements.linear_probe( 

87 method.encode, 

88 (self.eval_sets.probe_support.images, self.eval_sets.probe_support.labels), 

89 (query.images, query.labels), 

90 ) 

91 return metrics 

92 finally: 

93 method.train(was_training) 

94 

95 def _drift(self, z_query: torch.Tensor) -> dict[str, float]: 

96 """Compute drift of the fixed probe set vs. its first-checkpoint embeddings. 

97 

98 The first call stores the reference embeddings and reports zero drift; later calls 

99 compare against that stored reference. 

100 

101 Args: 

102 z_query: Current embeddings of the fixed probe-query set. 

103 

104 Returns: 

105 ``{"cka_drift": ..., "cosine_drift": ...}``. 

106 """ 

107 if self._z_ref0 is None: 

108 self._z_ref0 = z_query.clone() 

109 return {"cka_drift": 0.0, "cosine_drift": 0.0} 

110 return { 

111 "cka_drift": measurements.cka_drift(self._z_ref0, z_query), 

112 "cosine_drift": measurements.cosine_drift(self._z_ref0, z_query), 

113 }