cs.LGOct 21, 2026
SaveOptimal Multitask Representations Under a Rate Constraint: An Information Bottleneck Analysis
Organizations: Massachusetts Institute of Technology
Abstract
Which pretraining objective yields the best representation for downstream tasks under an information budget? We formalize this problem as a multitask information bottleneck, using a mutual information constraint to model incomplete information retention. For Gaussian inputs and tasks, we characterize optimal representations and show that a weakly task-aligned specialist can outperform a reconstruction generalist at low rates but lose at high rates. We also construct task families for which every fixed Gaussian prediction objective incurs arbitrarily large relative regret, driven by reversals in optimal allocation priorities. Conversely, the average task information matrix defines an objective with a uniform relative-regret bound controlled by maximum task SNR. CelebA and dSprites experiments illustrate that compression can change which pretraining objective performs best, and even task-aware training can reverse which downstream task is better preserved as rate increases.
Figures & tables
Figure 1: Illustration of Proposition 2 . Here, , , and the evaluation tasks span a -dimensional subspace for which squared cosines of the principal angles with all equal . SNR of all tasks equals 1, including the specialist tasks, reconstruction, and the evaluation tasks.
Figure 2: Rate–distortion gap (top) and optimal rate allocations (bottom) for the two-task example with . Dotted vertical lines mark and , and the dashed line marks the exact reversal .
Figure 3: CelebA specialist–generalist comparison. Left: Gaussian channels under fitted task models, against nominal Gaussian rate. Right: learned discrete codes, against code budget. Top: task-averaged excess NLL. Bottom: generalist minus specialist distortion. Shading shows 95% pointwise percentile bootstrap intervals from test-image resampling (left) and crossed resampling of test images and training seeds (right).
Figure 4: Factor recoverability reversal under direct multitask training on dSprites, . Top: common/moderate-precision and rare/high-precision factors’ recoverability against achieved variational rate. Points show individual seeds, vertical bars indicate recoverability spread, and horizontal bars show rate spread at fixed . Bottom: difference in recoverability indices. The dashed vertical line marks the Gaussian prediction bits, while the observed crossover is bits.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Finite-code rate accounting. Top: empirical joint code entropy versus nominal budget for 300 fitted candidate/budget/seed cells. The identity line denotes equality. The maximum observed entropy is 10.96578 bits, equal to the finite-test-set ceiling . Thus this estimator cannot resolve a 16- or 32-bit population entropy. Bottom: test-image codeword-occupancy histograms for the five-attribute specialist and reconstruction generalist at all six budgets, , using seed 42. Unobserved codewords are omitted. Sparse occupancy can bias plug-in entropy downward. The histograms use logarithmic occupancy and symlog count axes to display both singleton and heavily reused codewords.
Figure 6: Training covariance and whitening. Left: the 512 eigenvalues of the empirical covariance of the 10,000 training anchor features before PCA, using normalization . Eigenvalues range from 0.004730 to 16.451719. Right: covariance eigenvalues after applying the saved 32-dimensional PCA/whitening transform to the same training features. Their maximum absolute deviation from one is , confirming the numerical training-coordinate normalization.
Figure 7: Complete CelebA pretext-bank comparison. Both panels show the eight configured pretexts: expression attributes, hair attributes, joint attributes, grayscale prediction, lower-region prediction, reconstruction, semantic reconstruction, and two-region prediction. The task-aware mixture and the single-target Gaussian baseline with information matrix (“average task”) are also included. Curves show five-seed mean task-averaged measured excess NLL relative to the uncompressed probe, with markers at each evaluated rate. Bands are omitted here for legibility. Colors and markers identify the same objectives across panels. The task-aware curve is dotted orange in both panels, and the average-task reference is dashed. The left axis uses nominal Gaussian rate, the right code budget. The different curve shapes reveal objective-dependent changes in relative performance.
Figure 8: Pretext-bank rank instability. Left: Gaussian representations. Right: learned codes. Rows are ordered by the lowest-nonzero-rate rank. Darker cells indicate better ranks, and outlines identify column winners including ties. Zero rate and the task-aware/average-task references are excluded. Gaussian endpoints have , with 9 reversed pairs out of 28 comparable pairs and a maximum rank change of 6 places. Learned-code endpoints have , with 13 reversed pairs out of 28, and a maximum rank change of 4 places. Thus a pretext’s relative performance changes with nominal rate.
Figure 9: CelebA per-attribute Gaussian behavior. Nominal rates are bits. Panels show Smiling, Eyeglasses, Blond Hair, and Young. The ordinate is measured excess NLL relative to the uncompressed linear probe. Bands are 95% pointwise percentile intervals from 2,000 paired test-image resamples conditional on five fitted seeds. At rate 0.25 the specialist has lower mean NLL on all four tasks. At rate 32 the generalist is better on Eyeglasses, Blond Hair, and Young, while the specialist remains better on Smiling: the aggregate reversal is not uniform across tasks.
Figure 10: CelebA per-attribute learned-code behavior. Code budgets are bits/image. The same four evaluation tasks, specialist, generalist, and uncompressed-probe baseline are used as in Figure 9 . Bands are 95% pointwise percentile intervals from 2,000 paired training-seed/test-image resamples using five fitted seeds. These panels expose task-level variation behind the aggregate code-budget crossover reported in Table 1 .
Figure 11: Gaussian surrogate fidelity. Ordinary least squares of measured excess NLL on predicted Gaussian distortion, with a fitted intercept, gives across 80 objective/rate pairs. The fitted line is . The black line shows this fit. The dashed gray identity line remains a calibration reference. The score is descriptive and in-sample, not a held-out regression score. Strong linear association coexists with a non-unit slope, so it does not imply identity calibration.
Figure 12: Decoder-capacity sensitivity. Left: Gaussian channel. Right: learned binary codes. Linear and 64-hidden-unit MLP probes evaluate the same representations in each mechanism. The ordinate is generalist minus specialist excess NLL. The zero line denotes equal performance. Bands in both panels use 2,000 paired test-image resamples conditional on the five fitted seeds, rather than additional seed resampling. Gaussian crossings are 16.8182 and 19.2453 nominal bits for linear and MLP probes, respectively. The code-panel crossings are 21.2035 and 19.9533 bits/image. The two rate axes represent different budgets and are not shared.
| Quantity | Value | Units / interpretation |
|---|---|---|
| 0.414760 | fitted task SNR | |
| 0.737435 | fitted task SNR | |
| 0.884189 | fitted task SNR | |
| 0.375638 | fitted task SNR | |
| 32 | PCA dimension | |
| 5 | specialist rank |
Table 1: Fitted theory-to-data quantities. Values concern the five-source-attribute specialist and reconstruction generalist. The first four rows are per-task largest eigenvalues. and are fitted proxy quantities. The 16.2447-bit threshold is a sufficient isotropic-reference bound, not an equality prediction for the empirical crossover. The fitted specialist is anisotropic on its span. The saved strict isotropy diagnostic also rejects exact full-space isotropy for the reconstruction components. Consequently, the reference threshold is not asserted as a certificate for these fitted components. Crossover intervals condition on a first positive-to-negative crossing. All 2,000 resamples cross in each mechanism, with no multiple crossings.
| Gap | ||||
|---|---|---|---|---|
| 1.75 | 1.85 | 0.00 | 1.84 | 1.85 |
| 2.16 | 2.27 | 0.00 | 2.27 | 2.27 |
| 5.10 | 2.71 | 2.62 | 0.09 | 5.34 |
| 6.19 | 2.96 | 3.47 | -0.52 | 6.43 |
| 7.15 | 3.29 | 4.03 | -0.75 | 7.32 |
| 8.38 | 3.74 | 4.62 | -0.87 | 8.36 |
Table 2: dSprites factor recovery. The stated results are for , averaged over seeds 0,1, and 2.
| Run | Seed | Crossover | Lower grid rate | Upper grid rate |
|---|---|---|---|---|
| Primary | 0 | 5.2337 | 5.0378 | 6.1546 |
| Primary | 1 | 5.2432 | 5.1039 | 6.2125 |
| Primary | 2 | 5.3144 | 5.1595 | 6.2178 |
| Role swap | 0 | 5.1481 | 2.1223 | 5.1627 |
| Role swap | 1 | 5.1855 | 5.1095 | 6.1823 |
| Role swap | 2 | 5.2495 | 5.1796 | 6.1687 |
Table 3: Per-seed dSprites crossover audit. Rates are measured variational KL bounds in bits/image. The primary gap crosses positive-to-negative. The role-swap physical-factor gap crosses negative-to-positive. Every seed has one such crossing. Brackets are observed neighboring grid rates, not uncertainty intervals.
Figure 13: dSprites recoverability accounting. The 42 mixture runs compare the sum of per-run MSE indices with the variational KL bound over the full 14-setting grid. Maximum positive index-sum excess is 0.2470 bits. The sum is a recovery diagnostic, not a quantitative decomposition of the rate bound. The identity line is retained as a visual reference. No linear relationship or information-theoretic equality is assumed for these discrete factors.
Figure 14: dSprites role-swap control. Vertical position is common/moderate and horizontal position rare/precise. Individual points and thin gap curves show three seeds. Central curves show means and error bars give seed min–max ranges at fixed . The horizontal-minus-vertical gap crosses negative-to-positive at 5.1938 bits/image. All three seeds reverse. The displayed view covers 0–10 bits/image. The dotted reference is from the Gaussian surrogate, not an exact dSprites prediction. The changed physical-factor ordering supports a task-role effect rather than an explanation based solely on one fixed spatial direction being harder to recover.
Figure 15: Recovery under the full dSprites bank. The two panels show horizontal and vertical recoverability versus measured KL for both specialists, reconstruction, and mixture. Lines connect three-seed means at each of 14 settings. No confidence bands are drawn. At each objective’s highest mean achieved rate, the horizontal specialist has index pair , the vertical specialist , reconstruction , and mixture . These endpoints occur at different achieved rates and are not a matched-rate performance comparison. They illustrate that specialists primarily recover their own factor, while reconstruction and mixture retain both.
More in this topic
Representation Learning
Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.
See all 799 papers