diff --git a/simplexity/activations/activation_analyses.py b/simplexity/activations/activation_analyses.py index d192670c..8289b5a6 100644 --- a/simplexity/activations/activation_analyses.py +++ b/simplexity/activations/activation_analyses.py @@ -91,10 +91,18 @@ def __init__( use_probs_as_weights: bool = False, skip_first_token: bool = False, skip_deduplication: bool = True, + rcond: float | None = None, fit_intercept: bool = True, concat_belief_states: bool = False, compute_subspace_orthogonality: bool = False, ) -> None: + analysis_kwargs: dict[str, Any] = { + "fit_intercept": fit_intercept, + "concat_belief_states": concat_belief_states, + "compute_subspace_orthogonality": compute_subspace_orthogonality, + } + if rcond is not None: + analysis_kwargs["rcond"] = float(rcond) super().__init__( analysis_type="linear_regression", last_token_only=last_token_only, @@ -102,11 +110,7 @@ def __init__( use_probs_as_weights=use_probs_as_weights, skip_first_token=skip_first_token, skip_deduplication=skip_deduplication, - analysis_kwargs={ - "fit_intercept": fit_intercept, - "concat_belief_states": concat_belief_states, - "compute_subspace_orthogonality": compute_subspace_orthogonality, - }, + analysis_kwargs=analysis_kwargs, ) diff --git a/simplexity/analysis/layerwise_analysis.py b/simplexity/analysis/layerwise_analysis.py index c2ae0bd5..58ec9a1d 100644 --- a/simplexity/analysis/layerwise_analysis.py +++ b/simplexity/analysis/layerwise_analysis.py @@ -39,7 +39,14 @@ class AnalysisRegistration: def _validate_linear_regression_kwargs(kwargs: Mapping[str, Any] | None) -> dict[str, Any]: provided = dict(kwargs or {}) - allowed = {"fit_intercept", "concat_belief_states", "compute_subspace_orthogonality", "use_svd", "rcond_values"} + allowed = { + "fit_intercept", + "concat_belief_states", + "compute_subspace_orthogonality", + "use_svd", + "rcond", + "rcond_values", + } unexpected = set(provided) - allowed if unexpected: raise ValueError(f"Unexpected linear_regression kwargs: {sorted(unexpected)}") @@ -48,9 +55,12 @@ def _validate_linear_regression_kwargs(kwargs: Mapping[str, Any] | None) -> dict resolved_kwargs["concat_belief_states"] = bool(provided.get("concat_belief_states", False)) resolved_kwargs["compute_subspace_orthogonality"] = bool(provided.get("compute_subspace_orthogonality", False)) rcond_values = provided.get("rcond_values") + rcond = provided.get("rcond") use_svd = bool(provided.get("use_svd", False)) resolved_kwargs["use_svd"] = use_svd if use_svd: + if rcond is not None: + raise ValueError("rcond is only used when use_svd is False; use rcond_values instead") if rcond_values is not None: if not isinstance(rcond_values, (list, tuple)): raise TypeError("rcond_values must be a sequence of floats") @@ -58,8 +68,15 @@ def _validate_linear_regression_kwargs(kwargs: Mapping[str, Any] | None) -> dict raise ValueError("rcond_values must not be empty") rcond_values = tuple(float(v) for v in rcond_values) resolved_kwargs["rcond_values"] = rcond_values - elif rcond_values is not None: - raise ValueError("rcond_values are only used when use_svd is True") + else: + if rcond_values is not None: + raise ValueError("rcond_values are only used when use_svd is True") + if rcond is not None: + if isinstance(rcond, bool) or not isinstance(rcond, (int, float)): + raise TypeError("rcond must be a float or None") + if rcond < 0: + raise ValueError("rcond must be non-negative") + resolved_kwargs["rcond"] = float(rcond) return resolved_kwargs diff --git a/simplexity/analysis/linear_regression.py b/simplexity/analysis/linear_regression.py index db0f7d34..eab53e8b 100644 --- a/simplexity/analysis/linear_regression.py +++ b/simplexity/analysis/linear_regression.py @@ -29,6 +29,60 @@ def _design_matrix(x: jax.Array, fit_intercept: bool) -> jax.Array: return x +def default_rcond(n_columns: int, dtype: Any = jnp.float32) -> float: + """Relative singular-value cutoff used when a caller does not supply one. + + Scaled by the number of design columns, not by ``max(n_rows, n_columns)``. The column count is + the dimension that governs the conditioning of the solve; the row count is not, and using it + makes the cutoff grow with the dataset so that fitting on more data discards more of the design. + """ + return float(n_columns) * float(jnp.finfo(dtype).eps) + + +def _resolve_rcond(rcond: float | None, weighted_design: jax.Array) -> float: + if rcond is None: + return default_rcond(weighted_design.shape[1], weighted_design.dtype) + if rcond < 0: + raise ValueError("rcond must be non-negative") + return float(rcond) + + +def _prepare_problem( + x: jax.Array, + y: jax.Array, + weights: jax.Array | np.ndarray | None, + fit_intercept: bool, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Standardize inputs and build the weighted design shared by every solver here.""" + x_arr = standardize_features(x) + y_arr = standardize_targets(y) + if x_arr.shape[0] != y_arr.shape[0]: + raise ValueError("Features and targets must share the same first dimension") + if x_arr.shape[0] == 0: + raise ValueError("At least one sample is required") + w_arr = normalize_weights(weights, x_arr.shape[0]) + if w_arr is None: + w_arr = jnp.ones(x_arr.shape[0], dtype=x_arr.dtype) / x_arr.shape[0] + design = _design_matrix(x_arr, fit_intercept) + sqrt_w = jnp.sqrt(w_arr)[:, None] + return x_arr, y_arr, w_arr, design, design * sqrt_w, y_arr * sqrt_w + + +def _spectrum_diagnostics(s: jax.Array, threshold: float, rcond: float) -> dict[str, float]: + """Report what the cutoff kept, so a caller can see truncation instead of inferring it.""" + n_values = int(s.size) + effective_rank = int(jnp.sum(s > threshold)) if n_values else 0 + smallest = float(s[-1]) if n_values else 0.0 + largest = float(s[0]) if n_values else 0.0 + condition = largest / smallest if smallest > 0 else float("inf") + return { + "effective_rank": float(effective_rank), + "n_singular_values": float(n_values), + "rcond": float(rcond), + "condition_number": condition, + } + + def _regression_metrics( predictions: jax.Array, targets: jax.Array, @@ -53,49 +107,85 @@ def _regression_metrics( } +def _beta_arrays( + beta: jax.Array, + predictions: jax.Array, + targets: jax.Array, + fit_intercept: bool, +) -> dict[str, jax.Array]: + if fit_intercept: + return {"projected": predictions, "targets": targets, "coeffs": beta[1:], "intercept": beta[:1]} + return {"projected": predictions, "targets": targets, "coeffs": beta} + + def linear_regression( x: jax.Array, y: jax.Array, weights: jax.Array | np.ndarray | None, *, + rcond: float | None = None, fit_intercept: bool = True, ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]: - """Weighted linear regression using a closed-form least squares solution.""" - x_arr = standardize_features(x) - y_arr = standardize_targets(y) - if x_arr.shape[0] != y_arr.shape[0]: - raise ValueError("Features and targets must share the same first dimension") - if x_arr.shape[0] == 0: - raise ValueError("At least one sample is required") - w_arr = normalize_weights(weights, x_arr.shape[0]) - if w_arr is None: - w_arr = jnp.ones(x_arr.shape[0], dtype=x_arr.dtype) / x_arr.shape[0] - design = _design_matrix(x_arr, fit_intercept) - sqrt_w = jnp.sqrt(w_arr)[:, None] - weighted_design = design * sqrt_w - weighted_targets = y_arr * sqrt_w - beta, _, _, _ = jnp.linalg.lstsq(weighted_design, weighted_targets, rcond=None) + """Weighted linear regression using a closed-form least squares solution. + + Singular directions of the weighted design below ``rcond * max_singular_value`` are discarded. + That truncation is regularisation, so it is a parameter rather than a hidden default: pass + ``rcond=0.0`` for an untruncated solve, or a larger value to regularise a near-singular design. + The returned scalars report ``effective_rank`` and ``condition_number`` so that a caller can + see how much of the design was used. + + Args: + x: Features, shape ``(n_samples, n_features)``. + y: Targets, shape ``(n_samples, n_targets)``. + weights: Per-sample weights, or None for uniform weights. + rcond: Relative singular-value cutoff. None uses :func:`default_rcond`. + fit_intercept: Whether to prepend an intercept column to the design. + + Returns: + Scalar metrics (computed in-sample; use :func:`score_regression` for held-out data) and + arrays holding the fitted coefficients and projections. + """ + _, y_arr, w_arr, design, weighted_design, weighted_targets = _prepare_problem(x, y, weights, fit_intercept) + resolved_rcond = _resolve_rcond(rcond, weighted_design) + u, s, vh = jnp.linalg.svd(weighted_design, full_matrices=False) + max_singular = float(s[0]) if s.size else 0.0 + threshold = resolved_rcond * max_singular + beta = _compute_beta_from_svd(u, s, vh, weighted_targets, threshold) predictions = design @ beta - scalars = _regression_metrics(predictions, y_arr, w_arr) + scalars = dict(_regression_metrics(predictions, y_arr, w_arr)) + scalars.update(_spectrum_diagnostics(s, threshold, resolved_rcond)) + return scalars, _beta_arrays(beta, predictions, y_arr, fit_intercept) - if fit_intercept: - arrays = { - "projected": predictions, - "targets": y_arr, - "coeffs": beta[1:], - "intercept": beta[:1], - } - else: - arrays = { - "projected": predictions, - "targets": y_arr, - "coeffs": beta, - } - return scalars, arrays +def score_regression( + x: jax.Array, + y: jax.Array, + weights: jax.Array | np.ndarray | None, + coeffs: jax.Array, + intercept: jax.Array | None = None, +) -> Mapping[str, float]: + """Score fitted coefficients on data they were not fitted to. + :func:`linear_regression` reports metrics in-sample. Pass its ``coeffs`` and ``intercept`` + here with held-out features and targets to obtain an out-of-sample ``r2`` on the same + weighting convention, rather than reimplementing it per caller. -def _compute_regression_metrics( # pylint: disable=too-many-arguments + Args: + x: Held-out features, shape ``(n_samples, n_features)``. + y: Held-out targets, shape ``(n_samples, n_targets)``. + weights: Per-sample weights for the held-out rows, or None for uniform weights. + coeffs: Fitted coefficients, shape ``(n_features, n_targets)``. + intercept: Fitted intercept, shape ``(1, n_targets)``, or None if none was fitted. + + Returns: + The same scalar metrics as a fit, computed on the supplied data. + """ + fit_intercept = intercept is not None + beta = jnp.concatenate([intercept, coeffs], axis=0) if intercept is not None else coeffs + return _compute_regression_metrics(x, y, weights, beta, fit_intercept=fit_intercept) + + +def _compute_regression_metrics( x: jax.Array, y: jax.Array, weights: jax.Array | np.ndarray | None, @@ -104,20 +194,10 @@ def _compute_regression_metrics( # pylint: disable=too-many-arguments *, fit_intercept: bool = True, ) -> Mapping[str, float]: - x_arr = standardize_features(x) - y_arr = standardize_targets(y) - if x_arr.shape[0] != y_arr.shape[0]: - raise ValueError("Features and targets must share the same first dimension") - if x_arr.shape[0] == 0: - raise ValueError("At least one sample is required") - w_arr = normalize_weights(weights, x_arr.shape[0]) - if w_arr is None: - w_arr = jnp.ones(x_arr.shape[0], dtype=x_arr.dtype) / x_arr.shape[0] + _, y_arr, w_arr, design, _, _ = _prepare_problem(x, y, weights, fit_intercept) if predictions is None: - design = _design_matrix(x_arr, fit_intercept) predictions = design @ beta - scalars = _regression_metrics(predictions, y_arr, w_arr) - return scalars + return _regression_metrics(predictions, y_arr, w_arr) def _compute_beta_from_svd( @@ -141,20 +221,25 @@ def linear_regression_svd( rcond_values: Sequence[float] | None = None, fit_intercept: bool = True, ) -> tuple[Mapping[str, float], Mapping[str, jax.Array]]: - """Weighted linear regression solved via SVD with configurable rcond search.""" - x_arr = standardize_features(x) - y_arr = standardize_targets(y) - if x_arr.shape[0] != y_arr.shape[0]: - raise ValueError("Features and targets must share the same first dimension") - if x_arr.shape[0] == 0: - raise ValueError("At least one sample is required") - w_arr = normalize_weights(weights, x_arr.shape[0]) - if w_arr is None: - w_arr = jnp.ones(x_arr.shape[0], dtype=x_arr.dtype) / x_arr.shape[0] - design = _design_matrix(x_arr, fit_intercept) - sqrt_w = jnp.sqrt(w_arr)[:, None] - weighted_design = design * sqrt_w - weighted_targets = y_arr * sqrt_w + """Weighted linear regression solved via SVD with a search over candidate rcond values. + + The candidate chosen is the one with the lowest error **on the data it was fitted to**, which + cannot identify a value that regularises: a cutoff only helps out-of-sample. Select a cutoff + on held-out data with :func:`linear_regression` and :func:`score_regression` instead when the + design is near-singular. + + Args: + x: Features, shape ``(n_samples, n_features)``. + y: Targets, shape ``(n_samples, n_targets)``. + weights: Per-sample weights, or None for uniform weights. + rcond_values: Candidate relative cutoffs. None uses a single ``1e-15``. + fit_intercept: Whether to prepend an intercept column to the design. + + Returns: + Scalar metrics for the selected candidate, including ``best_rcond`` and the spectrum + diagnostics, and arrays holding its coefficients and projections. + """ + _, y_arr, w_arr, design, weighted_design, weighted_targets = _prepare_problem(x, y, weights, fit_intercept) u, s, vh = jnp.linalg.svd(weighted_design, full_matrices=False) max_singular = float(s[0]) if s.size else 0.0 rconds = tuple(rcond_values) if rcond_values else (1e-15,) @@ -182,22 +267,9 @@ def linear_regression_svd( assert best_beta is not None scalars = dict(best_scalars) scalars["best_rcond"] = float(best_rcond) + scalars.update(_spectrum_diagnostics(s, best_rcond * max_singular, float(best_rcond))) - if fit_intercept: - arrays = { - "projected": best_pred, - "targets": y_arr, - "coeffs": best_beta[1:], - "intercept": best_beta[:1], - } - else: - arrays = { - "projected": best_pred, - "targets": y_arr, - "coeffs": best_beta, - } - - return scalars, arrays + return scalars, _beta_arrays(best_beta, best_pred, y_arr, fit_intercept) def _process_individual_factors( @@ -229,6 +301,9 @@ def _merge_results_with_suffix( arrays.update({f"{key}/{suffix}": value for key, value in results_arrays.items()}) +_METRICS_KWARGS = frozenset({"fit_intercept"}) + + def _split_concat_results( layer_activations: jax.Array, weights: jax.Array, @@ -253,9 +328,7 @@ def _split_concat_results( else: intercepts_list = [None] * len(belief_states) - # Only recompute scalar metrics, reuse projections and coefficients - # Filter out rcond_values from kwargs (only relevant for SVD during fitting, not metrics) - metrics_kwargs = {k: v for k, v in kwargs.items() if k != "rcond_values"} + metrics_kwargs = {k: v for k, v in kwargs.items() if k in _METRICS_KWARGS} results = [] for factor, coeffs, intercept, projections, targets in zip( diff --git a/tests/analysis/test_layerwise_analysis.py b/tests/analysis/test_layerwise_analysis.py index 9124b015..37d3425a 100644 --- a/tests/analysis/test_layerwise_analysis.py +++ b/tests/analysis/test_layerwise_analysis.py @@ -291,3 +291,32 @@ def test_linear_regression_use_svd_defaults_false() -> None: params = validator({}) assert params["use_svd"] is False + + +def test_linear_regression_accepts_rcond() -> None: + """linear_regression validator should pass rcond through to the solver.""" + from simplexity.analysis.layerwise_analysis import ANALYSIS_REGISTRY + + validator = ANALYSIS_REGISTRY["linear_regression"].validator + assert validator({"rcond": 1e-6})["rcond"] == pytest.approx(1e-6) + assert "rcond" not in validator({}) + + +def test_linear_regression_rejects_invalid_rcond() -> None: + """A cutoff a caller cannot have meant must be refused, not silently coerced.""" + from simplexity.analysis.layerwise_analysis import ANALYSIS_REGISTRY + + validator = ANALYSIS_REGISTRY["linear_regression"].validator + with pytest.raises(ValueError, match="rcond must be non-negative"): + validator({"rcond": -1.0}) + with pytest.raises(TypeError, match="rcond must be a float or None"): + validator({"rcond": "small"}) + + +def test_linear_regression_svd_rejects_scalar_rcond() -> None: + """The SVD path searches rcond_values; a scalar rcond there would be silently ignored.""" + from simplexity.analysis.layerwise_analysis import ANALYSIS_REGISTRY + + validator = ANALYSIS_REGISTRY["linear_regression_svd"].validator + with pytest.raises(ValueError, match="rcond is only used when use_svd is False"): + validator({"rcond": 1e-6}) diff --git a/tests/analysis/test_linear_regression.py b/tests/analysis/test_linear_regression.py index 52c3b29a..f97112a7 100644 --- a/tests/analysis/test_linear_regression.py +++ b/tests/analysis/test_linear_regression.py @@ -13,13 +13,16 @@ import chex import jax import jax.numpy as jnp +import numpy as np import pytest from simplexity.analysis.linear_regression import ( + default_rcond, get_robust_basis, layer_linear_regression, linear_regression, linear_regression_svd, + score_regression, ) @@ -1211,3 +1214,114 @@ def test_get_robust_basis_single_vector(): expected_norm = jnp.linalg.norm(vector) parallel_tol = dim * eps * expected_norm assert jnp.allclose(jnp.abs(basis.T @ vector), expected_norm, atol=parallel_tol) + + +def _tall_well_conditioned_design( + n_rows: int, + n_features: int = 8, + condition: float = 1e4, + seed: int = 0, +) -> tuple[jax.Array, jax.Array]: + """A tall design whose condition number needs no truncation at all. + + ``max(n_rows, n_columns) * eps`` exceeds ``1 / condition`` once ``n_rows`` is large, so a + row-scaled cutoff discards well-conditioned directions purely because the dataset is big. + """ + rng = np.random.default_rng(seed) + scales = np.logspace(0.0, -np.log10(condition), n_features) + rotation, _ = np.linalg.qr(rng.standard_normal((n_features, n_features))) + x = rng.standard_normal((n_rows, n_features)) @ ((rotation * scales) @ rotation.T) + coeffs = rotation @ ((scales**-1.0)[:, None] * rng.standard_normal((n_features, 3))) + y = x @ coeffs + return jnp.asarray(x, dtype=jnp.float32), jnp.asarray(y / y.std(), dtype=jnp.float32) + + +def _float64_reference(x: jax.Array, y: jax.Array) -> np.ndarray: + design = np.concatenate([np.ones((x.shape[0], 1)), np.asarray(x, dtype=np.float64)], axis=1) + beta, *_ = np.linalg.lstsq(design, np.asarray(y, dtype=np.float64), rcond=None) + return beta + + +def test_linear_regression_does_not_truncate_a_well_conditioned_design() -> None: + """Regression for #201: the default cutoff must not discard usable directions. + + Fails on a row-scaled cutoff, which at 20000 rows zeroes every direction below 2.4e-3 of the + largest and so recovers coefficients far from the float64 solution. + """ + x, y = _tall_well_conditioned_design(n_rows=20000) + reference = _float64_reference(x, y) + + scalars, arrays = linear_regression(x, y, None) + beta = np.concatenate([np.asarray(arrays["intercept"]), np.asarray(arrays["coeffs"])], axis=0) + + assert scalars["condition_number"] < 1e5 + assert scalars["effective_rank"] == scalars["n_singular_values"] + chex.assert_trees_all_close(beta, reference.astype(np.float32), rtol=1e-2, atol=1e-2) + + +def test_linear_regression_cutoff_is_independent_of_row_count() -> None: + """Regression for #201: fitting on more rows must not discard more of the design.""" + ranks = [] + for n_rows in (2000, 20000, 60000): + x, y = _tall_well_conditioned_design(n_rows=n_rows) + scalars, _ = linear_regression(x, y, None) + ranks.append((scalars["effective_rank"], scalars["n_singular_values"])) + + assert len({rank for rank, _ in ranks}) == 1, f"effective rank varied with row count: {ranks}" + assert all(rank == total for rank, total in ranks) + + +def test_linear_regression_rcond_is_honoured() -> None: + """An explicit cutoff must take effect, so a caller can choose the regularisation.""" + x, y = _tall_well_conditioned_design(n_rows=20000) + + untruncated, _ = linear_regression(x, y, None, rcond=0.0) + truncated, _ = linear_regression(x, y, None, rcond=1e-1) + + assert untruncated["effective_rank"] == untruncated["n_singular_values"] + assert truncated["effective_rank"] < untruncated["effective_rank"] + assert truncated["r2"] < untruncated["r2"] + + +def test_linear_regression_rejects_negative_rcond() -> None: + x, y = _tall_well_conditioned_design(n_rows=200) + with pytest.raises(ValueError, match="rcond must be non-negative"): + linear_regression(x, y, None, rcond=-1.0) + + +def test_default_rcond_scales_with_columns_not_rows() -> None: + eps = float(jnp.finfo(jnp.float32).eps) + assert default_rcond(9) == pytest.approx(9 * eps) + assert default_rcond(9) == default_rcond(9) + + +def test_score_regression_evaluates_held_out_data() -> None: + """The scorer must read the data it is given, not the data the fit saw.""" + x, y = _tall_well_conditioned_design(n_rows=4000) + x_fit, y_fit, x_held, y_held = x[:3000], y[:3000], x[3000:], y[3000:] + + scalars, arrays = linear_regression(x_fit, y_fit, None, rcond=0.0) + held = score_regression(x_held, y_held, None, arrays["coeffs"], arrays["intercept"]) + in_sample = score_regression(x_fit, y_fit, None, arrays["coeffs"], arrays["intercept"]) + + assert in_sample["r2"] == pytest.approx(scalars["r2"], abs=1e-4) + assert held["r2"] > 0.9 + assert held["r2"] != in_sample["r2"] + + +def test_score_regression_without_intercept() -> None: + x, y = _tall_well_conditioned_design(n_rows=2000) + _, arrays = linear_regression(x, y, None, rcond=0.0, fit_intercept=False) + assert "intercept" not in arrays + + held = score_regression(x, y, None, arrays["coeffs"]) + assert held["r2"] > 0.9 + + +def test_linear_regression_svd_reports_spectrum_diagnostics() -> None: + x, y = _tall_well_conditioned_design(n_rows=2000) + scalars, _ = linear_regression_svd(x, y, None, rcond_values=[1e-10, 1e-1]) + + assert scalars["rcond"] == scalars["best_rcond"] + assert scalars["effective_rank"] <= scalars["n_singular_values"] + assert scalars["condition_number"] > 1.0