PaperlayerView exact revision on arXiv

Source-linked AI summary

UniTAC: Universal Task-Aware Compression via Weighted Distortion Measures

Homa Esfahanizadeh, Matin Mortaheb, Jinfeng Du, Harish Viswanathan

arXiv:2608.16696v1cs.LGcs.AIcs.ITcs.MM

TL;DR

Physical AI systems need compression that adapts as downstream tasks change, but retraining task-specific codecs is often infeasible in the field. UniTAC conditions one learned codec on a runtime importance vector and, at ≈0.034 bpp, retains 91.4% mouth-task accuracy versus 76.9% for a non-semantic codec.

  • Problem

    Physical AI tasks evolve under bandwidth-limited communication, making bespoke codec design and retraining for each task operationally expensive and often infeasible.

  • Method

    UniTAC trains one codec over importance vectors that condition both encoder and decoder, enabling runtime task re-targeting without retraining.

  • Results

    At comparable rates, task-matched conditioning improves task-region fidelity and downstream accuracy across tasks; on the mouth task, 91.4% accuracy versus 76.9% for the non-semantic codec.

  • Takeaways & Limitations

    A single shared backbone can span universal to task-specialized operation by swapping its weight vector rather than retraining per task.

  • Takeaways & Limitations

    The universality gap and the quality of the supplied importance map remain open questions for UniTAC.

Abstract

from arXiv · show

Physical AI systems such as autonomous vehicles and robots rely on timely exchange of high-dimensional sensory signals under tight bandwidth, latency, and energy budgets. Because the task driving downstream decisions evolves over time, a task-specific codec is brittle and retraining one per task is infeasible in the field. We propose UniTAC, a single learned image codec spanning universal (task-agnostic) to task-specialized operation, re-targeted at runtime without retraining. The task is abstracted as a per-component importance vector, derived, e.g., from gradient attribution of any downstream model, and transmitted as low-overhead side information that conditions both encoder and decoder. Trained once over a broad, randomized family of such vectors against weighted-reconstruction distortion, UniTAC keeps a fixed backbone and a single human-viewable reconstruction whose fidelity is steered to the active task by swapping the injected vector. We analyze the underlying weighted rate-distortion problem, characterizing when a diagonal weighted distortion is task-consistent and how weights relate to task sensitivity. Guided by this, we design a Vision Transformer (ViT) codec whose token-level conditioning natively realizes this weight-driven code. On a localized task at 0.034 bpp, a single UniTAC model reaches 91.4% accuracy, only 1.9% below a task-based codec (93.3%) and above universal codecs (76.9%).

I. INTRODUCTION … B. Task-oriented and semantic communication

UniTAC addresses changing Physical AI tasks with one weight-conditioned learned codec that redirects compression toward task-important components at runtime. It unifies universal and task-aware compression while avoiding task-specific retraining and heuristic region selection.

  • I. INTRODUCTION: Physical AI sensors produce high-rate, redundant streams, while downstream decisions need unequal fidelity for regions such as obstacles, grasp points, and safety-critical agents.Communication and control operate under tight constraints, making selective allocation of bits important.
  • I. INTRODUCTION: Generic codecs optimize PSNR or MS-SSIM and allocate bits uniformly, regardless of the receiver’s task.This task-agnostic behavior motivates task-aware compression.
  • I. INTRODUCTION: Changing tasks make bespoke codec design and retraining operationally expensive and often infeasible in the field.A robot may shift among navigation, mapping, and interaction.
  • I. INTRODUCTION: UniTAC represents task saliency as a per-component importance vector W that conditions both encoder and decoder, enabling runtime specialization of one trained compressor.The compressor is trained once over a broad family of importance vectors, with weighted distortion sum_i w_iD_i(x_i, x̂_i).
  • A. Universal image and feature compression: Classical and learned image codecs, including JPEG AI, are reconstruction-centric and inherently task-agnostic, allocating rate uniformly across the image.Their objectives use generic distortion measures such as PSNR and MS-SSIM.
  • B. Task-oriented and semantic communication: Task-oriented communication instead optimizes transmission for a downstream objective, with representative feature codecs jointly training extraction, task prediction, and entropy modeling.These single-task codecs combine task loss with a rate penalty and target task-sufficient representations.
  • B. Task-oriented and semantic communication: Information-bottleneck methods characterize rate versus retained information about a relevant target, but typically specialize representations to one fixed task without source reconstruction or runtime re-specialization.This limitation distinguishes UniTAC’s runtime task conditioning from fixed-task relevance optimization.
  • B. Task-oriented and semantic communication: Semantic and task-aware JSCC allocates unequal channel protection according to the semantic or task importance of underlying content rather than protecting all bits equally.JSCC maps source content directly to channel symbols instead of using separate compression and channel codes.

C. Multi-task and adaptive codecs · D. Task-to-importance mapping · III. A WEIGHTED-DISTORTION FRAMEWORK FOR TASK-CONDITIONED COMPRESSION

UniTAC extends multi-task compression with runtime importance-vector conditioning, enabling one human-viewable codec to retarget arbitrary tasks without retraining. Its weighted-distortion framework formalizes when component-wise importance weights align rate-constrained compression with task fidelity.

  • C. Multi-task and adaptive codecs: Prompt-based coding similarly uses task-driven importance maps and randomized training to support multiple tasks, whereas other codecs require task-specific side paths or fixed task hierarchies.The comparison distinguishes runtime map conditioning from predefined task sets and layered hierarchies.
  • C. Multi-task and adaptive codecs: UniTAC receives a transmitted importance vector as low-overhead side information, conditions both encoder and decoder, and retargets one universal codec to any task without retraining.The reconstructed image remains human-viewable and general-purpose while prioritizing fidelity for the targeted task.
  • D. Task-to-importance mapping: Explainability methods derive per-input importance or saliency maps through gradient- or activation-based, perturbation- or mask-based, and path-attribution approaches.These methods identify which input parts a trained model relies upon.
  • III. A WEIGHTED-DISTORTION FRAMEWORK FOR TASK-CONDITIONED COMPRESSION: The framework models a task as f: R^n → R^m and takes task fidelity to be mean squared error in task space.The original and reconstructed source vectors are denoted X and X̂.
  • III. A WEIGHTED-DISTORTION FRAMEWORK FOR TASK-CONDITIONED COMPRESSION: UniTAC optimizes separable weighted distortion, assigning each component a nonnegative distortion and a task-importance weight that may depend on the sample.Squared error is given as an example per-component distortion.
  • III. A WEIGHTED-DISTORTION FRAMEWORK FOR TASK-CONDITIONED COMPRESSION: Weighted distortion is task-consistent when minimizing it under a rate constraint also minimizes the task loss, which is the framework’s central criterion.The paper studies how to choose weights so rate-constrained compression reduces the true task distortion.
  • III. A WEIGHTED-DISTORTION FRAMEWORK FOR TASK-CONDITIONED COMPRESSION: The proposed training strategy exposes one fixed compressor backbone to a broad family of importance vectors, allowing runtime bit reallocation as the active task changes without retraining.Weights can be re-estimated through reinforcement learning, gradient or saliency attribution, or a user prompt.

IV. THEORETICAL ANALYSIS: TOWARD TASK-CONSISTENT WEIGHTED DISTORTION

This section analyzes when separable weighted distortion faithfully represents task loss and what weights that correspondence requires. It focuses on quadratic per-component distortions.

  • Task-consistent weighted distortion: The analysis characterizes when separable weighted distortion is a faithful surrogate for task loss and determines the corresponding weights.The framework uses quadratic per-component distortions, D_i(x_i, x̂_i) = ||x_i − x̂_i||^2.

A. From task loss to weighted distortion via sensitivity · B. Structural constraints: symmetry and irrelevance

The section derives task-aware weighted distortion from local task sensitivity and identifies when separable diagonal weights remain task-consistent. It then shows that task symmetry constrains weights symmetrically, while task-irrelevant components must receive zero weight under stated rate conditions.

  • A. From task loss to weighted distortion via sensitivity: A first-order local approximation of the task map provides the basis for translating task loss into distortion in the source space.
  • A. From task loss to weighted distortion via sensitivity: The ideal source-space distortion is a generally non-separable quadratic form weighted by G(X) = Jf(X)^T Jf(X).
  • A. From task loss to weighted distortion via sensitivity: A separable diagonal surrogate is faithful exactly when cross terms vanish per sample or in expectation.
  • A. From task loss to weighted distortion via sensitivity: This condition holds through task orthogonality, where distinct components have no first-order interaction, or through sufficiently uncorrelated errors.
  • B. Structural constraints: symmetry and irrelevance: Task structure constrains admissible weights: symmetry must be shared by the weights, while irrelevant components must receive zero weight.
  • B. Structural constraints: symmetry and irrelevance: For exchangeable sources and permutation-invariant tasks, task-consistent sample-dependent weights must be permutation-equivariant.
  • B. Structural constraints: symmetry and irrelevance: Sample-independent weights under the same symmetry assumptions must be equal across all components, w1 = · · · = wn.
  • B. Structural constraints: symmetry and irrelevance: Under positive variance for ignored components and a strictly decreasing optimal task loss with rate, task consistency requires wi(X) = 0 for every ignored component.

C. Linear tasks: exact reduction to weighted MSE · D. Nonlinear tasks: Jacobian sensitivity

For linear tasks, task loss exactly becomes weighted MSE under uncorrelated errors, with weights determined by task coefficients. For nonlinear tasks, Jacobian-based weighted distortion is approximate, while integrated gradients provide a practical per-sample sensitivity surrogate.

  • C. Linear tasks: exact reduction to weighted MSE: Linear task loss exactly reduces to weighted MSE when reconstruction errors across components are uncorrelated, yielding an optimal separable distortion.The result is stated for a single-output linear function and extends to multiple linear outputs.
  • C. Linear tasks: exact reduction to weighted MSE: For Gaussian sources with independent components, reverse water-filling gives the exact optimal allocation; for non-Gaussian sources, it is near-optimal.The closed-form allocation relies on source independence and Gaussianity for exactness.
  • C. Linear tasks: exact reduction to weighted MSE: For multiple linear outputs f(X) = SX, optimal separable weights are the squared column norms of S under uncorrelated errors.This generalizes the single-output weighting rule to vector-valued linear tasks.
  • C. Linear tasks: exact reduction to weighted MSE: Correlated source components or quantization errors require decorrelation, such as KLT or whitening, to recover separable weighted distortion.Without decorrelation, separable weighted distortion is only an approximation.
  • D. Nonlinear tasks: Jacobian sensitivity: Nonlinear tasks lose the exact reduction because G(X) = Jf(X)^T Jf(X) varies with X and is generally non-diagonal.The resulting separable surrogate discards off-diagonal cross terms involving inter-component error correlations.
  • D. Nonlinear tasks: Jacobian sensitivity: The approximation gap δ has two sources: first-order linearization error and discarded off-diagonal cross terms.These are the stated conditions governing δ-approximate task consistency.
  • D. Nonlinear tasks: Jacobian sensitivity: For nonlinear tasks, the exact coefficient acting on xi−x̂i is the path-averaged gradient along the segment from x̂ to x, not the point gradient alone.The path-average coefficient can be applied per output for vector-valued tasks.
  • D. Nonlinear tasks: Jacobian sensitivity: Integrated-gradient weights wi(X) = ḡi^2 yield a separable bound controlling task loss, using a fixed baseline X0 to obtain a design-time per-image surrogate.The baseline can be task-appropriate, such as a black input, and the resulting map is used in experiments.

V. CODEC DESIGN: VIT-BASED UNITAC · A. Tokenization and transformer backbone

UniTAC realizes task-aware rate allocation in a single ViT-based codec by conditioning encoder and decoder attention on an importance map. Its tokenized transformer backbone combines patch-aligned weighted distortion, local attention for efficient high-resolution processing, and a mirrored decoder for bounded image reconstruction.

  • V. CODEC DESIGN: VIT-BASED UNITAC: UniTAC instantiates task-aware rate allocation in a ViT codec whose token-level conditioning directly consumes the importance vector W without retraining.The codec is designed to realize a family of rate allocations across components according to their importance.
  • V. CODEC DESIGN: VIT-BASED UNITAC: A two-stage transformer autoencoder maps images to compact latents and back, while attention injection steers capacity toward components marked important by W.A hyperprior entropy model converts the latent representation into a bitstream and provides a differentiable rate estimate.
  • A. Tokenization and transformer backbone: The image is split into nonoverlapping 4×4 patches by a strided 2D convolution, producing tokens with embedding width C0 = 96.The importance map aligns with the token grid, so weighted distortion is measured on spatial patches.
  • A. Tokenization and transformer backbone: Each transformer block uses prenorm, residual connections, a two-layer GELU MLP with hidden dimension 4× the token width, and 4-head attention; each stage stacks D = 4 blocks.These architectural conventions define the repeated transformer stages throughout the codec.
  • A. Tokenization and transformer backbone: The encoder begins with neighborhood attention that restricts tokens to a local k×k window, reducing high-resolution attention complexity to linear in token count.The local window is k = 3, with learnable relative-position bias and cyclically increasing dilation indicated in the design.
  • A. Tokenization and transformer backbone: The decoder reverses the pipeline with D SGAT blocks, spatial upsampling, D NAT refinement blocks, and a sigmoid producing a bounded reconstruction ˆX ∈[0, 1]3×H×W.A 2D sinusoidal positional encoding supplies spatial reference at the decoder input.

B. Token-level weight conditioning · C. Entropy coding and rate estimation

UniTAC conditions ViT tokens with a low-resolution importance map, steering local and global attention so capacity concentrates on task-relevant regions. Its hyperprior entropy model codes latents and treats the fixed importance-map transmission as negligible common overhead.

  • B. Token-level weight conditioning: The importance map is a normalized G × G grid with G = 16, resized bilinearly to each stage’s token resolution.It is supplied at low resolution and shared across stages.
  • B. Token-level weight conditioning: Importance conditioning lets important tokens retain detail while unimportant tokens borrow detail from important regions, concentrating latent capacity where the task requires it.The two complementary mechanisms operate across token stages.
  • B. Token-level weight conditioning: In local neighborhood attention, an importance bonus s w_j biases each token toward important neighbors, while a second mechanism controls each token’s update strength.The bonus is added to standard query–key attention scores within each k × k window.
  • B. Token-level weight conditioning: At the bottleneck, each query samples an importance-weighted memory of T = 24 tokens, with cosine-annealed temperature shifting from exploration toward highest-importance selection.A learnable gate mixes each token’s self-representation with memory content, followed by full self-attention for global consistency.
  • C. Entropy coding and rate estimation: The latent Z is hyperprior-coded: a hyper-encoder produces Y under a factorized prior, and a hyper-decoder predicts Gaussian parameters (µ, σ) for Z.Y is transmitted as side information; training uses additive uniform noise, inference uses rounding, and arithmetic coding produces the bitstream.
  • C. Entropy coding and rate estimation: The rate term and reported test rate are derived from the conditional model p(Ẑ | µ, σ) and factorized prior p(Ŷ).These distributions define the entropy-coding rate used in optimization and evaluation.
  • C. Entropy coding and rate estimation: The G × G importance map is transmitted to the decoder as side information, but its fixed, resolution-independent overhead is omitted from the objective and reported bpp.With G = 16 and identical overhead across codecs, it contributes only a negligible common additive constant.

D. Training objective and randomized importance maps · VI. EXPERIMENTAL EVALUATION · A. Experimental setup

UniTAC is trained with randomized, task-weighted rate–distortion objectives so one ViT backbone can be specialized at test time through importance-map conditioning. Experiments evaluate this approach across gender and mouth classification using reconstruction, semantic-fidelity, and downstream-accuracy metrics against universal and task-specific codecs.

  • D. Training objective and randomized importance maps: The codec minimizes a rate–distortion objective combining estimated rate with normalized separable weighted distortion.The per-patch squared error uses the same importance map W that conditions the network, while λ trades rate against task fidelity.
  • D. Training objective and randomized importance maps: Each training image receives a freshly sampled importance map from synthetic Gaussian-blob mixtures spanning peaky-to-diffuse patterns and varied count, location, and scale.The maps are mean-normalized; brighter regions indicate larger wi.
  • D. Training objective and randomized importance maps: A single UniTAC backbone shifts fidelity toward task-important regions at nearly equal rates by swapping gender or mouth importance maps without retraining.The demonstrated reconstructions use one backbone and λ = 0.10.
  • VI. EXPERIMENTAL EVALUATION: The evaluation tests whether one weight-conditioned UniTAC backbone can match task-specific codecs while surpassing task-agnostic codecs across rate-dependent signal and task fidelity.Reported measures are PSNR, semantic PSNR, and classification accuracy.
  • A. Experimental setup: All codecs train on AffectNet and evaluate on CelebA, with square crops resized from 64 × 64 to 256 × 256 in 32-pixel steps.CelebA training images are used only for downstream classifiers and attribute-based importance maps.
  • A. Experimental setup: The tasks are global gender and localized mouth-state classification, whose clean-image ResNet-18 classifiers reach 98.5% and 94.3% Top-1 accuracy, respectively.The classifiers are fine-tuned from ImageNet-pretrained ResNet-18 models.
  • A. Experimental setup: Importance maps come from eight-step integrated gradients, then undergo absolute channel averaging, grid pooling, flooring at 0.02, and mean normalization.The same map conditions the codec and defines task-weighted fidelity.
  • A. Experimental setup: Comparisons include UniTAC, an evenly allocating non-semantic codec, and separately trained task-based codecs, all sharing a ViT architecture across rate–distortion trade-offs.Metrics versus measured bits per pixel are overall PSNR, semantic PSNR, and downstream Top-1 accuracy.

B. Rate–distortion: overall and semantic PSNR · C. Downstream task accuracy · VII. DISCUSSION AND CONCLUSION

UniTAC preserves generic reconstruction quality while substantially improving task-relevant semantic fidelity and downstream accuracy through runtime weight-vector conditioning. Across mouth and gender tasks, one shared backbone approaches dedicated task-based performance without retraining per task.

  • B. Rate–distortion: overall and semantic PSNR: UniTAC under a uniform map tracks the non-semantic codec on overall PSNR, indicating weight conditioning preserves generic reconstruction quality without a specified task.This comparison is reported on the overall-PSNR rate–distortion curve.
  • B. Rate–distortion: overall and semantic PSNR: Matched importance maps substantially improve semantic PSNR over both uniform-map and non-semantic operating points at comparable rates.The gain is measured under task-specific importance maps, including synthetic top-left ROI, mouth, and gender maps.
  • C. Downstream task accuracy: 91.4% mouth-task accuracy at ≈0.034 bpp exceeds 76.9% for the non-semantic codec at the same rate and 71.3% under a uniform map at 0.042 bpp.The task-matched map preserves downstream accuracy substantially better than both alternatives.
  • C. Downstream task accuracy: 92.2% gender-task accuracy at ≈0.043 bpp exceeds 85.3% under a uniform map using the same backbone re-conditioned on the gender map.The reported comparison is at a comparable rate.
  • C. Downstream task accuracy: A single shared model closely matches the dedicated per-task codec at equal rate.This result is reported for downstream task performance across the evaluated tasks.
  • VII. DISCUSSION AND CONCLUSION: UniTAC spans universal to task-specialized operation by swapping the injected weight vector W at runtime instead of retraining a codec per task.The same backbone matches the non-semantic codec on overall PSNR, improves task-region semantic PSNR and downstream accuracy, and approaches the per-task task-based upper baseline.

APPENDIX A … C. Theorem 1 (exact reduction to weighted MSE)

The appendix establishes structural conditions for task-consistent weighted distortion: symmetry requires equal weights across exchangeable coordinates, irrelevance requires zero weights, and the theorem begins an exact weighted-MSE reduction under uncorrelated errors.

  • A. Proposition 1 (symmetry forces symmetric weights): Permutation invariance and exchangeability force sample-independent weights to be equal across coordinates.Otherwise, permuting two coordinates preserves task loss and rate but changes weighted distortion, contradicting task consistency.
  • B. Proposition 2 (irrelevance forces zero weights): A coordinate with positive weight on a positive-probability set must receive positive rate when it has positive variance.Improving its reconstruction strictly reduces weighted distortion, so the rate-constrained minimizer spends information on that coordinate.
  • B. Proposition 2 (irrelevance forces zero weights): Replacing an ignored coordinate by a constant lowers mutual information without changing task loss.The replacement is a deterministic function of the reconstruction and therefore uses no more rate.
  • B. Proposition 2 (irrelevance forces zero weights): If the task loss is strictly decreasing at rate R, the freed rate improves task performance, contradicting task consistency.Therefore, all weights on task-irrelevant coordinates must be zero almost surely.
  • C. Theorem 1 (exact reduction to weighted MSE): Theorem 1 starts from the linear task relation f(X) − f(X̂) = A^T(X − X̂) = A^T E.This expresses task-output error through the reconstruction error vector E.
  • C. Theorem 1 (exact reduction to weighted MSE): Under the uncorrelated-error assumption, the theorem proceeds toward an exact reduction of task loss to weighted mean-squared error.The supplied proof passage introduces the assumption but does not include the resulting expression.

D. Proposition 3 (multiple linear outputs) · 2. Therefore,

Proposition 3 expresses the task loss for multiple linear outputs as a trace involving G ≜ S^T S. Under uncorrelated errors, only the diagonal entries of G affect the loss.

  • D. Proposition 3 (multiple linear outputs): The proof defines G ≜ S^T S and writes the task loss as L_task = E[E^T G E].This reformulates the multiple-output task loss using the output-transform matrix.
  • D. Proposition 3 (multiple linear outputs): Because a scalar equals its trace, the loss can be expressed in trace form.The proof then applies the cyclic property of the trace.
  • D. Proposition 3 (multiple linear outputs): The cyclic property of the trace enables rearrangement of the factors in the task-loss expression.This is the algebraic step used to expose the relevant matrix terms.
  • D. Proposition 3 (multiple linear outputs): Under the uncorrelated-error assumption, the error covariance E[E E^T] is diagonal.The assumption removes off-diagonal error-correlation terms.
  • 2. Therefore,: Consequently, only the diagonal of G contributes to the task loss.The relevant diagonal entries are denoted g_ii.
  • 2. Therefore,: The proposition therefore reduces the multiple-output weighting to the diagonal coefficients g_ii = P_m.This follows directly from the diagonal structure implied by uncorrelated errors.
Loading 2608.16696v1…