SubspaceInference.jl

The subspace inference method for Deep Neural Networks (DNN) and ordinary differential equations (ODEs) are implemented as a package named SubspaceInference.jl in Julia.

Subspace Construction

PCA based subspace construction

SubspaceInference.subspace_construction — Method

subspace_construction(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1)

Izmailov, P., Maddox, W. J., Kirichenko, P., Garipov, T., Vetrov, D., & Wilson, A. G. (2020, August). Subspace inference for Bayesian deep learning. In Uncertainty in Artificial Intelligence (pp. 1169-1179). PMLR.

To construct subspace from pretrained weights.

Input Arguments

model : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
cost : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
data : Inputs and outputs. Eg: X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
opt : Optimzer. Eg: opt = ADAM(0.1)

Keyword Arguments

T : Number of steps for subspace calculation. Eg: T= 1
c : Moment update frequency. Eg: c = 1
M : Maximum number of columns in deviation matrix. Eg: M= 2
print_freq: Loss printing frequency

Outputs

W_swa : Mean weights
P : Projection Matrix

source

Diffusion map based subspace construction

SubspaceInference.diffusion_subspace — Method

diffusion_subspace(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1)

To construct subspace from pretrained weights using diffusion maps.

Input Arguments

model : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
cost : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
data : Inputs and outputs. Eg: X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
opt : Optimzer. Eg: opt = ADAM(0.1)

Keyword Arguments

T : Number of steps for subspace calculation. Eg: T= 1
c : Moment update frequency. Eg: c = 1
M : Maximum number of columns in deviation matrix. Eg: M= 2
print_freq: Loss printing frequency

Outputs

W_swa : Mean weights
P : Projection Matrix

source

Autoencoder based subspace construction

SubspaceInference.auto_encoder_subspace — Method

auto_encoder_subspace(model, cost, data, opt, encoder, decoder; T = 10, c = 1, M = 3, print_freq = 1

)

To construct subspace from pretrained weights using autoencoders.

Input Arguments

model : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
cost : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
data : Inputs and outputs. Eg: X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
opt : Optimzer. Eg: opt = ADAM(0.1)
encoder : Encoder to generate subspace from NN or Neural ODE parameters
decoder : Decoder to generate NN or Neural ODE parameters from subspace

Keyword Arguments

T : Number of steps for subspace calculation. Eg: T= 1
c : Moment update frequency. Eg: c = 1
M : Maximum number of columns in deviation matrix. Eg: M= 2
print_freq: Loss printing frequency

Outputs

W_swa : Mean weights
decoder : Trained decoder to generate NN or Neural ODE parameters from subspace

source