[Proposal] Representation Geometry: causal inner products, concept directions, and categorical diagnostics #1880

Proposal

Add a maintained representation-geometry utility for TransformerBridge that derives a metric from centered unembedding covariance, constructs concept directions from counterfactual token pairs, and reports categorical geometry in explicitly declared measurement or intervention spaces.

Deliver this issue through two complementary outcomes:

1. A usable geometry/concept-analysis feature with tests, public documentation, and a download-free demo.
2. Independent research reproduction and controlled behavioral validation of metric-derived interventions.

The core implementation exists locally but is not merged. Independent empirical validation remains outstanding. This issue should remain open after the core feature merges and close only after both outcomes are delivered.

Motivation

TransformerLens exposes unembedding weights and internal activations, but comparing concept directions requires an explicit metric and coordinate-space contract. Measurement directions describe how a concept is read; intervention vectors describe how it is written. These spaces transform dually, so applying the same whitening map to both breaks their pairing.

A reusable geometry primitive would make contrast-based concept analysis less ad hoc and provide a common foundation for geometry-aware probing and steering experiments. It complements existing projection-kernel subspace analysis rather than duplicating it: issue #1720 and PR #1721 concern attention-head subspace overlap, not unembedding-covariance metrics.

The method is motivated by Park, Choe and Veitch, arXiv:2311.03658v2. Its causal interpretation depends on linear-representation and spanning conditions, plus Assumption 3.3 concerning uncorrelated concept measurements under uniform vocabulary sampling. Theorem 3.4 characterizes a family of metrics; the experiments choose D = I, giving the covariance-inverse metric. This is not a uniquely identified inner product or an unconditional property of every model.

Geometry alone does not establish causal separability, semantic categories, or causal model use. Raw cosine remains a useful comparison baseline, not a universally invalid measurement.

Pitch

A clear and concise description of what you want to happen. It is useful if you relate proposals to existing features or lackthereof as well as relevant mechanistic interpretability techniques or model architectures.

Geometry and concept-analysis API

Provide the following public types and operations:

• RepresentationGeometry: tensor fitting and non-mutating from_bridge(...) construction.
• Separate measurement/intervention transforms, inverse maps, inner products, and cosines.
• ConceptDirection: oriented pair means, raw/whitened differences, population dispersion, and explicitly labeled metric-derived interventions.
• CategoricalGeometry: centering, affine rank, Gram matrices, distances, angles, validity masks, and separate simplex/regular-simplex checks over supplied vertices.
• GeometryBasis: coordinate provenance for snapshots and their results.

For U = W_U with shape [d_model, d_vocab], estimate centered, uniformly weighted population covariance:

mean_g = mean_v U[:, v]
Sigma = (1 / V) sum_v (U[:, v] - mean_g)(U[:, v] - mean_g)^T

For an exact positive-definite fit, set M = Sigma. With explicit positive ridge, set M = Sigma + epsilon I. For column-vector measurement g and intervention/context vector h:

W = M^(-1/2)
g' = W g
h' = W^(-T) h
h'^T g' = h^T g

Measurement inner products use M^(-1) and intervention inner products use M. Exact whitening satisfies W Sigma W^T = I; with ridge, the identity applies to Sigma + epsilon I, not generally the original covariance. Reject singular/near-singular exact fits and insufficient explicit ridge rather than silently pseudo-inverting or clamping eigenvalues. Report population, precision, rank threshold, conditioning, and regularization.

Pair (lo, hi) contributes U[:, hi] - U[:, lo]. Return an unnormalized, equally weighted mean and population dispersion. A single pair has zero dispersion, not a confidence guarantee. Reject invalid, duplicate, reversed-duplicate, zero, or non-finite contrasts rather than silently filtering them.

String endpoints must encode to exactly one known token without implicit BOS/EOS, preserving leading-space spelling. IDs need no tokenizer. Snapshot the readout and optional tokenizer so subsequent source edits do not change results silently.

M^(-1) g is a metric-derived intervention. Its transformed-coordinate alignment is true by construction, not independent intervention evidence. Do not treat input-token embedding differences as contextual causal interventions.

Categorical analysis accepts explicit vertices in a declared raw space. Simplex checks require affine independence; regularity additionally checks pairwise distances under a declared relative tolerance. Retain degenerate vertices and mark undefined angles with NaNs and a validity mask. Arbitrary category-token rows are not automatically the representations constructed in arXiv:2406.01506; no token-to-category estimator or paper-specific construction is included in the core feature.

TransformerBridge support and basis contract

Use the actual linear unembedding input after unfolded final LN/RMS, including its learned gain and bias. Do not automatically fold normalization, apply a cached activation scale, enable compatibility mode, move the model, or change training flags. Unembedding bias is excluded from the metric and must still be accounted for in logit readouts.

The supported Bridge contract requires a raw causal decoder with a direct LN/RMS-to-plain-linear vocabulary head and consistent d_model/d_vocab_out. Reject compatibility processing, final output projections, active soft caps, custom post-readout transforms, parametrized/custom-forward heads, and meta/offloaded readouts.

Earlier residual hooks are not automatically in this basis. Tensor compute policies support CPU/CUDA, but the current geometry validation is CPU-based; the core does not claim MPS support or universal architecture coverage. A selected token population changes the covariance estimator, and the paper's full-vocabulary assumptions do not automatically transfer to arbitrary subsets. Ridge also changes the metric.

Independent reproduction and behavioral validation

Pin reference revisions, model/tokenizer configuration, contrast data, preprocessing, and dtype. Check licenses before redistributing reference code/data; provide regeneration routes when redistribution is unavailable.

The paper derives interventions through Equation (4.1) and then measures their effects. Reproduce that approach:

• Measure target and off-target changes across controlled counterfactual token configurations and held-out contexts.
• Include unedited, no-op, and appropriate matched controls.
• Compare raw-geometric baselines and run conditioning/ridge sensitivity analyses.
• Establish tolerances from measured reference evidence.
• Report negative findings and limitations instead of requiring every concept or small model to show clean orthogonality or simplex structure.

A separately fitted contextual intervention estimator is a possible research extension, not a required existing author algorithm to port. Such an estimator would need its own identification assumptions, dataset, and evaluation protocol.

The issue is complete when the core feature is merged with required repository validation resolved and independent reproduction, controlled intervention results, sensitivity evidence, and reproducible artifacts or commands are delivered.

Alternatives

• Raw Euclidean cosine only: retain it as a baseline, but it does not provide an explicit covariance-based metric or dual-space conversion contract.
• Activation-covariance whitening: useful for other questions, but data-dependent and a different estimator from the unembedding-covariance metric proposed here.
• Vendor the authors' notebooks: useful as inspected references where licensing permits, but not a substitute for a maintained API with independent tests and an explicit support contract.
• Put geometry only inside a steering or probing tool: reduces reuse and risks inconsistent space/normalization semantics across consumers. Keep the numerical core independent; consumer adapters can build on it.
• Treat whitening identity as empirical validation: it validates arithmetic by construction, not concept separability or behavioral intervention effects.

Additional context

Validation should distinguish algebraic correctness from empirical model evidence. Synthetic fixtures should verify covariance construction, dual-space transforms, regularization, and categorical diagnostics. Download-free integration tests should exercise the TransformerBridge and tokenizer contracts. Independent reproduction and controlled behavioral experiments are required before making claims about semantic geometry or causal separability in pretrained models.

Hierarchical tooling, automated lexicon/template construction, steering/probing adapters, unrestricted multi-token concepts, and a universal support matrix are outside this issue. Dense covariance/eigendecomposition can be expensive despite requiring no training; resource limits and numerical policies must be documented.
Hierarchical tooling, automated lexicon/template construction, steering/probing adapters, unrestricted multi-token concepts, and a universal support matrix are outside this issue. Dense covariance/eigendecomposition can be expensive despite requiring no training; resource limits and numerical policies must be documented.

Checklist

• I have checked that there is no similar issue in the repo (required).

Repository searches for "representation geometry", "causal inner product", and "whitening" found no matching issue or PR. Related projection-kernel work is linked above and has a different scope.