PaperlayerView exact revision on arXiv

Source-linked AI summary

Capsule Lens: Locating and Tracking Concept Geometry in Model Representations

Yiming Tang, Harshvardhan Saini, Samyak Jha, Huaming Chen, Xufeng Duan, Dianbo Liu

arXiv:2609.05575v1cs.LGcs.CV

TL;DR

Understanding how concepts occupy and change within model representations matters for mechanistic interpretability, but existing approaches offer limited direct characterization and validation of concept geometry. Capsule Lens fits and tracks capsule-shaped approximations of concept regions across static models and training settings. Its case studies find broad restructuring during CLIP pretraining and more localized, concept-specific changes during RL post-training.

  • Problem

    Interpretability lenses map representations into interpretable spaces, while geometric accounts typically do not directly estimate and validate occupied regions; representation changes during training are also less studied.

    Source: §1 · p. 1

  • Method

    Capsule Lens fits a closed-form capsule to a concept’s representation geometry, adds span and norm coverage curves, and validates fits on held-out samples.

    Source: §1 · p. 2

  • Results

    Across three case studies, CLIP pretraining produces broad network-wide restructuring, whereas RL post-training yields localized, concept-specific changes, including sparse RLVR sharpening at particular layers.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Capsule Lens provides a tool for locating and characterizing concept geometry in static representations and tracking its changes across training settings.

    Source: §6 · p. 9

    The geometric patterns from the three dynamic case studies merit substantially more systematic investigation.

    Source: B DISCUSSION · p. 20

Abstract

from arXiv · show

Understanding how concepts are encoded in the internal representations of machine learning models is a central problem in mechanistic interpretability, essential both for the science of deep learning and for the trustworthy deployment of increasingly capable models. Existing approaches to interpret model representations mainly map representations onto more interpretable spaces and do not directly characterize how concepts occupy representation space; various hypotheses have been proposed, but often lack of rigorous validation and largely focus on static representations. In this work, we introduce Capsule Lens, a framework that matches the region a concept occupies with a simple, trackable geometric form, a capsule, defined by several interpretable parameters, fitted in closed form to each concept's geometry and validated on held-out samples. We apply Capsule Lens in two major settings: static and dynamic representations. On static representations, we demonstrate how to locate concept geometry across various models, and how the span and norm curves uncover important geometric characteristics. On dynamic representations, we present three case studies tracking representation drifts induced by distinct training settings, CLIP pretraining, RL post-training on visual question answering, and RL post-training on mathematical reasoning. These analyses reveal qualitatively different geometric dynamics, ranging from broad network-wide restructuring in CLIP pretraining to localized and concept-specific changes in RL post-training. Our results include findings aligned with existing literature as well as novel observations. We believe Capsule Lens stands as a promising tool for locating, analyzing, and tracking concept geometry in both static and dynamic representations.

1 INTRODUCTION

Capsule Lens addresses the gap between mapping representations into interpretable spaces and directly estimating the geometry occupied by concepts. It fits and tracks capsules to characterize static representations and distinct training-induced changes.

1 INTRODUCTION

  • Existing interpretability lenses map activations into interpretable spaces, while geometric accounts often propose structures without directly estimating and validating occupied regions.Both lines of work largely focus on static representations, offering limited insight into how concept geometry changes during training.Source: §1 · p. 2
  • The framework treats a concept as a binary indicator on a dataset and approximates its positive samples’ representation-space region with a capsule whose parameters can be tracked.Source: §1 · p. 2
  • Capsule Lens fits a closed-form capsule to each concept’s representation cloud and uses span and norm coverage curves to describe its geometry.The framework validates fitted capsules on held-out samples.Source: §1 · p. 2
  • For static representations, Capsule Lens locates concept geometry across diverse model architectures and characterizes internal structure, disentanglement, and layer-wise evolution.Source: §1 · p. 2
  • Across three case studies, CLIP pretraining produces broad network-wide restructuring, while RL post-training yields more localized, concept-specific geometric changes.Under RLVR, a small subset of concepts sharpens at particular layers while most remain nearly unchanged.Source: §1 · p. 2

2 METHOD

Capsule Lens represents each concept’s positive samples as a geometric point set and summarizes that geometry with a fitted capsule. It uses coverage curves, held-out validation, and checkpoint comparisons to analyze and track concept geometry.

2.1 PRELIMINARIES: CONCEPTS AND CAPSULES

  • Capsule Lens treats a concept as a binary indicator and its geometry as the representation-space point set of positive samples.A fitted parametric shape makes this geometry readable through its parameters.Source: §2.1 · p. 3
  • The capsule is an angular cone intersected with a norm shell, whose axis, span, and norm bounds describe direction, angular spread, and magnitude range.The norm-band width is nhigh − nlow.Source: §2.1 · p. 3

2.2 CAPSULE MATCHING AND COVERAGE CURVES

  • Capsule matching fits each concept’s representation cloud in closed form, using the normalized mean direction as its axis and sample extrema for span and norm bounds.The resulting capsule covers all fitting samples by construction.Source: §2.2 · p. 4
  • Span and norm coverage curves show how concept samples populate the fitted capsule beyond its boundary parameters.A steep early span-curve rise indicates concentration near the axis, while a late rise indicates broader angular distribution; flat segments between rises suggest multiple angular clusters.Source: §2.2 · p. 4

2.3 CAPSULE VALIDATION

  • The validation protocol fits capsules on one positive-sample split, then assesses held-out coverage and separation from other concepts.Held-out coverage measures whether the fitted region generalizes to unseen samples of the same concept; disentanglement assesses overlap with other concepts.Source: §2.3 · pp. 4–5

2.4 TRACKING CONCEPT GEOMETRY

  • Comparing independently fitted capsules at training checkpoints tracks axis rotation, angular-span changes, and shifts or expansions in the norm band.Aggregating drift signals across concepts and layers localizes where training reshapes representations and which concepts it affects most.Source: Table 1 · p. 5, §2.4 · p. 5

3 VALIDATING CAPSULE LENS

Capsule Lens is validated on constructed visual and textual concepts through held-out coverage and disentanglement, and compared with alternative geometric approximations. Capsules combine high coverage with informative, separable geometric parameters.

3 VALIDATING CAPSULE LENS

  • The validation section evaluates fitted capsules using held-out coverage and disentanglement, compares alternative geometric approximations, and includes robustness analyses of sampling and boundary estimation.The protocol uses disjoint fitting and validation samples.Source: §3 · p. 5, §3.2 · p. 6

3.1 CONCEPT CONSTRUCTION

  • Concepts are positive-sample collections, constructed from annotated visual regions or textual token groups and human-labelled samples.Visual concepts use Visual Genome; textual concepts include lexical concepts from several corpora and labelled concepts from SemCor.Source: §3.1 · pp. 5–6

3.2 HELD-OUT COVERAGE AND DISENTANGLEMENT OF CAPSULES

  • 0.943 overall mean held-out coverage; visual concepts remain near the balanced disentanglement baseline, while textual concepts—especially lexical ones—separate more strongly.Table 2 reports mean and standard deviation of Cov and Dis across fitted capsules.Source: §3.2 · p. 6, Table 2 · p. 6

3.3 COMPARISON WITH ALTERNATIVE GEOMETRIC APPROXIMATIONS

  • Capsules achieve coverage comparable to cones and balls while attaining the highest disentanglement among approximations with meaningful coverage.Their parameters also separate angular extent from representation norm.Source: §3.3 · p. 6

4 ANALYZING CONCEPT GEOMETRY IN STATIC REPRESENTATIONS

In static representations, Capsule Lens reveals internal structure that capsule boundaries alone do not show and tracks how disentanglement varies with model depth. The observed depth profiles contrast across visual and textual representations.

4.1 INTERNAL STRUCTURE OF CONCEPT GEOMETRY

  • For stripes at CLIP ViT-B/32 layer 12, span coverage reveals two angular clusters while norm coverage rises smoothly, beyond what the fitted span alone shows.The fitted span is sa = 0.655; coverage curves show how representations populate the capsule boundary.Source: Figure 2 · p. 7, §4.1 · p. 7

4.2 LAYER-WISE EVOLUTION OF CONCEPT DISENTANGLEMENT

  • Several visual encoders become more disentangled toward later layers, while textual concepts—especially lexical concepts—become more entangled in later Qwen2.5-1.5B layers.The textual trend is consistent with prior findings that lexical identity and semantics are more explicit in earlier representations (Liu et al., 2024; Li & Subramani, 2025).Source: §4.2 · p. 7, Figure 3 · p. 8

5 TRACKING CONCEPT GEOMETRY IN DYNAMIC REPRESENTATIONS

Capsule Lens tracks distinct geometric changes across CLIP pretraining and two RL post-training settings. The observed dynamics range from broad network restructuring to changes concentrated in particular layers or concepts.

5.1 CASE STUDY I: CLIP PRETRAINING

  • During CLIP pretraining, axes rotate across probed blocks, spans widen in blocks 2–8 but contract in the final block, and norm bands contract steadily.The final-block span decreases from sa ≈0.61 to 0.28; per-concept drift is comparatively uniform.Source: Figure 4 · p. 8, §5.1 · p. 8

5.2 CASE STUDY II: RL POST-TRAINING ON VISUAL QUESTION ANSWERING

  • PPO leaves axis spans nearly unchanged but concentrates axis drift at the final layer, where the norm band contracts rapidly and then stabilizes as rotation continues.Final-layer drift reaches ∆a ≈0.026; earlier layers remain close to base geometry, making the restructuring smaller and more localized than pretraining.Source: Figure 5 · p. 9, §5.2 · p. 9

5.3 CASE STUDY III: RL POST-TRAINING ON MATHEMATICAL REASONING

  • RLVR changes are sparse across concepts and layers: most pairs stay near base geometry, while a subset shows span contractions or norm-bound shifts.For unit time minute at layer 27, span contracts from sa = 0.711 to 0.594 while axis orientation and norm geometry remain comparatively stable.Source: §5.3 · p. 9

6 CONCLUSION

The paper concludes that Capsule Lens can locate and characterize concept geometry in static representations and track its changes during training. It situates this approach against interpretability lenses and neural-geometry hypotheses that do not yet provide sufficient support for tracking training-induced change.

6 CONCLUSION

  • Capsule Lens locates and characterizes concept geometry in static representations and tracks its changes in dynamic representations across three training pipelines.The paper reports validation across models and case studies of geometric shifts under training.Source: §6 · p. 9

A.1 LENS TO INTERPRET MODEL INTERNALS

  • Existing lenses map activations to tokens, language, or features but do not directly report a concept’s spread, anisotropy, or norm structure in representation space.The related-work discussion also notes that current neural-geometry hypotheses often have suggestive but partial evidence, and training-induced changes lack a tracking method.Source: §A.1 · p. 19, §A.2 · p. 19

A.2 NEURAL GEOMETRY

  • The paper frames tracking training-induced neural-geometry changes as an open methodological gap, with existing characterizations largely resting on hypotheses not yet sufficiently supported.Neural geometry concerns how information is spatially organized in neural-network representation spaces.Source: §A.2 · p. 19

A.3 UNDERSTANDING HOW POST-TRAINING (SFT & RL) WORKS

The discussion frames representation tracking as a complement to static interpretation, while treating the training patterns found in three case studies as hypotheses needing broader validation. It also outlines future uses and the paper’s systematic validation plan for Capsule Lens.

A.3 UNDERSTANDING HOW POST-TRAINING (SFT & RL) WORKS

  • Post-training research offers conflicting accounts of RL and SFT, while most studies examine output behavior rather than changes in representation space.The discussion notes that effects may depend on optimization choices, data composition, and model capability.Source: §A.3 · p. 19

B DISCUSSION

  • Tracking geometry complements static interpretation by measuring where and how representations change across training states.The discussion contrasts this with asking only what a model represents.Source: B DISCUSSION · p. 20
  • Coverage acts like recall and disentanglement like precision, distinguishing how many unseen positive samples a capsule captures from how specifically it selects its concept.These metrics characterize different properties of a simple capsule approximation.Source: B DISCUSSION · p. 20
  • Capsules trade geometric flexibility for a compact, interpretable parameterization whose quantities can be compared across layers, models, and checkpoints.The axis describes orientation, span angular spread, and norm bounds radial extent.Source: B DISCUSSION · p. 20
  • The case studies suggest broad CLIP restructuring, localized VQA PPO changes, and sparse concept- and layer-specific RLVR concentration, but broader laws remain untested.Testing whether patterns generalize across models, datasets, optimization settings, and random seeds is beyond this work’s scope.Source: B DISCUSSION · pp. 20–21

C FUTURE WORK

  • Future directions include systematic studies of training pipelines and layer differences, capsule intersections for concept hierarchies, and representation steering.The proposed steering direction would use capsule axes, angular spans, or norm-band regions to study effects on model behavior.Source: C FUTURE WORK · p. 21

D SYSTEMATICALLY VALIDATING CAPSULE LENS

  • The validation program examines sample-count effects and sampling stability, percentile-based boundaries and their coverage–disentanglement trade-off, and alternative geometric approximations.The alternatives are compared under different boundary choices.Source: D SYSTEMATICALLY VALIDATING CAPSULE LENS · p. 22

D.1 EFFECT OF THE NUMBER OF SAMPLES PER CONCEPT

The validation experiments test sampling, boundary selection, and alternative shapes, then add concept-level views of CLIP pretraining drift. They show that capsule performance depends on sample count and boundary tightness, while the CLIP span contraction affects the population and concept drift varies.

D.1 EFFECT OF THE NUMBER OF SAMPLES PER CONCEPT

  • Held-out coverage reaches 0.878 at m = 25, 0.941 at m = 50, and 0.971 at m = 100, with diminishing gains beyond the main-experiment size of 50.The fitted axis span and norm band also increase with sample count.Source: §D.1 · p. 22, Figure 7 · p. 22
  • Mean axis distance between independently fitted capsules falls from 0.1150 at m = 5 to 0.0141 at m = 50 and 0.0071 at m = 100.The trend is consistent across model components, providing a reference scale for checkpoint drift.Source: §D.1 · p. 22, Table 4 · p. 23

D.2 ROBUSTNESS TO CAPSULE BOUNDARY ESTIMATION

  • At q = 95, mean axis span is 0.708–0.826 of its original value and mean norm band is 0.740–0.783, showing extrema substantially affect fitted boundaries.Both measures contract as the percentile decreases, across modalities and concept constructions.Source: Figure 8 · p. 23, §D.2 · p. 23
  • Reducing q from 100 to 80 lowers held-out coverage by approximately 0.32–0.34 while increasing disentanglement by 0.061 for visual, 0.140 for lexical, 0.201 for math, and 0.170 for labelled concepts.The authors retain min/max boundaries to prioritize generalization to unseen samples and report disentanglement separately.Source: Figure 9 · p. 23, §D.2 · pp. 23–24

D.3 ALTERNATIVE GEOMETRIC APPROXIMATIONS UNDER DIFFERENT BOUNDARIES

  • Across five boundary percentiles, capsules have higher disentanglement than cones and balls, while cones and balls retain higher coverage; bands are less disentangled and boxes have extremely low coverage.The results indicate directional information accounts for most observed separation, with the norm constraint adding a modest increase in disentanglement.Source: Table 5 · p. 24, §D.3 · p. 24

E CASE STUDY I: TRACKING CONCEPT GEOMETRY IN CLIP PRETRAINING

  • The appendix extends the CLIP-pretraining analysis with within-concept views of block-11 span contraction and concept-level drift rankings.These views use the same training run and checkpoints as the main-text analysis.Source: E CASE STUDY I: TRACKING CONCEPT… · p. 25

E.1 AXIS-SPAN CONTRACTION IS POPULATION-WIDE

  • Block-11 span contraction occurs across the concept population, with the distribution shifting downward and narrowing during training.The mean span changes from approximately 0.60 to 0.28 without a separate non-contracting group at any checkpoint.Source: Figure 10 · p. 25, §E.1 · p. 25

E.2 CONCEPT-LEVEL STRUCTURE

  • Concept-level rankings reveal that block-11 axis drift can be ordered by magnitude, while the 20 most-drifted concepts have totals spanning only approximately 2.88–2.92.The rankings sum drift across the five probed blocks.Source: Figure 11 · p. 25, §E.2 · p. 25

E.3 AXIS DRIFT VS. SPAN CHANGE

The CLIP analysis relates individual concepts’ axis drift to span change, and the next appendix extends the PPO VQA analysis with concept-level views of drift at layer 27.

E.3 AXIS DRIFT VS. SPAN CHANGE

  • The 20 most-stable concepts have total drift of approximately 2.36–2.46, only about 20% below the most-drifted group’s 2.88–2.92.The gap between the extremes is present but modest relative to the overall drift scale.Source: Figure 12 · p. 26
  • At block 11, all plotted concepts have negative span change (−0.57 to −0.14) and axis drift in a narrow range of approximately 0.65–0.87.The figure labels examples with the smallest span change, largest drift, and largest span change magnitude.Source: §E.3 · p. 25, Figure 13 · p. 26

F CASE STUDY II: PPO FOR VLM ON VISUAL QUESTION ANSWERING

F.1 CONCEPT-LEVEL HETEROGENEITY

The PPO VQA appendix examines concept-level drift and its relationship to span change at layer 27. The RLVR appendix finds little population-wide change in the mathematical concepts, while noting that population averages can conceal individual trajectories.

F.1 CONCEPT-LEVEL HETEROGENEITY

  • At PPO step 200, layer 27 accounts for most of the total drift among the 20 most-drifted concepts, while the most- and least-drifted groups differ modestly.The figure comparison gives y-axis maxima of approximately 0.041 and 0.027, respectively.Source: Figure 15 · p. 27, §F.1 · p. 27, Figure 17 · p. 28

F.2 AXIS DRIFT VS. SPAN CHANGE

  • At layer 27, plotted concept drift ranges from 0.018 to 0.033 alongside positive and negative span changes; the figure also tracks the top 60 concepts across PPO steps.The panels distinguish examples by span-change magnitude and encode cosine distance from the base axis over training.Source: §F.2 · p. 27, Figure 16 · p. 28, Figure 18 · p. 29

G CASE STUDY III: RLVR ON MATHEMATICAL REASONING

  • The RLVR appendix follows all 18 mathematical concepts under one GRPO run, with held-out coverage averaging 0.942 and ranging from 0.80 to 1.00.Capsules use 50 fitting samples per concept and a disjoint set of 50 for evaluation.Source: G CASE STUDY III: RLVR ON MATHEMATICAL… · p. 30

G.1 POPULATION-AVERAGED GEOMETRY

  • Under RLVR, mean axis drift stays below 7 × 10−6 at every layer; spans, norm-band midpoint, and width remain largely stable, with only modest late contraction at layer 27.These averages across 18 mathematical concepts indicate little global restructuring.Source: Figure 19 · p. 30, §G.1 · p. 30

G.2 PER-CONCEPT TRAJECTORIES

RLVR produces heterogeneous, concept- and layer-specific changes rather than broad restructuring of the mathematical concept space. Most tracked concepts remain geometrically stable, with a few notable exceptions.

G.2 PER-CONCEPT TRAJECTORIES

  • At layer 27, 7 of 18 concepts change axis span by more than 0.005; unit time minute shows the largest contraction, while population averages remain nearly unchanged.The trajectories indicate selective changes to a small number of concepts at particular layers, not broad restructuring of the mathematical concept space.Source: §G.2 · pp. 30, 32
  • Arithmetic-operation concepts remain largely stable under RLVR, although op percentage varies somewhat more than most and other operation concepts show only limited geometric change.Addition, subtraction, multiplication, division, fraction ratio, and comparison remain largely stable.Source: Figure 21 · p. 31, Figure 20 · p. 31
  • Entity people shows a layer-27 span contraction and substantial intermediate-layer norm-bound changes, while entity countable object remains comparatively stable.Entity people is among the clearest geometric changes outside unit time minute.Source: Figure 22 · p. 32
  • Unit dollar, distance, weight, and time hour remain largely stable; time month and year contract at layer 27, less than time minute.Day, week, month, and year show heterogeneous but generally limited changes.Source: Figure 23 · p. 32, Figure 24 · p. 33

H GEOMETRY-GUIDED REPRESENTATION STEERING WITH CAPSULE LENS

Capsule Lens geometry provides directions for steering representations within a concept and across concepts. Qualitative examples show shifts toward different meanings or concepts, while systematic evaluation remains future work.

H GEOMETRY-GUIDED REPRESENTATION STEERING WITH CAPSULE LENS

  • The method derives unit directions between mode or concept centroids, calibrates intervention magnitude with the fitted norm shell, and adds the intervention to the representation.For source and target centroids, the direction is their difference normalized; magnitude is set using the norm-shell midpoint and a multiplier from 0.6 to 0.9.Source: H GEOMETRY-GUIDED REPRESENTATION… · p. 34
  • For intra-concept steering, opposite modes can selectively elicit different meanings of polysemous textual concepts in representative Qwen2.5-1.5B generations.Table 6 presents examples of steering toward different modes within the same concept geometry.Source: H GEOMETRY-GUIDED REPRESENTATION… · p. 34, Table 6 · p. 34
  • Inter-concept steering shifts generations toward distinct textual and visual concepts using directions between centroids from different concepts.Table 7 shows each base generation followed by a generation steered toward the target concept.Source: H GEOMETRY-GUIDED REPRESENTATION… · p. 34, Table 7 · p. 34
  • Internal modes support steering toward different meanings of the same concept, while relative concept locations support steering across concepts; systematic evaluation remains future work.The reported examples are qualitative.Source: H GEOMETRY-GUIDED REPRESENTATION… · p. 34
Loading 2609.05575v1…