Overview

Neural networks trained on the same task but with different random initializations, architectures, or training procedures learn representations that appear different in their raw coordinate systems. However, these representations often share deep structural similarities that are hidden beneath superficial differences in scale, rotation, and coordinate systems.

This work addresses a fundamental question: What invariances should we build into neural representations to reveal and leverage these structural similarities?

We propose a framework that systematically infuses invariances into neural representations with respect to three key factors:

1. Initialization: Different random seeds leading to different weight initializations
2. Architecture: Variations in network design (depth, width, activation functions)
3. Training modality: Different optimization procedures, hyperparameters, or data augmentation strategies

By carefully choosing which invariances to enforce, we can:

• Reveal structural similarities between independently trained models
• Enable zero-shot model communication and stitching
• Predict which models will align well without expensive alignment procedures
• Design representations that are robust to specific variations while remaining sensitive to others

The Core Insight

The key observation is that different sources of variation induce different types of transformations on learned representations. Some transformations (like rotations from different initializations) can be factored out through appropriate invariances, while others (like semantic differences from different tasks) should be preserved.

Our framework provides tools to:

• Identify which transformations affect representations for a given task
• Quantify the degree of structural similarity despite these transformations
• Infuse appropriate invariances to enable cross-model communication

Method: Framework for Infusing Invariances

Transformations in Neural Representations

Consider a neural network encoder \(E_\theta : \mathcal{X} \rightarrow \mathbb{R}^d\) that maps inputs to a \(d\)-dimensional latent space. When we change factors like initialization \(\phi\), architecture \(\mathcal{A}\), or training procedure \(\mathcal{T}\), the learned representation changes:

where \(T\) is a transformation matrix and \(b\) is a bias term.

The central question is: What structure does \(T\) have, and how can we factor it out?

Invariance Classes

We consider three types of invariances, each targeting different transformation properties:

1. Scale Invariance

• Normalizes representations to unit norm
• Factors out amplitude differences between models
• Implemented via: \(r_i = \frac{e_i}{\|e_i\|}\)

2. Rotation Invariance (Relative Representations)

• Makes representations invariant to orthogonal transformations
• Encodes data relative to anchor points using cosine similarity
• Implemented via: \(r_i = [\cos(e_i, a_1), \cos(e_i, a_2), \ldots, \cos(e_i, a_k)]\)

3. Affine Invariance

• Handles both rotation and translation
• Learned linear transformation between spaces
• Implemented via: \(r_i = W \cdot e_i + b\) (optimized on anchors)

Discovering Optimal Transformation Classes

The framework operates in two phases:

Phase 1: Hypothesis Testing

1. Train multiple models with variations in the target factor (initialization, architecture, or training)
2. For each invariance class, transform representations accordingly
3. Measure alignment between transformed representations using similarity metrics

Phase 2: Selection

• Choose the invariance class that maximizes cross-model alignment
• If no invariance improves alignment, the variations may encode task-relevant information that should be preserved

Metrics for Measuring Alignment

We evaluate alignment using multiple complementary metrics:

• Cosine Similarity: Measures angular alignment between representation spaces
• CKA (Centered Kernel Alignment): Captures linear relationships between representations
• Procrustes Distance: Measures similarity after optimal orthogonal alignment
• Stitching Performance: Direct evaluation via zero-shot encoder-decoder combination

Experiments and Results

We conducted extensive experiments across 8 benchmark datasets spanning computer vision and natural language processing:

Vision: MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, ImageNet-1k Language: TREC, DBpedia, Amazon Reviews (multi-lingual)

For each dataset, we varied initialization, architecture, and training modalities to understand which invariances are most effective.

Invariances to Initialization

We trained 10 models per architecture with different random seeds and measured alignment across different invariance schemes:

Key Finding: Rotation invariance (relative representations) consistently achieves the highest alignment for initialization variations, suggesting that different random seeds primarily induce orthogonal transformations on learned representations.

Invariances to Architecture

We compared representations across different architectures (varying depth, width, activation functions):

Key Finding: Affine invariance performs best for architectural variations, as different architectures introduce both rotational and translational shifts in representation spaces. The learned alignment can adapt to these more complex transformations.

Invariances to Training Modality

We varied training procedures (optimizers, learning rates, data augmentation) while keeping architecture fixed:

Key Finding: The optimal invariance depends on the specific training variation. Optimizer and learning rate changes are well-captured by rotation invariance, while data augmentation strategies may require affine alignment.

Zero-Shot Model Stitching Performance

To validate that improved alignment translates to practical benefits, we evaluated zero-shot stitching performance (connecting encoders and decoders without fine-tuning):

Key Finding: Infusing appropriate invariances dramatically improves zero-shot stitching performance, reducing reconstruction error by up to 35× and maintaining 75-85% of classification performance compared to non-stitched baselines.

Key Findings

1. Structural similarities are pervasive: Independently trained neural networks exhibit strong structural similarities in their learned representations, even when trained with different initializations, architectures, or training procedures. These similarities are hidden beneath coordinate system differences.

2. Optimal invariances are task and variation-dependent: There is no one-size-fits-all invariance. Initialization variations are best handled by rotation invariance, architectural variations by affine invariance, and training modality variations require case-by-case analysis.

3. Invariances enable zero-shot communication: By infusing appropriate invariances, we can enable zero-shot model stitching with performance approaching non-stitched baselines, without any fine-tuning or additional training.

4. Predictive power: The alignment scores under different invariance schemes can predict which models will stitch successfully, providing a practical tool for model selection and ensemble design.

5. Generalization across domains: The framework applies equally well to computer vision and natural language processing tasks, demonstrating broad applicability across modalities and architectures.

Citation

@inproceedings{cannistraci2023infusing,
  title     = {Infusing invariances in neural representations},
  author    = {Cannistraci, Irene and Fumero, Marco and Moschella, Luca and
               Maiorca, Valentino and Rodol{\`a}, Emanuele},
  booktitle = {ICML 2023 Workshop on Topology, Algebra, and Geometry in Machine Learning},
  year      = {2023},
  url       = {https://openreview.net/forum?id=mCm4iiNoNc}
}

Authors

Irene Cannistraci¹ · Marco Fumero¹ · Luca Moschella¹ · Valentino Maiorca¹ · Emanuele Rodolà¹

¹Sapienza University of Rome