Canonical locks that encode part-whole hierarchies

Rajat Modi; Yogesh Singh Rawat

cs.AISep 22, 2026

Canonical locks that encode part-whole hierarchies

Authors: Rajat Modi, Yogesh Singh Rawat

Organizations: Department of Computer Science, University of Central Florida, United States.

Abstract

One of the challenges in representational learning is how to encode part-whole hierarchies in a neural net. Prior works rely on flattening tree-like structures into string-like sequences and training a sequence-to-sequence model via autoregression. While such a representation works for parse-trees in NLP, it is not entirely clear how to make it work for images. Thus, we propose a geometric primitive called canonical locks. The key idea is that parts/wholes can be modelled as higher-dimensional vectors ( $d \geq 4$ ), and information can be encoded in their relative phase differences. Inductively, the net consists of positionally-bound bottom-up and top-down neural fields, which drive each other to achieve a state of thermal equilibrium. Additionally, we show the existence of a few symmetrical configurations in the net. The computational iterations taken to break these symmetries depend on the angle between parts/wholes arranged on a disk (or more precisely a ring) in higher dimensions. It also appears to have connections to the psychological phenomenon of mental rotation.

Figures & tables

Figure 2 : Two ways of encoding part-whole hierarchies. (i) Encoding hierarchy inside the neural net; (ii) Encoding hierarchy outside the neural net. A darker color on the neural connection indicates a higher probability that a given part chooses a whole in a higher layer as its parent.

Figure 3 : (i) Face is at $\mathtt{90}^{\circ}$ w.r.t lip; (ii) An alternative hypothesis where face is at $\mathtt{45}^{\circ}$ w.r.t lip. Visually, there is no constraint about which of the angles $\mathtt{45}^{\circ}/\mathtt{90}^{\circ}$ will be preferred.

Figure 4 : A hypersphere analogy. Each disk is a great circle of the largest possible radius lying on the sphere.

Figure 5 : A 2D projection of the great circle taken from the hypersphere. Part’s and Whole’s arrows can be thought of as ‘needles’ oscillating in a compass.

Figure 6 : A perfect canonical lock.

Figure 7 : Some tough symmetrical configurations of parts/whole.

Figure 9 : Resolving the collapsed symmetry.

Figure 10 : Fitting a linear line model to fundamental symmetries. x axis is no of iterations. y axis is the angle of center of gravity of part vectors w.r.t whole vector.

Figure 11 : Any random configuration of part-whole vectors on the disk appears to cut $y=0$ multiple times. Each x-axis coordinate (eg, 40, 60) may be thought of as a fundamental symmetry.

Figure 12 : The tri-fold symmetry

Figure 13 : Synchrony across levels results in loss of information.

Figure 14 : A successful case of encoding information in a hierarchy.

Figure 15 : An arm analogy. Representing coordinate frames with three layers ( Hinton, 1977 )

Figure 16 : Representing coordinate frames with a single layer.

Figure 17 : A human stick figure showing the hierarchy of body parts. Each part has chosen its own canonical axis. Similarly, the part corresponding to the left hand knows that its whole is the left forearm. We want a neural net to learn similar hierarchy. Figure gratefully inspired from ( Hinton, 1977 ) .

Figure 18 : Lattice of a learning machine. Solid lines are connections already formed. Dotted lines are connections currently being formed. The lattice is equivariant around the origin.

Figure 19 : A single part sending a vote to the higher level whole

Figure 20 : The buggy capsule-voting mechanism

Figure 21 : Eliminating part-to-whole transformation matrices

Figure 22 : A two level hierarchy showing formation of whole vectors from part vectors

Figure 23 : Representational collapse in mean field approximation

Figure 24 : An alternate solution involves bottom-up net

W_{f}

Symbol	Description	Symbol	Description
$\mathbf{x}\in\mathbb{R}^{H\times W\times C}$	Input image	$W_{f}=\{W_{f,1},\dots,W_{f,L}\}$	All bottom-up weight matrices
$l\in\{1,\dots,L\}$	Current level in the hierarchy, $L$ total levels	$W_{b}=\{W_{b,1},W_{b,2},\dots,W_{b,L}\}$	All top-down weight matrices
$(i,j,l)$	Spatial grid coordinates at level $l$	$W_{outer}=\{W_{e},W_{f}\}$	Outer-loop weights
$f_{e}(\cdot)$	Encoder	$W_{inner}=\{W_{b}\}$	Inner-Loop weights
$g_{f,l}(\cdot)$	Bottom-up Net at level $l$	$\text{Opt}_{outer}$	Optimizer for $W_{outer}$
$W_{f,l}$	Bottom-up Net’s weights	$\text{Opt}_{b,l}$	Optimizer for top-down net in level $l$

Table 1 : Mathematical notation used in the rest of this paper

Figure 25 : Asynchronous Perception Machine.

W_{e}

is encoder,

W_{f}

is the bottom-up net and

W_{b}

is the top-down net. They may be implemented with either convolutional or attention based connectivity.

Figure 27 : The slurp update.

Figure 28 : An analysis of the case of perfect agreement.

Figure 29 : Combining bottom-up and top-down nets together. They are identical in structure. Each box shows the number of neurons present in that layer. Each neuron is marked as a ‘double circle’ to indicate that it may be used to ‘both’ present an input to the net, as well read output from it. Intermediate activation functions are omitted for clarity.

Figure 30 : Hinton Diagrams for synchronization: The top-down and bottom-up net can drive each other to identical weights. They follow signed symmetry: the magnitude of sign of the weight in each box is the same. However, the magnitude of the weight may be different. For eg, the thickness of blue highlighted square, is different in layer L3 of both nets.

Figure 31 : The fixed class-name hierarchy (L1 finest to L5 coarsest) that the model’s 5 classification heads are trained against, shown as a dendrogram.

Explore similar work

Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic
Jun 8, 2026Hu Tan, Kuo Gai, Shihua ZhangNeural Collapse Neural Representations
Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; and School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China · Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS), Shanghai, China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China; and Key Laboratory of Systems Health Science of Zhejiang Province, School of Life Science, Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Chinese Academy of Sciences, Hangzhou 310024, China
Topological Signatures of Grokking
May 7, 2026Yifan Tang, Qiquan Wang, Inés García-Redondo +1Persistent Homology Neural Network Representations
Imperial College London · Queen Mary University of London · University of Fribourg
Transformation Laws in Neural Representations: Structure, Realisability, and Construction
Sep 16, 2026Yuan SunNeural Representations Transform Concepts
School of Mathematical Sciences, Beijing Normal University, Beijing 100875, China

Jun 8, 2026, cs.LGJ/K move · Enter open · S save

Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic

Hu Tan, Kuo Gai, Shihua Zhang

Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; and School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China · Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS), Shanghai, China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China; and Key Laboratory of Systems Health Science of Zhejiang Province, School of Life Science, Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Chinese Academy of Sciences, Hangzhou 310024, China

Neural Collapse Neural Representations

Open paper AskPDF arXiv

While neural collapse (NC) predicts that a

K

-class-balanced classifier should organize terminal representations as a

(K-1)

-dimensional simplex equiangular tight frame (ETF), modular addition consistently enters a different regime: networks compress to a two-dimensional cyclic geometry in which both classifier weights and token embeddings lie on circles. We refine the explanation of this phenomenon in three directions. First, we formalize a layerwise non-uniform training mechanism: downstream classifier weights are driven by dense cross-entropy gradients into a rank-2 equiangular configuration before upstream embeddings fully reorganize, and once this classifier plane forms, backpropagated feature gradients constrain embedding motion to the same plane while weight decay suppresses orthogonal components. Second, after this subspace locking, the induced in-plane dynamics admit an entropy-regularized transport interpretation on

S^1

; combined with modular-addition labels, this reduces embedding formation to phase alignment, whose minimizers are single-frequency characters of

\mathbb{Z}/P\mathbb{Z}

and hence equal-angle points on a circle. Third, we quantify why this solution prevails over NC: a simplex ETF gains only an

O(1)

advantage in cross-entropy, whereas the cyclic rank-2 solution enjoys a

Θ(K)

advantage under Schatten or weight-decay surrogates, yielding a critical threshold

λ_{\mathrm{crit}} = Θ(1/K)

. Our results explain both why classifier weights move first and why embeddings subsequently align with them, showing that grokking on modular arithmetic is governed not by maximal separation alone but by a task-structured trade-off between separation, symmetry, and complexity.

Abstract

Explore similar work

Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic

Topological Signatures of Grokking

Transformation Laws in Neural Representations: Structure, Realisability, and Construction