Harmonic Loss Trains Interpretable AI Models

Abstract: In this paper, we introduce harmonic loss as an alternative to the standard cross-entropy loss for training neural networks and LLMs. Harmonic loss enables improved interpretability and faster convergence, owing to its scale invariance and finite convergence point by design, which can be interpreted as a class center. We first validate the performance of harmonic models across algorithmic, vision, and language datasets. Through extensive experiments, we demonstrate that models trained with harmonic loss outperform standard models by: (a) enhancing interpretability, (b) requiring less data for generalization, and (c) reducing grokking. Moreover, we compare a GPT-2 model trained with harmonic loss to the standard GPT-2, illustrating that the harmonic model develops more interpretable representations. Looking forward, we believe harmonic loss has the potential to become a valuable tool in domains with limited data availability or in high-stakes applications where interpretability and reliability are paramount, paving the way for more robust and efficient neural network models.

• The paper proposes harmonic loss that computes logits via Euclidean distance, achieving finite convergence targets and reducing grokking.
• The paper demonstrates that harmonic loss improves data efficiency and generalization across various algorithmic tasks with fewer training samples.
• The paper reveals enhanced interpretability through embeddings that align with class centers, uncovering clear geometric structures in both vision and language tasks.

The paper introduces harmonic loss as an alternative to the conventional cross-entropy loss for training neural networks and LLMs (LLM LLM). The key innovation is to compute logits using the Euclidean distance rather than inner products, leading to a loss formulation with several desirable theoretical and empirical properties. In particular, by defining the harmonic logit as

$d_i = \|\mathbf{w}_i - \mathbf{x}\|_2,$

with probability computed as

$p_i = \frac{1/d_i^n}{\sum_{j} 1/d_j^n},$

(where $n$, the harmonic exponent, is ideally chosen as $n \sim \sqrt{D}$ with $D$ being the intrinsic data dimensionality), the approach naturally enforces scale invariance and a finite convergence target. This design enables each weight vector $\mathbf{w}_i$ to converge to the corresponding class center, promoting better interpretability.

Several detailed experiments and analyses underscore the benefits of harmonic loss:

• Faster Convergence and Reduced Grokking The paper demonstrates through toy examples that while standard cross-entropy loss requires driving logits to infinity (due to its inherent demand for exponential separation of classes), harmonic loss requires only that the Euclidean distance between feature representations and corresponding weight vectors tend to zero. This finite target accelerates convergence and results in plateauing $l_2$ weight norms, which in turn minimizes the phenomenon of grokking (i.e., delayed test generalization relative to training loss convergence).
• Data Efficiency and Robust Generalization In a series of algorithmic tasks—including in-context learning on integer lattices, modular addition, equivalence class identification, genealogy learning, and permutation composition—models trained with harmonic loss not only generalize with fewer training samples but also exhibit minimal grokking. Empirical results consistently show that harmonic models achieve high test accuracy with a lower training fraction compared to their cross-entropy counterparts. For example, in the in-context learning task, harmonic models recover the underlying 2D lattice structure in nearly 100% of the runs.
• Interpretable Representations
  • In modular addition tasks, harmonic networks consistently uncover a clean, circular 2D structure with nearly all variance explained by the first two principal components.
  • In genealogy learning and permutation composition tasks, the learned embeddings depict clear clustering corresponding to the underlying combinatorial or hierarchical structures.
  • On MNIST, a simple one-layer network trained with harmonic loss yields weight matrices where non-informative pixels are nearly zero and the weights align well with class centroids, making the model inherently more interpretable.
• Scalability to Language Modeling The paper also extends the approach to a GPT-2 small model trained on OpenWebText. Here, harmonic loss leads to slightly lower final validation loss and more stable training dynamics. Moreover, when evaluating interpretability using parallelogram consistency across token representations (a proxy for functional regularity in the embedding space), the harmonic GPT model consistently exhibits lower parallelogram losses. This result suggests that the harmonic training procedure encourages embedding structures that better capture linguistic analogies, such as morphological transformations (e.g., present-to-past verb forms).
• Theoretical Considerations The authors rationalize the choice of the harmonic exponent by noting that if embeddings are initialized from a Gaussian distribution in a $D$-dimensional space, then normalizing distances with $n \sim \sqrt{D}$ keeps the loss behavior invariant in high dimensions. This theoretical underpinning reinforces the empirical findings regarding the stability and convergence of harmonic loss.

In summary, the paper provides a comprehensive evaluation of harmonic loss, highlighting that it:

• Accelerates convergence by having a finite target for the logits.
• Enhances interpretability through the natural emergence of class centers in the weight space.
• Improves data efficiency and reduces the disruptive effects of grokking.
• Scales effectively to large-scale LLMs, where it refines the internal representations in line with linguistic and geometric analogies.

The work suggests that harmonic loss could become an important tool in settings where model interpretability, robustness, and efficient generalization are critical, particularly in domains with limited data availability or high-stakes application environments.