DOCSEARCHbyAlgolia
Submit Docs

WeightWatcher Documentation

repository·master·Indexed 23 days ago

https://github.com/calculatedcontent/weightwatcher

A diagnostic tool for analyzing Deep Neural Networks (DNNs) using Random Matrix Theory and Statistical Mechanics. WeightWatcher allows developers to monitor model training, predict test accuracies, and detect over-fitting or under-fitting without training or test datasets. It supports PyTorch, TensorFlow/Keras, and HuggingFace models, including LLMs, Transformers, and CV architectures. Key features include the alpha metric for generalization, Correlation Trap detection, SVDSharpness and SVDSmoothing transforms, and experimental support for PEFT/LoRA models.

Tokens
14.3K
Snippets
54
Records
89
Agent score
83%

What's inside WeightWatcher

  1. Understand WeightWatcher Generalization Metrics

    master

    WeightWatcher uses metrics derived from the Theory of Heavy-Tailed Self-Regularization (HT-SR) to predict test accuracy without needing test data. Metrics are categorized into Scale, Shape, and Direct Correlation:

    Scale Metrics

    • log_spectral_norm: $\log_{10}\Vert\mathbf{W}\Vert^{2}_{\infty}$
    • stable_rank: $\Vert\mathbf{W}\Vert^{2}{F}/\Vert\mathbf{W}\Vert^{2}{\infty}$
    • mp_softrank: $\lambda_{MP}/\lambda_{max}$

    Shape Metrics

    • alpha: The Power Law (PL) exponent (slope of the tail on a log-log scale). Smaller alpha values generally indicate better generalization.
    • alpha_weighted: A scale-adjusted form of alpha ($\hat{\alpha}=\alpha\log_{10}\lambda_{max}$).
    • log_alpha_norm: Shatten norm ($\log_{10}\Vert\mathbf{X}\Vert^{\alpha}_{\alpha}$).
    • D: Kolmogorov Smirnov Distance (quality of the PL fit).

    Direct Correlation Metrics

    • rand_distance: Distance of the layer ESD from the ideal RMT Marchenko-Pastur (MP) ESD.
    • ww_maxdist / ww_softrank: Related correlation metrics.

    Summary Statistics

    Calling get_summary() returns the average of these metrics across layers. Use average alpha to compare models with different hyperparameters, and average alpha_weighted to compare models with different architectures and depths simultaneously.

  2. Use the alpha metric for Early Stopping

    master
    The alpha metric can serve as an experimental signal for early stopping. When the average alpha (summary statistic) drops below 2.0, it indicates the model may be over-trained. This is most effective for very well-trained models where the optimal alpha is approximately 2.0.

  3. WeightWatcher Requirements and Supported Frameworks

    master

    Requirements

    • Python: 3.7+
    • Frameworks: PyTorch 1.x, TensorFlow 2.x / Keras, and HuggingFace. (Note: Current versions require both tensorflow and torch to be installed).

    Supported Layers

    • Dense / Linear / Fully Connected (including Conv1D)
    • Conv2D

  4. Analyze Correlation Flow in CV models

    master

    For Computer Vision (CV) models, WeightWatcher can be used to visualize and analyze Correlation Flow and test accuracies:

    • VGG: Compare VGG test accuracies against AlphaHat $\hat{\alpha}$ and view Correlation Flow plots (WW-VGG.ipynb).
    • ResNet and DenseNet: View examples of Correlation Flow specifically for these architectures (WW-ResNet.ipynb, WW-DenseNet.ipynb).
    • Full PyTorch CV: A comprehensive computation of all CV models (WW-Full-PyTorchCV.ipynb).

  5. Apply WeightWatcher to LLMs and Transformer models

    master

    WeightWatcher can be used to analyze and compare various Transformer-based models. Available examples include:

    • BERT, RoBERTa, and XLNet: Compare these models using layer Alphas $\alpha$'s (WW-BERT-BlogExample.ipynb).
    • Legal NER: Compare two different LLMs specifically for Legal Named Entity Recognition (WW-LegalNER.ipynb).
    • GPT vs GPT2: Compare GPT and GPT2 models, which share architecture but were trained with different data amounts (WW-GPT.ipynb).
    • Sentence Transformers: Select the best model from a large set (e.g., 100 different HuggingFace models) (WW_Sentence_Transformers.ipynb).

  6. Validate correlation trap removal

    master

    Because the correlation trap workflow is experimental, follow this validation checklist to ensure model integrity:

    1. Establish Baseline: Save baseline model metrics and downstream task scores.
    2. Analyze: Run analyze_traps(...) with a fixed seed and inspect the generated plots.
    3. Iterative Removal: Remove only one trap at a time (e.g., trap_indices=[1]) and iterate.
    4. Re-evaluate: Re-run WeightWatcher metrics and downstream evaluation on the clean_model.
    5. Safety: Always maintain a checkpointed version of the original model for rollback.

  7. Basic Usage of WeightWatcher

    master

    To perform a basic analysis of a PyTorch model, initialize a WeightWatcher object with your model, call .analyze() to generate layer-wise details, and .get_summary() to get generalization metrics.

    import weightwatcher as ww
import torchvision.models as models

model = models.vgg19_bn(pretrained=True)
watcher = ww.WeightWatcher(model=model)
details = watcher.analyze()
summary = watcher.get_summary(details)

  9. Use the fast randomized trap ablation workflow

    master

    For faster iterative ablation loops, use a cached workflow that randomizes the model once and reuses the state across multiple analyze_traps and remove_traps calls. This avoids the overhead of repeated full randomizations.

    Key requirements for the cached workflow:

    • Use randomize_model(..., return_state=True) to obtain the trap_state.
    • Pass the trap_state and permuted_ids into analyze_traps and remove_traps.
    • Use randomized_model=... in subsequent calls.
    • Use layers=sorted(trap_state["permuted_ids"].keys()) to ensure you are only analyzing randomized layers.
    • Performance Tip: Use trap_burden_mode="fast" for approximate overlap and bulk-reference metrics during exploration. Switch to trap_burden_mode="full" only for final verification.
    • Performance Tip: Use a small bulk_mode_sample value during exploration.
    randomized_model, trap_state = watcher.randomize_model(
    model=model, layers=layers, rng=seed, return_state=True, pool=False
)
permuted_ids = trap_state["permuted_ids"]
randomized_layers = sorted(permuted_ids.keys())

trap_df, trap_state = watcher.analyze_traps(
    randomized_model=randomized_model,
    layers=randomized_layers,
    trap_state=trap_state,
    permuted_ids=permuted_ids,
    return_artifacts=True,
    trap_burden=True,
    trap_burden_mode="fast",
    bulk_mode_sample=10,
    plot=False,
    pool=False,
)
ablated_model = watcher.remove_traps(
    randomized_model=randomized_model,
    traps=trap_df.iloc[[0]],
    trap_state=trap_state,
    plot=False,
    pool=False,
,

  11. Use SVDSharpness and SVDSmoothing transforms

    master

    WeightWatcher provides transforms to manipulate weight matrices:

    • SVDSharpness: Used to remove Correlation Traps (WW-SVDSharpness-VGG11.ipynb).
    • SVDSmoothing: Used to create lower-rank approximations for each layer. Examples are available for VGG16 and Keras-based VGG16 (WW-SVDSmoothing.ipynb, WW-SVDSmoothing-VGG16.ipynb, WW-SVDSmoothing-VGG16-Keras.ipynb).

  12. Understanding trap_id and bulk_id indexing

    master

    When working with WeightWatcher outputs and ablation methods, be aware of the following indexing conventions:

    • Public IDs: Both trap_id and bulk_id are 1-based. Use these when calling remove_modes or remove_traps via the trap_indices or mode_ids_by_layer arguments.
    • Internal Indices: The svd_mode_index is 0-based. This is used for internal mapping and debugging.
    • State Mapping: The trap_state['layers'][layer_id] object contains the mappings between public IDs and internal indices: trap_id_to_svd_index and bulk_id_to_svd_index.