physioex.explain.foundational#

Explainability methods for foundation model embeddings.

Available methods:

  • CSD (Conceptual Spectral Decomposition): time-frequency attribution of class-specific embedding dimensions via masked SpectralGradients.

class physioex.explain.foundational.CSDResult(concepts, class_attribution, mask, selected_dims, band_frequencies, target_class, specificity_name, probe_weights, contributions)[source]#

Bases: object

Result of a CSD explanation.

Parameters:
top_band_hz(dim)[source]#

Return the center frequency of the top band for a concept.

Parameters:

dim (int)

Return type:

float

class physioex.explain.foundational.CohenDSpecificity(tau=1.0)[source]#

Bases: SpecificityStrategy

Static specificity via Cohen’s d effect size.

For each dimension d, measures the standardized difference between the class-conditional mean and the out-of-class mean:

d_cohen(d, c) = (μ_d|c - μ_d|¬c) / σ_pooled
mask_d(c) = sigmoid(d_cohen / tau)

Requires dataset-level embeddings and labels (computed once, reusable).

Parameters:

tau (float)

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W – (n_classes, D) linear probe weights.

  • embeddings – (N, D) dataset embeddings (for static strategies).

  • labels – (N,) class labels (for static strategies).

  • x – (B, D) current input embeddings (for dynamic strategies).

  • target_class – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor

class physioex.explain.foundational.ConceptualSpectralDecomposition(model, probe_weights, fs, freq_step=4.0, bands=None, steps=10, path='marginal', specificity=None, mask_threshold=0.5, skip_ig=False, n_perms=20, device=None)[source]#

Bases: Module

CSD: per-concept time-frequency explanation of foundation model embeddings.

For each class-specific embedding dimension (selected by the specificity strategy), runs SpectralGradients to produce an independent (n_bands, T) attribution map. The class-level map is the weighted sum of all concept maps.

Parameters:

  • model (nn.Module) – FoundationEncoder (pure encoder without head).

  • probe_weights (dict) – Dictionary containing ‘ln’ (LayerNorm) and ‘W’ (Linear weights) from the trained linear probe.

  • fs (float) – Sampling rate of the preprocessed signal (Hz).

  • freq_step (float) – SpectralGradients band width (Hz). Default 4.0.

  • bands (Optional[List[FrequencyBand]]) – Custom frequency bands. If None, uniform bands of width freq_step.

  • steps (int) – Integration steps per local IG. Default 10. Ignored if skip_ig=True.

  • path (str) – SG ablation direction (‘marginal’, ‘cumulative_add’, ‘cumulative_remove’, ‘both_cumulative’, ‘shapley’). Default ‘marginal’.

  • specificity (Optional[SpecificityStrategy]) – Strategy for selecting class-specific dimensions. Default: MarginSpecificity(tau=0.5).

  • mask_threshold (float) – Dimensions with mask > threshold get their own map. Default 0.5.

  • skip_ig (bool) – If True, skip IG computation and return fast band importance only. Returns (B, n_bands) instead of (B, n_bands, C, T). Default False.

  • n_perms (int) – Number of permutations for Shapley path. Default 20.

  • device (Optional[str]) – Torch device. Default: auto.

explain(x, target_class, embeddings=None, labels=None, max_concepts=None)[source]#

Compute per-concept CSD explanation.

Parameters:

  • x (Tensor) – (B, C, T) or (B, L, C, T) raw signal.

  • target_class (int) – Class index to explain.

  • embeddings (Tensor | None) – (N, D) dataset embeddings for static strategies.

  • labels (Tensor | None) – (N,) labels for static strategies.

  • max_concepts (int | None) – Limit number of concepts to explain (None = all above threshold). Selects top by |mask * W[c,d]|.

Returns:

CSDResult with per-concept maps and aggregated class map.

Return type:

CSDResult

class physioex.explain.foundational.MarginSpecificity(tau=1.0)[source]#

Bases: SpecificityStrategy

Dynamic per-input specificity via class margin.

For each dimension d and input x, measures how much the contribution to the target class exceeds the best competing class:

margin_d(x, c) = W[c,d]·ĥ_d(x) - max_{c'≠c} W[c',d]·ĥ_d(x)
mask_d(x, c) = sigmoid(margin / tau).detach()

Captures input-specific discriminative dimensions (e.g., K-complex present in this particular N2 epoch).

Parameters:

tau (float)

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W – (n_classes, D) linear probe weights.

  • embeddings – (N, D) dataset embeddings (for static strategies).

  • labels – (N,) class labels (for static strategies).

  • x – (B, D) current input embeddings (for dynamic strategies).

  • target_class – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor

class physioex.explain.foundational.NoFilter[source]#

Bases: SpecificityStrategy

No filtering — all dimensions contribute equally.

Baseline strategy: SpectralGradients on the raw class logit without any specificity masking.

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W – (n_classes, D) linear probe weights.

  • embeddings – (N, D) dataset embeddings (for static strategies).

  • labels – (N,) class labels (for static strategies).

  • x – (B, D) current input embeddings (for dynamic strategies).

  • target_class – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor

class physioex.explain.foundational.SoftmaxSpecificity(tau=1.0)[source]#

Bases: SpecificityStrategy

Dynamic per-input specificity via softmax normalization.

For each dimension d, computes the softmax over class contributions:

v_d(x) = [W[0,d]·ĥ_d(x), ..., W[C-1,d]·ĥ_d(x)]
mask_d(x, c) = softmax(v_d / tau)[c].detach()

mask ≈ 1/C for generic dimensions, ≈ 1 for class-specific ones.

Parameters:

tau (float)

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W – (n_classes, D) linear probe weights.

  • embeddings – (N, D) dataset embeddings (for static strategies).

  • labels – (N,) class labels (for static strategies).

  • x – (B, D) current input embeddings (for dynamic strategies).

  • target_class – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor

class physioex.explain.foundational.SpecificityStrategy[source]#

Bases: object

Base class for specificity strategies.

Subclass and implement compute_mask() to define custom strategies.

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W (Tensor) – (n_classes, D) linear probe weights.

  • embeddings (Tensor | None) – (N, D) dataset embeddings (for static strategies).

  • labels (Tensor | None) – (N,) class labels (for static strategies).

  • x (Tensor | None) – (B, D) current input embeddings (for dynamic strategies).

  • target_class (int | None) – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor

class physioex.explain.foundational.TopKSpecificity(k=10)[source]#

Bases: SpecificityStrategy

Hard selection of top-K dimensions by absolute contribution.

Selects the K dimensions with highest |W[c,d] · ĥ_d(x)| for the target class. All other dimensions are zeroed out.

Parameters:

k (int)

compute_mask(W, embeddings=None, labels=None, x=None, target_class=None)[source]#

Compute specificity mask.

Parameters:

  • W – (n_classes, D) linear probe weights.

  • embeddings – (N, D) dataset embeddings (for static strategies).

  • labels – (N,) class labels (for static strategies).

  • x – (B, D) current input embeddings (for dynamic strategies).

  • target_class – class index to explain.

Returns:

(D,) for static strategies or (B, D) for dynamic strategies. Values in [0, 1].

Return type:

Tensor