compute_grouped_similarity_matrices

Compute one similarity matrix per parameter group.

This is the grouped variant of :func:compute_similarity_matrix. It first selects GRADIEND coordinates via part and topk, assigns each decoded base-model parameter name to a group, and then computes one square model-by-model matrix for every group.

Parameters:

Name	Type	Description	Default
models	Dict[str, object]	Mapping from display/model id to a trained model with a .gradiend attribute. At least two models are required.	required
measure	str	Similarity measure. Supported values are "cosine", "cosine_signed", "spearman", and "spearman_signed". Unsigned variants take the absolute value.	'cosine'
part	Optional[str]	GRADIEND part to compare. Supported values are "encoder-weight", "decoder-weight", "decoder-bias", and "decoder-sum". Defaults to "encoder-weight".	None
topk	Optional[Union[int, float]]	Optional coordinate selection. None uses all coordinates, an integer selects that many top coordinates, and a float in (0, 1] selects that fraction using the model's own get_topk_weights.	None
group_by	GroupingSpec	Grouping strategy. "param" keeps decoded parameter names separate, "layer" extracts layer numbers from common transformer names, and "component" uses common transformer component naming conventions. Component grouping is heuristic and may not classify every architecture perfectly. A dict maps exact parameter names to labels, and a callable receives each parameter name and returns a label. Coordinates whose parameter name cannot be decoded are put into "other".	'param'

Returns:

Type	Description
Dict[str, Dict[str, Any]]	A dict keyed by group name. Each value contains measure, group,
Dict[str, Dict[str, Any]]	group_by, model_ids, matrix, part, and topk.

Raises:

Type	Description
TypeError	If models or topk has an invalid type.
ValueError	If fewer than two models are passed, measure is unsupported, part is unsupported, or group_by is invalid.

Source code in gradiend/comparison/similarity.py

def compute_grouped_similarity_matrices(models: Dict[str, object], *, measure: str = "cosine", part: Optional[str] = None, topk: Optional[Union[int, float]] = None, group_by: GroupingSpec = "param") -> Dict[str, Dict[str, Any]]:
    """Compute one similarity matrix per parameter group.

    This is the grouped variant of :func:`compute_similarity_matrix`. It first
    selects GRADIEND coordinates via ``part`` and ``topk``, assigns each decoded
    base-model parameter name to a group, and then computes one square
    model-by-model matrix for every group.

    Args:
        models: Mapping from display/model id to a trained model with a
            ``.gradiend`` attribute. At least two models are required.
        measure: Similarity measure. Supported values are ``"cosine"``,
            ``"cosine_signed"``, ``"spearman"``, and ``"spearman_signed"``.
            Unsigned variants take the absolute value.
        part: GRADIEND part to compare. Supported values are
            ``"encoder-weight"``, ``"decoder-weight"``, ``"decoder-bias"``,
            and ``"decoder-sum"``. Defaults to ``"encoder-weight"``.
        topk: Optional coordinate selection. ``None`` uses all coordinates, an
            integer selects that many top coordinates, and a float in ``(0, 1]``
            selects that fraction using the model's own ``get_topk_weights``.
        group_by: Grouping strategy. ``"param"`` keeps decoded parameter names
            separate, ``"layer"`` extracts layer numbers from common transformer
            names, and ``"component"`` uses common transformer component naming
            conventions. Component grouping is heuristic and may not classify
            every architecture perfectly. A dict maps exact parameter names to
            labels, and a callable receives each parameter name and returns a
            label. Coordinates whose parameter name cannot be decoded are put
            into ``"other"``.

    Returns:
        A dict keyed by group name. Each value contains ``measure``, ``group``,
        ``group_by``, ``model_ids``, ``matrix``, ``part``, and ``topk``.

    Raises:
        TypeError: If ``models`` or ``topk`` has an invalid type.
        ValueError: If fewer than two models are passed, ``measure`` is
            unsupported, ``part`` is unsupported, or ``group_by`` is invalid.
    """
    _validate_models(models)
    _validate_topk_optional(topk)
    measure = (measure or "cosine").lower()
    if measure not in {"cosine", "cosine_signed", "spearman", "spearman_signed"}:
        raise ValueError("Grouped similarities currently support cosine/cosine_signed/spearman/spearman_signed")
    resolved_part = (part or "encoder-weight").lower()
    extracted = {mid: _extract_sparse_blocks(model, part=resolved_part, topk=topk, include_param_names=True) for mid, model in models.items()}
    grouped = _grouped_blocks(extracted, group_by)
    out: Dict[str, Dict[str, Any]] = {}
    model_ids = list(models.keys())
    for group_name, group_blocks in grouped.items():
        matrix = [[0.0] * len(model_ids) for _ in range(len(model_ids))]
        for i, mi in enumerate(model_ids):
            for j, mj in enumerate(model_ids):
                blocks_i = group_blocks.get(mi, {})
                blocks_j = group_blocks.get(mj, {})
                if measure.startswith("cosine"):
                    score = _cosine_from_blocks(blocks_i, blocks_j)
                else:
                    vec_i, vec_j = _vector_from_union(blocks_i, blocks_j)
                    score = 0.0 if not vec_i else _pearson(_rankdata_average(vec_i), _rankdata_average(vec_j))
                matrix[i][j] = abs(score) if measure in {"cosine", "spearman"} else score
        out[group_name] = {
            "measure": measure,
            "group": group_name,
            "group_by": "param" if group_by is None else str(group_by),
            "model_ids": model_ids,
            "matrix": matrix,
            "part": resolved_part,
            "topk": topk,
        }
    return out