PAFT: Auxiliary Bottlenecks Make Representations More Projectable

Visual analysis of learned representations plays an important role in interpreting deep neural networks. Current practice predominantly focuses on applying post hoc dimensionality-reduction techniques, such as PCA, t-SNE, and UMAP. However, the strength and prevalence of these tools largely frame dimensionality reduction for visualization as a downstream stage, leaving the geometry learned during training unaddressed. In this work, we propose PAFT (Projection-Aware Fine-Tuning), a fine-tuning method built around an auxiliary branch consisting of a linear bottleneck and a classification head attached to the penultimate layer of a classifier. The classifier and the auxiliary branch are jointly trained with three additional terms in the training objective: prediction consistency between the main and auxiliary heads, together with CKA alignment and neighborhood preservation between the penultimate and bottleneck representations. During training, PAFT learns a linear projection from the penultimate representation to a low-dimensional bottleneck that preserves predictive and geometric structure, while simultaneously reshaping the high-dimensional representation itself so that its structure is more faithfully preserved in lower dimensions. At d = 2, this bottleneck also provides a directly inspectable two-dimensional view without requiring a separate post hoc projection. Across five image-classification datasets and three classifier backbones, we evaluate two-dimensional projections using 20 established projection quality measures. PAFT gives the strongest results, outperforming both the untreated baseline and rank-reduction alternatives overall, while maintaining comparable classification accuracy.