mapping-networks 0.2.0

No project description provided

Mapping Networks

mapping_networks is a model-agnostic, production-oriented PyTorch package for training target models through low-dimensional parameter manifolds. It is inspired by the paper Mapping Networks. The library decouples model architecture from parameter representation, allowing you to optimize neural networks by updating compact, trainable latent coordinates instead of mutating the target module's parameters.

---

Installation

Install mapping_networks from your local workspace:

pip install mapping_networks

Or add it using Poetry:

poetry add mapping_networks

---

Quickstart

Train a standard PyTorch model through a low-dimensional layer-wise latent representation:

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from mapping_networks import MappingModel, MappingLoss, ClassificationLoss, MappingTrainer

# 1. Create target model and training data
target = nn.Sequential(nn.Linear(10, 16), nn.ReLU(), nn.Linear(16, 2))
data = TensorDataset(torch.randn(100, 10), torch.randint(0, 2, (100,)))
loader = DataLoader(data, batch_size=16)

# 2. Wrap target model with low-dimensional mapping strategy
model = MappingModel(
    target_model=target,
    latent_dim=32,
    strategy="layerwise"
)

# 3. Configure composite loss and trainer
loss_fn = MappingLoss(task_loss=ClassificationLoss())
trainer = MappingTrainer(
    model=model,
    train_loader=loader,
    loss_fn=loss_fn,
    learning_rate=1e-3,
)

# 4. Train the latent parameters
trainer.fit(epochs=5)

---

Core Concepts

• Latent Manifold: Trainable coordinates $z$ which undergo transformation to map to target parameter space.
• BaseMapper: Projects $z$ to flat generated parameter descriptors (e.g., using fixed, orthogonal projections).
• BaseModulation: Integrates generated descriptors back into target parameters (e.g., additive $W_{ij} \leftarrow W_{ij} + \alpha z_i$, or affine).
• Generation Strategy: Defines the mapping scope.
  • "slvt" (Single Latent Vector Training) projects the entire model from a single global latent vector.
  • "layerwise" constructs independent smaller latent coordinates per layer.
  • "grouped" allows custom parameter subdivision.
• MappingLoss: A composite loss that balances task loss (classification/regression) with stability, smoothness, and cosine alignment regularization components.

---

Paper Correspondence

The code maps directly to the concepts defined in the paper:

• Fixed Projection Matrices: Registered as PyTorch buffers inside MLPMapper so they stay frozen and are excluded from DDP/optimizer updates.
• Additive Modulation: Implemented in AdditiveModulation representing $w_{ij} \leftarrow w_{ij} + \alpha \cdot z_i$.
• Regularization Terms: Fully implemented in MappingLoss:
  • Stability Loss ($L_{\text{stability}}$): penalizes changes in output predictions when adding small noise perturbation to the latent vector (StabilityLoss).
  • Smoothness Loss ($L_{\text{smoothness}}$): penalizes the Jacobian norm of the mapper to enforce a smooth manifold (SmoothnessLoss).
  • Alignment Loss ($L_{\text{alignment}}$): maximizes alignment via cosine distance between latent vectors and weight summaries (AlignmentLoss).

---

Limitations

• Stateless Execution: The target model is executed via torch.func.functional_call. It is never mutated, but any custom operations inside the target model must support stateless calls.
• Stateful Buffers (e.g., BatchNorm): Persistent target buffers (like running mean/variance) are isolated. During training, running stats do not mutate.
• Parameter Tying: Target models with tied parameters (e.g., shared embeddings in language models) are not supported out of the box and will fail validation explicitly.
• torch.compile: Compatible with torch.compile in non-fullgraph mode; fullgraph compilation may fail due to dynamic shapes or functional execution limits depending on the target model.

---

Documentation

The official documentation is hosted on Read the Docs:

👉 mapping-networks.readthedocs.io