Category: Machine Learning

Summary: Testing whether magnitude pruning can recover out-of-distribution compositional generalization in wide networks by increasing representation rank.

Earlier results suggested that wide networks can lose compositional generalization as their hidden representations become more redundant and effectively lower-rank. This experiment asks whether pruning away small-magnitude weights can reverse that effect and restore better out-of-distribution behavior.

The script trains multilayer perceptrons on a synthetic compositional task, applies several levels of global magnitude pruning after training, and then compares in-distribution accuracy, out-of-distribution accuracy, and effective rank before and after pruning. By sweeping width and pruning rate, it tests whether compression can selectively remove the redundancy associated with poor generalization.

That makes pruning a mechanistic probe rather than only a compression tool. If pruning improves compositionality, it would support the idea that extra parameters hurt here by collapsing representation geometry rather than by merely overfitting in a generic way.

Method: Post-training global magnitude pruning of MLPs across widths, combined with compositional-task evaluation and representation-rank diagnostics.

What is measured: In-distribution accuracy, out-of-distribution accuracy, compositional gap, effective rank before and after pruning, and best pruning level by width.