Cosine-Hutchinson Probe Objective

• Cosine-Hutchinson probe objective is a self-supervised training loss that aligns Hutchinson-style probe vectors with their preconditioned images to shape spectral clustering.
• It leverages cosine similarity to emphasize angular alignment over Euclidean reconstruction, enabling effective conditioning in critical spectral subspaces for PCG.
• Integrated in the Hierarchical Transformer Preconditioner, the objective balances fine-scale diagonal correction with low-rank far-field structure to enhance solver performance.

The cosine-Hutchinson probe objective is a self-supervised training loss for neural preconditioners $M \approx A^{-1}$ in preconditioned conjugate gradient (PCG). Introduced in the context of the Hierarchical Transformer Preconditioner, it evaluates the cosine similarity between a batch of Hutchinson-style probe vectors and their preconditioned images under $MA$, thereby optimizing angular alignment on convergence-relevant spectral subspaces rather than enforcing Euclidean reconstruction at a fixed scale (Osborne et al., 13 May 2026). Its defining motivation is that PCG depends on the spectrum of $MA$ up to scaling, so the objective is designed to shape spectral clustering without imposing unnecessary spectral-placement constraints.

1. Formal definition

Given a batch of probe vectors $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$, the method first computes $A\mathbf{Z}$ and then applies the learned preconditioner to form $MA\mathbf{Z}$. The training loss is

$$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$

Equivalently, if the batch is vectorized, the loss is the cosine distance between $\operatorname{vec}(\mathbf{Z})$ and $\operatorname{vec}(MA\mathbf{Z})$. The paper emphasizes that this is a single global angle over all $N \cdot K_z$ entries, not an average of per-probe cosine losses. That global normalization couples the probes through one denominator, which the authors found to make gradients more stable when $MA$0 is small (Osborne et al., 13 May 2026).

The objective has a clear range. A perfect preconditioner yields $MA$1, while anti-alignment yields $MA$2. In this formulation, training does not ask $MA$3 to reproduce $MA$4 in magnitude; it asks the two to point in the same direction.

2. Hutchinson probes as isotropic spectral diagnostics

The probe vectors are Hutchinson-style random vectors $MA$5 satisfying

$MA$6

In this setting they are not used as a trace estimator. Their role is spectral whiteness: because the probe distribution is isotropic, every eigenmode of $MA$7 is sampled fairly. The paper states that for any eigenvector $MA$8,

$MA$9

so the objective sees all modes without bias toward smooth or oscillatory components (Osborne et al., 13 May 2026).

This isotropy makes the probe batch a global diagnostic of preconditioner quality. If $MA$0 acts well on these random directions, the operator is likely behaving well spectrally. The paper further modifies the probes by applying a small amount of damped Jacobi smoothing,

$MA$1

with $MA$2 and two smoothing steps. This reduces high-frequency energy and redistributes probe power toward lower-frequency, coarse-scale modes, which helps train the off-diagonal hierarchical blocks more evenly. Within the architecture under study, that reweighting is important because those blocks represent low-rank far-field structure rather than only fine-scale diagonal correction.

3. Projective geometry and relevance to PCG

The key geometric shift is from a vector-distance objective to a subspace-distance objective. A Frobenius or sparse approximate inverse (SAI) loss tries to make $MA$3 numerically close to $MA$4, whereas the cosine-Hutchinson objective only requires directional agreement. The paper formalizes this by defining

$MA$5

so that

$MA$6

where $MA$7 is the principal angle between the lines spanned by $MA$8 and $MA$9. In that sense, the loss optimizes a distance on projective space $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$0 (Osborne et al., 13 May 2026).

A central property is positive-scale invariance:

$\mathbf{Z} \in \mathbb{R}^{N \times K_z}$1

This matches the paper’s claim that PCG does not need a particular absolute normalization of $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$2. The relevant quantity for CG is the relative spread of the spectrum. The cited convergence bound is

$\mathbf{Z} \in \mathbb{R}^{N \times K_z}$3

Within that framework, the paper explicitly states that CG mainly cares about how tightly the eigenvalues of $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$4 cluster, and that convergence can be rapid even when those clusters are not near $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$5, provided they are tight and away from zero. The cosine-Hutchinson objective therefore targets the action of $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$6 on convergence-critical spectral subspaces rather than demanding exact approximation of $\mathbf{Z} \in \mathbb{R}^{N \times K_z}$7 in a strict Euclidean sense.

4. Contrast with SAI-style Frobenius objectives

The paper contrasts the cosine-Hutchinson loss with SAI-style objectives such as

$\mathbf{Z} \in \mathbb{R}^{N \times K_z}$8

or the sample-based form

$\mathbf{Z} \in \mathbb{R}^{N \times K_z}$9

The first difference is scale. The SAI formulation penalizes Euclidean deviation from $A\mathbf{Z}$0, which implicitly asks that $A\mathbf{Z}$1 match $A\mathbf{Z}$2 at a chosen normalization. That introduces a preferred spectral location, namely around $A\mathbf{Z}$3. The cosine-Hutchinson objective removes that requirement by caring only about angular alignment of $A\mathbf{Z}$4 with $A\mathbf{Z}$5 (Osborne et al., 13 May 2026).

The second difference is invariance. The SAI-style loss is not invariant under $A\mathbf{Z}$6, so it does not respect the same geometry that motivates PCG. The third difference is representational allocation. The authors argue that forcing eigenvalues to cluster near a prescribed anchor can waste capacity on spectral placement, whereas CG only requires tight clusters away from zero. This is the basis for the claim that the cosine objective removes unnecessary spectral-placement constraints and improves conditioning on irregular spectra.

5. Role within the Hierarchical Transformer Preconditioner

The objective was introduced as the key training component of the Hierarchical Transformer Preconditioner. That preconditioner is anchored to a weak-admissibility $A\mathbf{Z}$7-matrix partition with dense diagonal leaves and coarsening off-diagonal tiles, enabling full-graph approximate-inverse computation with $A\mathbf{Z}$8 scaling at fixed block sizes. The network models the inverse through low-rank far-field factors and uses highway connections, specifically axial buffers plus a global summary token, to propagate context across transformer depth (Osborne et al., 13 May 2026).

Within this architecture, training proceeds by producing structured factors for diagonal and off-diagonal blocks, pushing probe vectors through $A\mathbf{Z}$9, applying the learned $MA\mathbf{Z}$0, and measuring alignment between $MA\mathbf{Z}$1 and $MA\mathbf{Z}$2. The paper states that this objective drives the architecture to learn strong diagonal correction at fine scale, low-rank far-field structure at coarse scale, and spectral behavior that improves PCG iteration counts.

The experiments isolate the role of the loss by holding the architecture fixed and changing only the objective. The same architecture trained with the SAI loss performs much worse; the cosine loss keeps improving even when the SAI loss starts to worsen; and the learned spectrum of $MA\mathbf{Z}$3 becomes more tightly clustered under cosine training. The paper further reports that the cosine-trained model can achieve a much larger condition-number reduction than the SAI-trained version, with the spectrum clustering wherever angular alignment is easiest rather than near a fixed anchor such as $MA\mathbf{Z}$4. In the full solver, because both inference and apply are dense, dependency-free tensor programs, the solve loop is captured as a single CUDA Graph.

6. Spectral interpretation, empirical signature, and scope

The paper presents the objective not as a generic similarity metric but as a surrogate for spectral shaping. Because the probes are isotropic, all subspaces are sampled fairly, and the cosine objective encourages $MA\mathbf{Z}$5 to preserve probe directions across the spectrum. The resulting behavior is described as compression of effective eigenvalue spread, which the paper connects directly to better PCG convergence, tighter clustering of $MA\mathbf{Z}$6, and improved behavior on convergence-critical modes (Osborne et al., 13 May 2026).

In the reported training figures, decreases in $MA\mathbf{Z}$7 coincide with collapse of the angle between $MA\mathbf{Z}$8 and $MA\mathbf{Z}$9 toward zero and with tighter eigenvalue clustering. That is the spectral signature used to motivate the loss. In the surrounding application, the paper evaluates stiff multiphase Poisson systems with up to $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$0 density contrast and $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$1, reporting frame rates from approximately $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$2 to $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$3 fps. At $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$4, the solver reaches $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$5 ms/frame, with $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$6 speedup over GPU Jacobi, approximately $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$7 over GPU IC/DILU (AMGX multicolor_dilu), and $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$8 over neural SPAI retrained per scale on the same benchmark. These figures characterize the computational setting in which the objective was proposed rather than constituting a direct standalone metric of the loss itself.

A common misconception is to treat any cosine-based training signal as a reliable proxy for downstream performance. A separate study in latent visual reasoning reports the opposite outcome in a different setting: cosine alignment between supervised latents and visual targets is negatively correlated with accuracy across five variants, with Pearson $$\mathcal{L}_{\cos} = 1 - \frac{\langle \mathbf{Z},\, MA\mathbf{Z} \rangle_F} {\|\mathbf{Z}\|_F\,\|MA\mathbf{Z}\|_F}, \qquad \langle X,Y\rangle_F = \operatorname{tr}(X^\top Y).$$9, and the supervised latents are largely bypassed under corruption tests (Zhang et al., 4 Jun 2026). A plausible implication is that the utility of cosine similarity is domain-specific: in the cosine-Hutchinson objective, cosine is tied to the projective and spectral geometry relevant to PCG, whereas in latent visual reasoning it measures latent-target fidelity without guaranteeing that the aligned latent is causally load-bearing.