Fully-Connected Layer Decomposition

Factorize fully-connected layers with a truncated SVD

FC_Decomposer replaces each Linear layer by two smaller ones, taken from the truncated SVD of its weight matrix: Linear(n, m) becomes Linear(n, k) followed by Linear(k, m).

from fastai.vision.all import *
from fasterai.misc.all import *
from torchvision.models import vgg16_bn, VGG16_BN_Weights

1. A trained VGG16

VGG16 is the usual example for this: almost all of its parameters sit in the three Linear layers of classifier.

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64))
model = vgg16_bn(weights=VGG16_BN_Weights.IMAGENET1K_V1)
model.classifier[6] = nn.Linear(4096, 2)

learn = Learner(dls, model, metrics=accuracy)
learn.fit_one_cycle(1, 1e-4)
epoch train_loss valid_loss accuracy time
0 0.261939 0.158083 0.937754 00:05 
total = sum(p.numel() for p in learn.model.parameters())
head = sum(p.numel() for p in learn.model.classifier.parameters())

print(learn.model.classifier)
print(f'\n{total:,} parameters, {head:,} of them in the classifier')
Sequential(
  (0): Linear(in_features=25088, out_features=4096, bias=True)
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Linear(in_features=4096, out_features=4096, bias=True)
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): Linear(in_features=4096, out_features=2, bias=True)
)

134,277,186 parameters, 119,554,050 of them in the classifier

2. Decompose

percent_removed=0.5 keeps half of the singular values of each layer. classifier.6 maps 4096 features to the 2 classes, so its SVD has rank 2 at most; exclude leaves it untouched.

decomposed = FC_Decomposer().decompose(learn.model, percent_removed=0.5, exclude=['classifier.6'])

decomposed.classifier
Sequential(
  (0): Sequential(
    (0): Linear(in_features=25088, out_features=2048, bias=False)
    (1): Linear(in_features=2048, out_features=4096, bias=True)
  )
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Sequential(
    (0): Linear(in_features=4096, out_features=2048, bias=False)
    (1): Linear(in_features=2048, out_features=4096, bias=True)
  )
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): Linear(in_features=4096, out_features=2, bias=True)
)
before = sum(p.numel() for p in learn.model.parameters())
after = sum(p.numel() for p in decomposed.parameters())

print(f'{before:,} -> {after:,} parameters ({100 * (1 - after / before):.1f}% fewer)')
134,277,186 -> 91,285,570 parameters (32.0% fewer)

Each decomposed Linear is now two: 25088 -> 2048 -> 4096 in place of 25088 -> 4096. The model goes from 134,277,186 to 91,285,570 parameters, 32.0% fewer, all of it in the classifier.

3. Accuracy on the same validation set

The decomposition is an approximation, so the two models are scored on the same 1478 validation images, without any fine-tuning.

def wilson(k, n, z=1.96):
    "k out of n and its Wilson 95% interval, as a printable string"
    p, d = k / n, 1 + z**2 / n
    c, h = (p + z**2 / (2 * n)) / d, z * ((p * (1 - p) / n + z**2 / (4 * n**2))**0.5) / d
    return f'{k}/{n} = {100 * p:.2f}% [{100 * (c - h):.2f}, {100 * (c + h):.2f}]'

n = len(dls.valid_ds)

def report(name, m):
    lrn = Learner(dls, m, metrics=accuracy)
    with lrn.no_bar(), lrn.no_logging():
        acc = lrn.validate()[1]
    print(f'{name:12s} {wilson(round(acc * n), n)}')

report('original', learn.model)
report('decomposed', decomposed)
original     1386/1478 = 93.78% [92.43, 94.90]
decomposed   1385/1478 = 93.71% [92.35, 94.84]

The original gets 1386 of the 1478 images right and the decomposed model 1385: 93.78% [92.43, 94.90] against 93.71% [92.35, 94.84], Wilson 95%. The two intervals overlap almost exactly, so this single run does not resolve a difference between them. One seed, no fine-tuning after the decomposition.

4. Let the energy choose the rank

energy_threshold replaces percent_removed: each layer keeps the smallest rank that retains that fraction of the squared singular values, so the rank differs from layer to layer.

by_energy = FC_Decomposer().decompose(learn.model, energy_threshold=0.9, exclude=['classifier.6'])

print(by_energy.classifier)
print(f'\n{sum(p.numel() for p in by_energy.parameters()):,} parameters')
Sequential(
  (0): Sequential(
    (0): Linear(in_features=25088, out_features=1928, bias=False)
    (1): Linear(in_features=1928, out_features=4096, bias=True)
  )
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Sequential(
    (0): Linear(in_features=4096, out_features=1233, bias=False)
    (1): Linear(in_features=1233, out_features=4096, bias=True)
  )
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): Linear(in_features=4096, out_features=2, bias=True)
)

81,107,010 parameters

At 90% of the energy the two layers get different ranks, 1928 and 1233 instead of a fixed 2048, and the model lands at 81,107,010 parameters.

5. Activation-aware SVD

With data=, the decomposer first collects the per-input-channel activation RMS over n_batches batches and scales the weight matrix by it before the SVD (ASVD). The shapes are the same; the retained subspace is not.

asvd = FC_Decomposer().decompose(learn.model, percent_removed=0.5, data=dls.train,
                                 n_batches=5, exclude=['classifier.6'])

print(f'{sum(p.numel() for p in asvd.parameters()):,} parameters')
report('ASVD', asvd)
91,285,570 parameters
ASVD         1387/1478 = 93.84% [92.50, 94.96]

The ranks, and so the parameter count, are those of the plain SVD; the retained subspace is not. This run scores 1387/1478, 93.84% [92.50, 94.96], against 1385/1478 for the plain SVD. The intervals overlap, so one run on one calibration set does not separate them.

Summary

Call What it gives you
FC_Decomposer().decompose(model, percent_removed=0.5) a copy of the model with each Linear(n, m) replaced by Linear(n, k) + Linear(k, m)
energy_threshold=0.9 one rank per layer, the smallest retaining 90% of the squared singular values
data=dls.train, n_batches=5 activation-aware SVD: the weights are scaled by the measured input RMS before the factorization
layers=[...], exclude=[...] restrict or skip layers by name

The numbers above come from one run: an ImageNet-pretrained VGG16-bn whose last Linear was replaced and fine-tuned for one epoch on the PETS cats-vs-dogs split at 64 px, scored on n = 1478 validation images. No latency was measured.

See Also