Skip to content

SufficientStats API

SufficientStats is the accumulator the GaussianMixture M-step reads from (nk, s1, s2, n_samples); see the sufficient-statistics design doc for the derivation. Most users won't construct one directly -- it's public mainly so later mixture families and streaming/mini-batch fitting can share the same accumulator type.

SufficientStats dataclass

Accumulated sufficient statistics for one (chunk of a) mixture E-step, per component :math:k.

Attributes:

Name Type Description
nk (Tensor, shape(n_components))

:math:n_k = \sum_i r_{ik}.
s1 (Tensor, shape(n_components, n_features))

:math:\sum_i r_{ik} x_i.
s2 Tensor

Covariance-type-specific second moment, centered on each component's empirical mean :math:\bar y_k = s1_k / n_k for this accumulation (not on whatever the model's current means_ holds): full -> (n_components, n_features, n_features) = :math:\sum_i r_{ik} (x_i - \bar y_k)(x_i - \bar y_k)^T; diag -> (n_components, n_features); spherical -> (n_components,). Tied covariance types use the same shape as their non-tied base type -- the pooling across components happens in the M-step, not here.
n_samples int

Number of rows this instance summarizes.
aux (dict[str, Tensor], optional)

Family-specific accumulators that don't fit the Gaussian-shaped nk/s1/s2 fields above -- e.g. StudentTMixture's :math:u-weighted sums. None (default) for every family that doesn't need it, in particular GaussianMixture, whose behavior is completely unaffected by this field's existence. __add__/ scale merge/scale it key-by-key when present on either side.
Source code in tgmm/stats.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@dataclass
class SufficientStats:
    r"""
    Accumulated sufficient statistics for one (chunk of a) mixture E-step,
    per component :math:`k`.

    Attributes
    ----------
    nk : torch.Tensor, shape (n_components,)
        :math:`n_k = \sum_i r_{ik}`.
    s1 : torch.Tensor, shape (n_components, n_features)
        :math:`\sum_i r_{ik} x_i`.
    s2 : torch.Tensor
        Covariance-type-specific second moment, centered on each
        component's empirical mean :math:`\bar y_k = s1_k / n_k` for
        *this* accumulation (not on whatever the model's current
        ``means_`` holds): ``full`` -> ``(n_components, n_features,
        n_features)`` = :math:`\sum_i r_{ik} (x_i - \bar y_k)(x_i - \bar
        y_k)^T`; ``diag`` -> ``(n_components, n_features)``; ``spherical``
        -> ``(n_components,)``. Tied covariance types use the same shape as
        their non-tied base type -- the pooling across components happens
        in the M-step, not here.
    n_samples : int
        Number of rows this instance summarizes.
    aux : dict[str, torch.Tensor], optional
        Family-specific accumulators that don't fit the Gaussian-shaped
        ``nk``/``s1``/``s2`` fields above -- e.g. `StudentTMixture`'s
        :math:`u`-weighted sums. ``None`` (default) for every family that
        doesn't need it, in particular `GaussianMixture`, whose behavior is
        completely unaffected by this field's existence. ``__add__``/
        ``scale`` merge/scale it key-by-key when present on either side.
    """

    nk: torch.Tensor
    s1: torch.Tensor
    s2: torch.Tensor
    n_samples: int
    aux: Optional[Dict[str, torch.Tensor]] = None

    def __add__(self, other: "SufficientStats") -> "SufficientStats":
        if not isinstance(other, SufficientStats):
            return NotImplemented
        return SufficientStats(
            nk=self.nk + other.nk,
            s1=self.s1 + other.s1,
            s2=self.s2 + other.s2,
            n_samples=self.n_samples + other.n_samples,
            aux=self._merge_aux(other),
        )

    def scale(self, factor: float) -> "SufficientStats":
        r"""
        Scale ``nk``/``s1``/``s2``/``aux`` by ``factor``; ``n_samples`` is
        left untouched since it tracks how many physical rows contributed,
        not a quantity meaningful to reweight. Used by online/mini-batch EM
        (Phase 2) to blend a running accumulator with a new batch's
        statistics.
        """
        return SufficientStats(
            nk=self.nk * factor,
            s1=self.s1 * factor,
            s2=self.s2 * factor,
            n_samples=self.n_samples,
            aux=None if self.aux is None else {k: v * factor for k, v in self.aux.items()},
        )

    def _merge_aux(self, other: "SufficientStats") -> Optional[Dict[str, torch.Tensor]]:
        if self.aux is None and other.aux is None:
            return None
        self_aux = self.aux or {}
        other_aux = other.aux or {}
        keys = set(self_aux) | set(other_aux)
        merged = {}
        for key in keys:
            if key in self_aux and key in other_aux:
                merged[key] = self_aux[key] + other_aux[key]
            else:
                merged[key] = self_aux.get(key, other_aux.get(key))
        return merged

Methods:

scale(factor)

Scale nk/s1/s2/aux by factor; n_samples is left untouched since it tracks how many physical rows contributed, not a quantity meaningful to reweight. Used by online/mini-batch EM (Phase 2) to blend a running accumulator with a new batch's statistics.

Source code in tgmm/stats.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def scale(self, factor: float) -> "SufficientStats":
    r"""
    Scale ``nk``/``s1``/``s2``/``aux`` by ``factor``; ``n_samples`` is
    left untouched since it tracks how many physical rows contributed,
    not a quantity meaningful to reweight. Used by online/mini-batch EM
    (Phase 2) to blend a running accumulator with a new batch's
    statistics.
    """
    return SufficientStats(
        nk=self.nk * factor,
        s1=self.s1 * factor,
        s2=self.s2 * factor,
        n_samples=self.n_samples,
        aux=None if self.aux is None else {k: v * factor for k, v in self.aux.items()},
    )