This talk was presented as part of the Trends and Advances in Monte Carlo Sampling Algorithms Workshop.

1. Stochastic Gradient MCMC for
Independent and Correlated Data
Yi-An (Yian) Ma
University of California, Berkeley
Tianqi ChenEmily Fox Nick Foti Felix Ye

2. Issue = Efficiency
multiple modes
strongly correlated
across dimensions
parameters observations
complex
Goal: Large-Data Posteriors

3. Goal: Large-Data Posteriors
largeparameters observations
Issue = Scalability
E.g.:
Wikipedia corpus analysis,
Human genome sequence,
Ion channel recordings

4. Classical Approach:
MCMC via Jump Processes
“Standard” method = Metropolis-Hastings (MH)
– Propose θ’ from kernel depending on past value θ
– Accept or reject
Example of a jump process
Often, inefficiently explores posterior

5. Continuous dynamic based samplers
aka Grad-MCMC
Use (stochastic) dynamics on energy landscape
to simulate distant proposal

6. Example: Hamiltonian Monte Carlo (HMC)
Hamiltonian
(total energy H)
Target posterior of θ
Add auxiliary “momentum”
variable r
Focus on Potential Energy
Kinetic Energy

7. Simulate Hamiltonian Dynamics
Use Hamiltonian dynamics to collect samples on a fixed
(continuous-time) interval

8. . . .
Resample Momentum for Ergodicity

9. Continuous Dynamic-Based Samplers
Dynamic 1 Dynamic 2 Dynamic 3
invariant under these dynamics
Simulated samples are desired posterior samples
Want:
(+ ergodicity)

10. Is there a general
recipe for
construction?
All Continuous
Markov Processes
Exploring “Correct” Dynamics
Processes with
ps(θ) = π(θ)
Langevin
dynamics (LD)
HMC
Riemann
Manifold
HMC
Riemann Manifold LD

11. Assume target distribution
A Recipe for Continuous Dynamics MCMC
d-dim Wiener process
parameters auxiliary vars
Target total energy
Ma, Chen, Fox, NIPS 2015.
SDE
PSD D
skew-sym Q
total energy H
is invariant

12. Assume target distribution
parameters auxiliary vars
Target total energy
Ma, Chen, Fox, NIPS 2015.
A Recipe for Continuous Dynamics MCMC

13. Recipe is Complete
All Continuous
Markov Processes
ps(z) =
π(z)
SDE defined by
D(z),
Q(z)
Ma, Chen, Fox, NIPS 2015.
All existing samplers can
be written in framework:
– HMC
– Riemann HMC
– LD
– Riemann LD
Any valid sampler
has a D and Q
in our framework

14. Continuous dynamic based samplers
aka Grad-MCMC
The nitty gritty practical issues

15. A Practical Algorithm
Consider ε – discretization
Some discretization error
m steps
Construct irreversible MH to correct this bias

16. Example: Metropolis-Hastings
Propose sample using
proposal dist. q
Calculate
accept-reject
ratio
Ratio of reverse
and forward
proposals
Accept or reject

17. Irreversible Jump: Naïve Revision
f
f
g
Forward proposal
Reverse proposal
Leads to wrong stationary distribution
Naïve Irreversible

18. Correcting the Algorithm
Jump to
adjoint process
upon rejection
Introduce auxiliary variable zp ~ U{1,-1}:
zp = 1
zp = -1
I-Jump Algorithm: Ma, Fox, Chen, Wu, arXiv 2016.

19. Irreversible MALA Algorithm
Continuous dynamical process
(e.g., irreversible SDE)
Adjoint process
Use as
f(z|y)
Use as
g(z|y)
Compute for one
step of dynamics
I-MALA Algorithm: Ma, Fox, Chen, Wu, arXiv 2016.
Q → - Q

20. Consider ε – discretization
m steps
Must compute gradient!
– Computations costly for large data
– Cannot handle streaming data
DATA
A Practical Algorithm

21. Scaling Grad-MCMC:
Handling large datasets efficiently

22. Scaling up: Stochastic Gradients
Compute noisy gradient based on minibatch
consisting of n i.i.d. observations:
– For minibatches sampled uniformly at random from data,
Only requires examining n data points
True gradientNoisy gradient
DATAand we assume (appealing to CLT):

23. Scalable Version of Algorithm
Original update rule:
Modified stochastic gradient update rule:
As , SG noise decreases and bias  0
Use small, finite in practice (allow some bias)
Subtract estimate of variance of SG noise

24. Example D and Q of Past Algorithms
D(z)
Q(z)
(D, Q) Space not
previously
explored
SGHMC:
SGLD:
SGRLD:
SGNHT:
SGLD
SGRLD

25. • Use existing D, Q building blocks to define new samplers
• SGRHMC existence previously only speculated
– Naïve(ish) approach has wrong stationary distribution
• Take D and Q of SGHMC and make state dependent
SGRiemannHMC
Ma, Chen, Fox, NIPS 2015.
State-dependent
pos. def. matrix

26. Perplexity
Iteration
SGLD
SGHMC
SGRHMC
SGRLD
Applied SGRHMC (using Fisher info metric) to online LDA
– Latent Dirichlet allocation (LDA) = mixed membership
document model
Scraped Wikipedia entries in a streaming manner
– Each entry was analyzed on-the-fly
Streaming Wikipedia Analysis
Ma, Chen, Fox, NIPS 2015.
Step size
selected via
grid search

27. Scaling Approach 2:
Stochastic Gradient MCMC
for Correlated datasets

28. Scaling up: Stochastic Gradients
Compute noisy gradient based on minibatch
consisting of n i.i.d. observations:
– For minibatches sampled uniformly at random from data,
Only requires examining n data points
DATA
Welling, Teh, ICML 2011. Chen, Fox, Guestrin, ICML 2014.
and we assume (appealing to CLT):

29. Scaling up: Stochastic Gradients
Compute noisy gradient based on minibatch
consisting of n i.i.d. observations:
– For minibatches sampled uniformly at random from data,
Only requires examining n data points
DATA
Welling, Teh, ICML 2011. Chen, Fox, Guestrin, ICML 2014.
and we assume (appealing to CLT):
i.i.d.
What about non i.i.d data?
E.g., time series:
:
OK
. . . .
How?
:

30. Hidden Markov Models (HMMs)
discrete state sequence
observations
transition probabilities,
observation parameters

31. Batch learning for HMMs:
A quick review

32. Batch Learning for HMMs
Use current to form local state beliefs:
– Propagate info forwards to form

33. Use current to form local state beliefs:
– Propagate info backwards
Batch Learning for HMMs

34. Batch Learning for HMMs
Combine to form smoothed local state belief:

35. Given local beliefs, update global parameter
Batch Learning for HMMs
Issue: Cost is O(K2T) per global update!
Costly when using uninformed initializations
or observations are redundant

36. Minibatch learning for HMMs
via SG-MCMC

37. Why is this not so straightforward?
SG-MCMC assumes continuous parameter space
– Typical HMM MCMC algorithms iterate on
(sampling) latent discrete-valued state sequence
Need to prove that correct stationary distribution
is maintained in presence of:
– Incomplete observations in subsequences
– Mutually correlated subsequences per minibatch
Ma, Foti, Fox, ICML 2017

38. Marginal Likelihood Representation
Marginalize x

39. Marginal Likelihood Representation
…

40. Rewriting in terms of a specific subsequence
…
…

41. Rewriting in terms of a specific subsequence
Random dynamical system:
Synchronizes as
Ye, Ma, Qian, arXiv 2017

42. Potential energy for Grad-MCMC

43. Issues with the gradient computation
q,π calculations
involve touching
(nearly) all T obs!
sum over all
subsequences

44. A Stochastic Gradient Approach
q,π calculations
involve touching
(nearly) all T obs!
sum over all
subsequences
𝑠

45. Approximating Gradient Terms with Buffering
𝑠

46. 𝑠
Approximating Gradient Terms with Buffering

47. How much buffering is sufficient?
B𝐿 𝑠B
Set buffer length B by estimating the Lyapunov exponent
of the underlying random dynamical system
Ye, Ma, Qian, arXiv 2017

48. Minibatch = Set of Subsequences
Subsequences are correlated! Reduces efficiency
B𝐿 𝑠B

49. Mitigating Subsequence Correlations
Minimum gap: ν
Based on
2nd largest eig(A)
B𝐿 𝑠B

50. Resulting Gradient Approximation
Plugs into
SG-MCMC
theory
B𝐿 𝑠B

51. Ion Channel Analysis – Segmentations
716.19 sec 7245.14 sec2124.45 sec
44.05 sec 138.51 sec 466.82 sec
1 MHz recording of single alamethicin channel
Our dataset: 209,634 observations
BatchGrad-
MCMC
SG-MCMC
[Rosenstein et al. 2013]

52. Ion Channel Analysis – Estimation

53. Consequence of non-i.i.d. data
& importance of buffering
DiagonallyDominantReversedCycles
log predictive prob || A − Atrue ||F

54. Thank You

55. Irreversibility:
Increasing Sampler Efficiency

56. When Q(z) is non-zero, process is irreversible
– i.e., time-reversed process is statistically
distinguishable from forward process
– Saw greater efficiency for such processes
(e.g., HMC, Riemann HMC, SGHMC, SGRHMC,…)
Reversibility: Continuous Dynamics
Skew-symmetric
Hwang, Hwang-Ma and Cheu (1993, 2005); Rey-Bellet and Spiliopoulos (2014)

57. Reversibility: Jump Processes
Reversibility
Irreversibility
explores distribution in a directed manner
asymmetric, cyclic motion
Chen, Lovasz and Pak (1999); Diaconis, Holmes and Neal (2000); Bierkens(2015)
symmetric dynamics

58. Reversibility: Jump Processes
Reversibility
Irreversibility
Chen, Lovasz and Pak (1999); Diaconis, Holmes and Neal (2000); Bierkens(2015)
Easy!
e.g., Metropolis-Hastings
Hard

59. Example: Metropolis-Hastings
Propose sample using
proposal dist. q
Calculate
accept-reject
ratio
Ratio of reverse
and forward
proposals
Accept or reject

60. Irreversible Jump: Naïve Revision
f
f
g
Forward proposal
Reverse proposal
Leads to wrong stationary distribution
Naïve Irreversible

61. What can we do?
For continuous dynamic part:
SDE(f(z))
Defines diffusion
Hard to specify!

62. What can we do?
For continuous dynamic part:
SDE(D(z),Q(z))
Defines
irreversibility
PSD Skew-sym
For jump part:
MJP(W(z|x))
Transition
kernel
Hard to specify!
MH is one choice,
but reversible

63. What can we do?
For continuous dynamic part:
SDE(D(z),Q(z))
Defines
irreversibility
PSD Skew-sym
For jump part:
MJP(S(x,z),A(x,z))
Antisymmetric
kernel
Symmetric
kernel
Defines
irreversibility
“Only” need
Ma, Fox, Chen, Wu, arXiv 2016.

64. Correcting the Algorithm
f
f
g
Naïve Irreversible

65. Correcting the Algorithm
Jump to
adjoint process
upon rejection
Introduce auxiliary variable zp ~ U{1,-1}:
zp = 1
zp = -1
Lifting Method: Turitsyn, Chertkov and Vucelja (2011)

66. Simple Irreversible Jump Algorithm
Keep track of
process in use
Flip the process
Can show:
Ma, Fox, Chen, Wu, arXiv 2016.

67. Simple Irreversible Jump Algorithm
Possible choice:
f
g
f
g
Ma, Fox, Chen, Wu, arXiv 2016.

68. zp
Proposal
rejected
(z,1) (z*,1)1
(z*,1)2
(z*,1)3
(z,-1)

69. zp
(z,1) (z1,1)
(z4,-1)
(z2,1)
(z2,-1)
(z3,-1)

70. Moving to Higher Dimensions
zp
In 1D, there’s only one choice
of direction: In higher dimensions:
1. Let
uniformly distributed
on unit ball
2. Flip sign of
upon rejection
3. After multiple rejections,
resample

71. Approach 1:
jump process
with π
invariant
Approach 2:
continuous
dynamics with π
invariant
Combined:
irreversible MALA
Combining Approach 1 & 2

72. Irreversible MALA Algorithm
Metropolis Adjusted Langevin Algorithm (MALA) algorithm: Xifara, Sherlock ,
Livingstone, Byrne and Girolami (2014); Zig-Zag: Bierkens and Roberts (2016)
Continuous dynamical process
(e.g., irreversible SDE)
Adjoint process
Use as
f(z|y)
Use as
g(z|y)
Can compute for one
step of dynamics

73. Thank You