Source-linked AI summary
Optimal Multitask Representations Under a Rate Constraint: An Information Bottleneck Analysis
Nicholas Stiles, Gregory W. Wornell
TL;DR
Under a limited information budget, it may be impossible to preserve all information needed for downstream tasks, leaving open which features should receive capacity. The paper analyzes this allocation with a multitask information bottleneck and shows that optimal priorities and pretext rankings can change with rate.
Problem
When a representation budget makes full sufficiency impossible, it is unclear which information should receive the available capacity across downstream tasks.
Source: §1 · p. 1Method
The paper uses a mutual-information-constrained multitask bottleneck and characterizes optimal Gaussian representations alongside fixed Gaussian prediction pretexts.
Source: §1 · pp. 1–2Results
Some task families make every fixed Gaussian pretext incur arbitrarily large relative regret; for M=256, distortion exceeds 5.67 times the optimum at at least one of R1=8 or R2=32.
Source: §1 · p. 2, §4.1 · p. 5Takeaways & Limitations
Representation objectives should be evaluated across the intended deployment bitrate range, especially when compression is for transmission or storage.
Source: §7 · p. 8The experiments use a fixed backbone for compression-objective selection and a deliberately constructed synthetic task pair, so they do not establish that the mechanism is common in natural task suites.
Source: §6.3 · p. 8
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper asks how to allocate a limited information budget among downstream-relevant features when a fully sufficient representation may be impossible. It analyzes this choice through a multitask information bottleneck and finds that both optimal allocation and pretext rankings can depend on the rate.
1 INTRODUCTION
- The central question is which information should receive scarce representation capacity when preserving a fully sufficient representation is impossible.The budget measures information retained about each observation, not parameter count or computational cost.Source: §1 · p. 1
- The analysis complements prior work on task-aware compression, view selection, and representation generalization by asking whether a surrogate allocates capacity appropriately across budgets.It distinguishes this population-level allocation question from the statistical benefits of learning shared structure across tasks.Source: §1 · p. 2
- Task directions can reverse in allocation priority as rate changes, so one fixed Gaussian prediction objective may incur substantial relative regret while an average-information objective has a bounded guarantee.The guarantee is controlled by maximum task SNR, with a second-order regret bound at low SNR.Source: §1 · p. 2
- CelebA compression experiments show rate-dependent pretext rankings, while dSprites experiments find a reversal in relative factor recoverability near the rate predicted by the Gaussian analysis.The dSprites experiment uses task-aware end-to-end training, but does not establish empirical optimality.Source: §1 · p. 2
2 FROM SINGLE-TASK TO MULTITASK INFORMATION BOTTLENECK
The paper formulates representation quality as task-averaged excess log-loss under a mutual-information rate constraint. It moves from single-target IB to shared multitask representations, then uses a Gaussian model to obtain an exact benchmark.
2.1 Single-Task IB: What Does a Pretext Preserve?
- For a single Gaussian target, the optimal representation allocates information along a fixed eigenbasis, and increasing rate cannot reverse the allocation ordering set by the target eigenvalues.The rate constraint limits retained information, and the minimized quantity is the increase in optimal predictive log-loss.Source: §2.1 · p. 2
2.2 A Shared Representation for Multiple Tasks
- A shared encoder with task-dependent decoders is evaluated by the smallest average compression-induced excess log-loss at rate R; a fixed pretext substitutes its target for the evaluation tasks.The task-indexed target yields the weighted sum of task information, though it is generally not jointly Gaussian with X.Source: §2.2 · p. 3
2.3 Gaussian Task Model and Exact Benchmark
- For Gaussian tasks, Gaussian extremality reduces the multitask problem to a conditional-covariance optimization whose unique minimizer exhausts the rate budget, providing an exact benchmark.The stated result assumes H ≻ 0; it also holds when restricted to the task-relevant subspace.Source: §2.3 · p. 3
3 TASK KNOWLEDGE AND PRETEXT CHOICE
The section compares a specialist that concentrates on a task-relevant subspace with a generalist that preserves information across the input. Under stated relevance and omission conditions, their relative performance reverses as the rate grows.
3.1 When Specialists Beat a Generalist
- The comparison models the specialist as allocating rate across a strict subspace and the generalist as allocating it isotropically across the input.The objectives represent focused prediction versus broader input preservation, respectively.Source: §3.1 · pp. 3–4, §3 · p. 3
- A specialist with above-average evaluation relevance per dimension and omitted evaluation-relevant information beats the generalist at sufficiently low rates but loses at sufficiently high rates.The two distortion curves cross at a finite positive rate; the specialist's initial advantage reflects greater downstream benefit per bit, while limited coverage creates a loss floor.Source: §3.1 · p. 4
- The crossover shows that transferability of these pretexts depends on bitrate, but does not establish that every fixed Gaussian pretext is suboptimal for multitask representation.The section uses this distinction to motivate a stronger analysis of fixed-pretext regret.Source: §3.1 · p. 4
4 REVERSALS IN OPTIMAL ALLOCATION PRIORITIES
Optimal allocation priorities can reverse as the information budget increases, even when the task suite is fixed. The section shows how this reversal affects fixed Gaussian pretexts and characterizes a rate-dependent pretext that matches the multitask optimum.
4 REVERSALS IN OPTIMAL ALLOCATION PRIORITIES
- The section’s central result is that fixed-rate-independent priorities may misallocate information as the budget changes, whereas a matched pretext can track the multitask optimum.The reversal can occur with the task suite held fixed.Source: §4 · p. 4, §4.2 · p. 5
4.1 Example: task frequency versus precision
- At low rates, task frequency sets a direction’s marginal value; at high rates, residual signal strength can reverse the relative priority of the two directions.Neither allocation decreases with total rate, and active supports remain nested; it is their relative priority that reverses.Source: §4.1 · p. 4
- For M = 256, every fixed Gaussian pretext has distortion exceeding 5.67 times the optimum at R1 = 8 or R2 = 32.The rate–distortion gap and optimal allocations are shown in Figure 2, with the exact reversal near R* ≈ 24.Source: §4.1 · p. 5, Figure 2 · p. 5
4.2 Rate-Dependence of Optimal Pretext Geometry
- For any task family and rate, a rate-dependent single-task Gaussian information bottleneck can share the multitask solutions, and its matched covariance is a global optimum.The matching rule depends on rate, and the inducing pretext need not be unique.Source: §4.2 · p. 5
- The optimal pretext’s directional priorities change with rate: strong task signals saturate at low rates, while task precision enters linearly at high rates.This changes the relative priority of directions across budgets.Source: §4.2 · pp. 5–6
5 WHEN A FIXED PRETEXT SUFFICES
The section asks when one Gaussian pretext fixed across rates can approximate the multitask information bottleneck. It shows that bounded task SNR yields a uniform relative-regret guarantee for a pretext based on the weighted average task matrix.
5 WHEN A FIXED PRETEXT SUFFICES
- Under bounded task SNR, the pretext using the weighted average task matrix has a regret bound uniform over all finite rates.The fixed pretext requires knowledge of the task suite and bounds the cost of approximating the multitask solution by its high-rate limit.Source: §5 · p. 6
- For L = 2, S(3) ≈ 1.160 bounds the regret of the optimal fixed Gaussian pretext by 16%.The bound is independent of task count and input dimension, and applies even when optimal directions rotate.Source: §5 · p. 6
- At low SNR, the guaranteed relative loss is second order in task strength because weak Gaussian task channels have nearly linear distortion in K.Strong channels can depart substantially from this approximation through saturation.Source: §5 · p. 6
6 EXPERIMENTS
The experiments show that compression can reverse which pretext performs best, and that direct multitask training can reverse the relative recoverability of task factors as rate increases. These results are illustrative rather than claims of empirical optimality or prevalence across natural task suites.
6.1 Pretext Selection Under Compression
- The CelebA comparison fixes a pretrained backbone and evaluates four held-out attributes using excess NLL relative to probes on uncompressed features.The primary comparison pairs feature reconstruction with a five-attribute specialist; subsets were fixed before inspecting rate-distortion curves.Source: §6.1 · pp. 6–7
- Across the full bank, 9 of 28 pretext pairs reverse order across Gaussian budgets and 13 of 28 across discrete-code budgets.The corresponding Kendall correlations are τb = 0.357 and 0.07; the primary crossover also persists with an MLP probe.Source: §6.1 · p. 7
- The specialist performs better at low budgets, but the generalist overtakes it as the budget grows, with crossovers near 16.8 bits for Gaussian channels and 21.2 bits for discrete codes.The crossovers had a sign change in 100% of resamples.Source: Figure 3 · p. 7, §6.1 · p. 7
6.2 Allocation Reversal Under Direct Multitask Training
- The dSprites experiment trains a variational-bottleneck encoder on common, moderately precise horizontal-position and rare, highly precise vertical-position tasks.The task weights are (0.75, 0.25), their SNRs are (16, 256), and three seeds are used per setting.Source: §6 · p. 6, §6.2 · p. 7
6.3 Experimental Scope
- The experiments demonstrate distinct reversals: which objective ranks best depends on budget, while which input direction receives priority within the multitask solution also changes with rate.The CelebA study compresses fixed-backbone features, and the dSprites study uses a synthetic, deliberately constructed noisy task pair; neither establishes prevalence in natural task suites.Source: §6.3 · p. 8, §6.2 · p. 8
6.2 Allocation Reversal Under Direct Multitask Training
- Both factors become more recoverable as rate increases, but their ordering reverses near 5.3 bits in all three seeds.The Gaussian construction predicts a reversal at R⋆≈5.30 bits, but the agreement is not a calibrated rate prediction.Source: Figure 4 · p. 8, §6.2 · p. 8
7 CONCLUDING REMARKS
The concluding remarks argue that representation objectives should be evaluated across the intended deployment bitrate range. The analysis links allocation priorities to task frequency and precision, while bounding the cost of a rate-independent task-aware objective under bounded task SNR.
7 CONCLUDING REMARKS
- Representations compressed for transmission or storage should be evaluated across the intended deployment bitrate range.An objective that uses scarce capacity effectively may become limited by the information it omits at higher rates.Source: §7 · p. 8
- Task frequency and precision jointly determine allocation priorities, with rare tasks requiring precise predictions potentially gaining priority as rate grows.For Gaussian tasks with bounded SNR, the average task information matrix provides a principled rate-independent objective with a uniform relative-regret bound.Source: §7 · p. 8
- The experiments leave unresolved whether these allocation effects are prevalent in broader task suites.Extending the analysis to restricted decoders requires distinguishing information loss from information inaccessible to those decoders; weight quantization additionally requires linking weight precision to retained information.Source: §7 · pp. 8–9
A PROOF OF THE GAUSSIAN REDUCTION
The proof reduces the Gaussian multitask information-bottleneck problem to optimization over a conditional covariance, then establishes the optimizer’s existence, uniqueness, and rate properties. It also derives optimal specialist and reconstruction-generalist covariances.
A.1 Gaussian extremality in Lagrangian form
- Gaussian extremality reduces the weighted linear-Gaussian task problem to a Gaussian test channel parameterized by its residual covariance.The residual covariance is K = Cov(X | Z).Source: §A.1 · p. 11
A.2 The Gaussian covariance parameterization
- Every feasible residual covariance K can be realized by a Gaussian encoder, with its rate determined by −1/2 log2 det K.The construction sets B = (K−1 − I)1/2 and yields Cov(X | Z) = K.Source: §A.2 · pp. 11–12
A.3 Geodesic geometry
- The feasible covariance set is geodesically convex, and the task distortion is strictly geodesically convex under the stated positive-definiteness condition.The rate function is geodesically affine.Source: §A.3 · pp. 12–13
A.4 Existence, uniqueness, and rate exhaustion
- For every finite rate, the matrix program has a unique minimizer; at positive rates it exhausts the rate constraint, while at zero rate K0 = I.Source: §A.4 · pp. 13–14
A.5 Constrained reduction
- The covariance optimizer is globally optimal for the constrained problem, and a Gaussian encoder attains its value even against arbitrary encoders.A supporting multiplier converts the constrained optimizer into a minimizer of a geodesically convex objective.Source: §A.5 · pp. 14–15
A.6 Numerical optimality conditions
- The optimizer’s first-order conditions have a water-filling interpretation: encoded directions share a common normalized marginal value, while unencoded directions are bounded by it.Source: §A.6 · pp. 15–16
B.1 Optimal Specialist and Generalist Covariances
- The specialist allocates rate uniformly over its task subspace, whereas squared-error reconstruction has an isotropic optimal covariance across all dimensions.The specialist covariance is 2−2R/k Ik; the reconstruction covariance is 2−2R/d Id.Source: §B.1 · p. 17
B.2 Proof of Proposition 2
The specialist and reconstruction generalist have opposite distortion advantages at small and large rates when relevant task information is missing from the specialist’s subspace; continuity guarantees a crossing.
- At finite small and large rates, the specialist and generalist have opposite strict distortion advantages, so continuity guarantees at least one positive finite crossing.Source: §B.2 · p. 17
- As R → ∞, specialist distortion approaches Dmiss > 0 while generalist distortion approaches zero.Source: §B.2 · p. 17
- The high-rate advantage applies when task information lies outside the specialist subspace, expressed by tr((I − P)H) > 0.Source: §B.2 · p. 18
C FIXED-PRETEXT REGRET BOUND FOR THE TWO-MODE EXAMPLE
A two-mode task family shows that fixed Gaussian pretexts can be badly mismatched across rates because the optimal allocation ordering reverses.
C FIXED-PRETEXT REGRET BOUND FOR THE TWO-MODE EXAMPLE
- The optimal coordinate ordering reverses across rates, but a fixed pretext cannot change its ordering and is therefore wrong or tied at one of the two rates.Source: C FIXED-PRETEXT REGRET BOUND FOR THE… · p. 19
- For the rank-one tasks, removing off-diagonal covariance preserves distortion while reducing rate, so the regret bound reduces to minimizing on the wrongly ordered side of r = R/2.Source: C FIXED-PRETEXT REGRET BOUND FOR THE… · p. 19
- Every fixed Gaussian pretext can have arbitrarily large worst-rate relative regret, even for commuting tasks with perfectly nested supports.The construction uses increasing signal strengths and unequal task weights; it is an existence result, not a claim about typical data.Source: C FIXED-PRETEXT REGRET BOUND FOR THE… · p. 19
D.2 Proof of Theorem 1
The proof establishes that a rate-dependent single-task pretext can share the multitask optimizer, while clarifying that this matched pretext is not unique and the optimum uses the full rate.
D.2 Proof of Theorem 1
- The downstream-optimal covariance is globally optimal for the matched pretext, and evaluating the downstream objective there gives the same optimum.The matched pretext is task-dependent but is not the unique pretext inducing that covariance.Source: §D.2 · p. 21
D.3 Endpoint behavior of the downstream optimum
- At positive rates, the downstream-optimal covariance exhausts the rate constraint; otherwise scaling it down would improve distortion while remaining feasible.Source: §D.3 · p. 22
E PROOF OF THE BOUNDED-SNR APPROXIMATION
The bounded-SNR proof compares the average-task pretext with the multitask optimum and establishes one rate-uniform approximation. The experiment appendix documents the CelebA and dSprites procedures, while noting that its fixed settings were not broadly optimized.
E PROOF OF THE BOUNDED-SNR APPROXIMATION
- The average-task pretext’s trace distortion is at most S(1 + L) times the multitask optimum, using the same matrix A at every rate.For small L, the bound is second order: ln S(e^x) = x^2/8 + O(x^3).Source: E PROOF OF THE BOUNDED-SNR APPROXIMATION · p. 26
- The comparison applies on the task-relevant subspace; requiring a positive-definite pretext on extra task-null dimensions need not preserve the guarantee.For singular H, the result restricts the input to range(H), or permits a positive-semidefinite pretext that ignores the common task-nullspace.Source: E PROOF OF THE BOUNDED-SNR APPROXIMATION · p. 26
F EXPERIMENT DETAILS
- The recorded experiments use shared architectures, optimization settings, data partitions, and rate grids across pretexts, but do not establish that these fixed settings are optimal.Test losses are used for reporting, not checkpoint or hyperparameter selection.Source: F EXPERIMENT DETAILS · p. 27
F.1 CelebA experiment
- CelebA compares a five-attribute specialist with feature reconstruction on four disjoint evaluation attributes, alongside a broader bank of pretexts and two reference objectives.The generalist reconstructs anchor features, and the average-task baseline uses the mean evaluation-task information matrix.Source: §F.1 · p. 27
- CelebA compression cost is measured as task-averaged excess NLL against uncompressed probes, with uncertainty intervals that do not include new-dataset or model-selection uncertainty.Gaussian intervals resample test images; learned-code intervals resample test images and fitted seeds.Source: §F.1 · p. 29
- Individual-attribute evaluations show that the aggregate crossover need not occur uniformly across tasks; separate analyses assess Gaussian surrogate distortion against measured excess NLL and sensitivity to replacing linear probes with MLP probes.Source: §F.1 · p. 29
F.2 dSprites experiment
The dSprites experiment tests whether rate-dependent allocation reversals also appear in learned representations, including after swapping task roles. The appendix also details the Gaussian covariance solver and limits of its rate and optimization diagnostics.
F.2 dSprites experiment
- The dSprites tasks are horizontal and vertical position, weighted 3/4 and 1/4, with independent noise variances 1/16 and 1/256; other rendering factors remain nuisance variation.The factors are standardized using their known uniform 32-level mean and standard deviation.Source: §F.2 · p. 29
- The variational KL is an upper bound on mutual information, while the per-factor recoverability index is a diagnostic transformation of MSE, not an allocation of bits.The index equals −1/2 log2(1 − R2) only when the empirical target variance is one; the three-seed ranges are not confidence intervals.Source: §F.2 · p. 31
- 5.2638 and 5.1938 bits/image are the primary and role-swapped mean-curve crossover rates, respectively; all three seeds reverse in each run.The role swap changes task roles while retaining task-index weights and noise variances; the Gaussian surrogate reference is 5.2996 bits.Source: §F.2 · p. 31, Figure 14 · p. 39
- For the role-swap audit, crossover rates come from piecewise-linear interpolation on the observed grid without extrapolation; neighboring-rate brackets are not uncertainty intervals.Each seed has exactly one expected-direction crossing in both the primary and swapped runs.Source: §F.2 · p. 31, Table 3 · p. 31
F.3 Gaussian covariance optimization
- The Gaussian solver reduces optimization to the numerically relevant task span and uses log-covariance coordinates with projected steps that enforce the rate constraint.The reported residual-based gap is an a posteriori certificate for the surrogate program, not a statistical error bound for learned image representations.Source: §F.3 · pp. 32–33
- The finite-test-set entropy estimator tops out at 10.96578 bits, so it cannot resolve a 16- or 32-bit population entropy.The ceiling equals log2(2000); sparse codeword occupancy can also bias plug-in entropy downward.Source: Figure 5 · p. 34