Learning on Model Weights using Tree Experts

TL;DR

This paper introduces ProbeX, a lightweight probing method that leverages the structure of Model Trees to classify models and predict their training data categories using only model weights, enabling zero-shot model search.

Contribution

The paper presents ProbeX, a novel probing technique designed specifically for single-layer model weights, and demonstrates its effectiveness in classifying models and their training data.

Findings

ProbeX accurately predicts training data categories from model weights.

Model weights within the same Model Tree show less nuisance variation.

ProbeX enables zero-shot model classification via weight-language embedding.

Abstract

The number of publicly available models is rapidly increasing, yet most remain undocumented. Users looking for suitable models for their tasks must first determine what each model does. Training machine learning models to infer missing documentation directly from model weights is challenging, as these weights often contain significant variation unrelated to model functionality (denoted nuisance). Here, we identify a key property of real-world models: most public models belong to a small set of Model Trees, where all models within a tree are fine-tuned from a common ancestor (e.g., a foundation model). Importantly, we find that within each tree there is less nuisance variation between models. Concretely, while learning across Model Trees requires complex architectures, even a linear classifier trained on a single model layer often works within trees. While effective, these linear classifiers are computationally expensive, especially when dealing with larger models that have many parameters. To address this, we introduce Probing Experts (ProbeX), a theoretically motivated and lightweight method. Notably, ProbeX is the first probing method specifically designed to learn from the weights of a single hidden model layer. We demonstrate the effectiveness of ProbeX by predicting the categories in a model's training dataset based only on its weights. Excitingly, ProbeX can map the weights of Stable Diffusion into a weight-language embedding space, enabling model search via text, i.e., zero-shot model classification.