Fundamental Limits and Tradeoffs in Invariant Representation Learning
†‡§?

‡‡

Han Zhao , Chen Dan , Bryon Aragam , Tommi Jaakkola , Geoffrey Gordon , Pradeep Ravikumar†‡?
§University of Illinois Urbana Champaign, Carnegie Mellon University, University of Chicago, MIT
hanzhao@illinois.edu, {cdan,pradeepr,ggordon}@cs.cmu.edu, bryon@chicagobooth.edu, tommi@csail.mit.edu
Example and Applications

Invariant Representation Learning

Regression

Domain Adaptation/Generalization: mitigating distribution shifts

Problem Setup: Given a joint distribution µ over (X, A, Y ), wewould like to learn a (randomized) mapping g : X × A 7→ Z suchthat the marginal distribution of Z is invariant to the attribute A:

Under the noiseless setting, the tradeoff problem has the following
form: maxZ=g(X,A) Var E[Y |Z] − λ Var E[A|Z]. Define the information plane for regression problems as the feasible region of the tradeoff problem: RLS := {(Var E[Y |Z], Var E[A|Z]) : Z = g(X, A)}.

Target data

TargetOptimalPredictor

Pr(Z|A=a)=Pr(Z),∀a∈AµµSource
OptimalPredictor• X ⊆ Rd: input space of the data
• A: attribute space (classification: A = {0, 1}; regression: A = R)
• Y: label space (classification: Y = {0, 1}; regression: Y = R)

Invariant
Representations

Practical Implementation via Adversarial Training:

Fair Representations: mitigating bias in data
0min max0

Male or Female ?

1g,hageBC
education
BC
B
Crace
B
CB
C
genderBC
B annual income CBCB repaying history CBC@defaulting historyAcredit score

Applicants`Y (h(g(X, A)), Y ) − λ · `A(h (g(X, A)), A)• Classification: `Y = `A = cross-entropy loss
• Regression: `Y = `A = mean-squared error• λ: tradeoff parameter: λ = 0 (accuracy); λ → ∞ (invariance)=<latexit sha1_base64="DtnBsMFJKt0GW8M0Mud6iLS7M5Q=">AAAETnicfZPNbhMxEMfdboEQvlo4clkRVeIUJQWpXCJVLZXCrQ1NWymJKq93NrHij63tbbKx9gm4wntx5UW4IfCmQUkcwUgrjeb3nxl71hOljGrTaPzY2g52Hjx8VHlcffL02fMXu3svL7XMFIEukUyq6whrYFRA11DD4DpVgHnE4Coan5T86g6UplJcmDyFAcdDQRNKsHGh89bNbq1Rb8wt3HSaC6eGFnZ2sxcc9mNJMg7CEIa17jUbqRlYrAwlDIpqP9OQYjLGQ+g5V2AOemDnJy3CfReJw0Qq9wkTzqOrGRZzrXMeOSXHZqR9Vgb/xcyIF36s7KKLcG5rzImNlEyvZ5R1N8O9zCQfBpaKNDMgyP0tkoyFRoblTMOYKiCG5c7BRFE3iJCMsMLEuMlX90O/RYpnsqiuteB4DEpK7iaFBQHmDiBgQiTnWMS2T0Vc2H6ZG0X2U+HRCdDhyPwVJHbiC6bTJZz6MM+XMPfhbLaEs42+k//01KlYUIKZ/byJjQJYUVz4inZ7hbZ9Ckq5B2X7d1hBqimTwhPEbnlWCnz0C8SQwG1ZgJQv2T1E0fJb3DpNYVsrEk9xOk2Xv+XU79DpLGGnKNyqNf3F2nQuD+rNd/WD8/e1o+PF0lXQa/QGvUVNdIiOUBudoS4iCNAX9BV9C74HP4Nfwe976fbWIucVWrOdyh/3CoRU</latexit>hTo focus on the fundamental tradeoff, we focus on the noiseless set0ting and we assume both h and h have infinite capacity, i.e, they are
perfect predictors.

Learning to Rank: matching query/document distributions

q1 d1
q1 d2q1 d3
Source

μS

q2 d1
q2 d2
q2 d3
TargetμT(Z1, Z2, Z3)
Feature map

Ranking scores
Projection

Aligned featuredistributions
ZZ
μS μT

0.1
0.9
0.6 0.4
0.3
0.7

SoG

Ranks
312 231

Research Questions
Question:Is there any limitation on using invariant representations? If yes,
what is the fundamental tradeoff between utility (accuracy) andinvariance (distribution matching)?
Our Answer: Yes in general, and we can characterize the tradeoff
on an information plane, where the optimal tradeoff depends on thecoupling between the target Y and the attribute A.

Under the noiseless setting, the tradeoff problem has the following
form: maxZ=g(X,A) I(Y ; Z) − λI(A; Z). We define the information
plane for classification problems as the feasible region of the tradeoff
problem: RCE := {(I(Y ; Z), I(A; Z)) : Z = g(X, A)}.

Z=X

Var(A) · ρ2Y A
Z≡c

Var(Y ) · ρ2Y A

Var E[Y | Z]
Var(Y )In regression problems, we can also precisely characterize the Pareto
frontier of the following problem:max

Z=g(X,A)

Var E[Y | Z],

s.t. Var E[A | Z] ≤ c

Var E[A | Z]

I(A; Z)

Var(A)

I(A; Y )

Z=X

Z=X

2
Var(A) · ρY A

I(A; Y )Z≡c

I(Y ; Z)∆Y |A · H(A)H(Y )• ∆Y |A := | Prµ(Y = 1 | A = 0) − Prµ(Y = 1 | A = 1)|.

NeurIPS-2023 (JMLR 2022)

Var(A)• ρY A := correlation coefficient between Y and A.

Classification

H(A)

Var(Y ) · ρ2AY

∆A|Y · H(Y )(X1, X2, X3)0

Var E[A | Z]

Var(A) · ρ2AY

Source dataPresenter: Han ZhaoZ≡c2Var(Y ) · ρY A

Var E[Y | Z]
Var(Y )

Equation for the Pareto frontier:
Var E[Y |Z] ≤ Var(Y )q2222ρY A (1 − ρY A)α(1 − α) + 1 − α − ρY A + 2αρY A• α := c/ Var(A).

Dec. 14, 2023