Ternary weight networks

Introduction

这两天看了一下这篇文章，我就这里分享一下，不过我还是只记录一下跟别人blog上没有，或者自己的想法(ps: 因为有时候翻blog时候发现每篇都一样还是挺烦的= =) 。为了不重复前人的工作，我post一个不小心翻到的博客权值简化（1）：三值神经网络（Ternary Weight Networks），整个论文内容及实现都讲的很全面了，可以翻阅一下，我也借鉴一下。

文中主要工作的点在三个方面:

• 增加了网络的表达力（expressive ability）。在{1，0，1}基础上增加了 \(\alpha\) 作为scaled factor；
• 压缩模型大小。当然主要是weight的压缩。比起FPWN（full precision weight network）有16~32x的提升，但是BPWN（binary precision weight network）的2x大小（ps:当然在TWN的caffe代码里面，都由float double类型存储，因为这需要在应该上方面来实现）；
• 减少计算需求。主要相比于BPWN增多了0，当然这方面也需硬件来获得提升，在该caffe代码里面并没有；

Ternary Quantization

在我的理解看来，文中最核心的内容是：将有约束的并且两变量之间互相依赖的优化问题，逐步拆分最后用具有先验的统计方法来近视解决。

最初的优化问题：



将\(W^{t}\)的约束具体化为：



并将其带入公式（1），将\(W^{t*}\)的优化转化为\(\Delta^*\)的优化：



然后对公式（4）中的\(\alpha\)求偏导，得到：



因为\(\alpha\)和\(\Delta\)相互依赖，将（5）代入（4）消去\(\alpha\)：



但问题来了，公式（6）依然没法求，而文中就根据先验知识，假设\(W_i\)服从\(N(0,\sigma^2)\)分布，近视的\(\Delta^*\)为\(0.6\sigma\)（\(0.6\sigma\)等于\(0.75E(|W|)\)）。因此作者采用粗暴的方法，把\(\Delta^*\)设为\(\Delta^*\approx0.7E(|W|)\approx\frac{n}{0.7}\sum_{i=1}^n|W_i|\)

//caffe-twns
//blob.cpp
template <typename Dtype>
void Blob<Dtype>::set_delta(){
  float scale_factor = TERNARY_DELTA * 1.0 / 10; //delta = 0.7
  Dtype delta = (Dtype) scale_factor * this->asum_data() / this->count(); // 0.7*(E|W_i|)/num
  delta = (delta <= 100) ? delta : 100;
  delta = (delta >= -100) ? delta : -100;
  this->delta_ = delta;
}

template <typename Dtype>
void Blob<Dtype>::set_delta(Dtype delta){
  delta = (delta <= 100) ? delta : 100;
  delta = (delta >= -100) ? delta : -100;
  this->delta_ = delta;
}

Implement

我借用一张图



步骤3~5，其中第5步代码在上面：

template <typename Dtype>
void Blob<Dtype>::ternarize_data(Phase phase){

if(phase == RUN){

// if(DEBUG) print_head();

 //LOG(INFO) << "RUN phase...";
 // caffe_sleep(3);
 return; // do nothing for the running phase
}else if(phase == TRAIN){
 //LOG(INFO) << "TRAIN phase ...";
 // caffe_sleep(3);
}else{
 //LOG(INFO) << "TEST phase ...";
 // caffe_sleep(3);
}

  // const Dtype delta = 0; // default value;
  // const Dtype delta = (Dtype) 0.8 * this->asum_data() / this->count();
  this->set_delta();  //defualt 0.7*(E|W_i|)/num or set by user
  const Dtype delta = this->get_delta();
  Dtype alpha = 1;

  if (!data_) { return; }
  switch (data_->head()) {
  case SyncedMemory::HEAD_AT_CPU:
{
	caffe_cpu_ternary<Dtype>(this->count(), delta, this->cpu_data(), this->mutable_cpu_binary()); //quantized weight to ternary
	alpha = caffe_cpu_dot(this->count(), this->cpu_binary(), this->cpu_data());  //scale-alpha: (E |W_i|)   i belong to I_delta
	alpha /= caffe_cpu_dot(this->count(), this->cpu_binary(), this->cpu_binary()); //(1/num_binary)*alpha
	caffe_cpu_scale(this->count(), alpha, this->cpu_binary(), this->mutable_cpu_binary());
	// this->set_alpha(alpha);
}
    return;
  case SyncedMemory::HEAD_AT_GPU:
  case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
{
    caffe_gpu_ternary<Dtype>(this->count(), delta, this->gpu_data(), this->mutable_gpu_binary());
	Dtype* pa = new Dtype(0);
	caffe_gpu_dot(this->count(), this->gpu_binary(), this->gpu_data(), pa);
	Dtype* pb = new Dtype(0);
	caffe_gpu_dot(this->count(), this->gpu_binary(), this->gpu_binary(), pb);

	alpha = (*pa) / ((*pb) + 1e-6);
	this->set_alpha(alpha);

	caffe_gpu_scale(this->count(), alpha, this->gpu_binary(), this->mutable_gpu_binary());
	// this->set_alpha((Dtype)1);

    // LOG(INFO) << "alpha = " << alpha;
	// caffe_sleep(3);
}
    return;
#else
    NO_GPU;
#endif
  case SyncedMemory::UNINITIALIZED:
    return;
  default:
    LOG(FATAL) << "Unknown SyncedMemory head state: " << data_->head();
  }
}

步骤6~7，其中在第6步作者在caffe-twns直接采用传统caffe的方法，而$Z=XW\approx X(\alpha W^t)=(\alpha X)\bigoplus W^t $更偏向与在硬件加速的优化（因为本身在caffe-twns的ternary就采用float或者double，并且用blas或cudnn加速也无法直接跳过0值）：

//conv_layer.cpp
template <typename Dtype>
void ConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {

  // const Dtype* weight = this->blobs_[0]->cpu_data();
if(BINARY){
  this->blobs_[0]->binarize_data();
} 

if(TERNARY){
  this->blobs_[0]->ternarize_data(this->phase_);  //quantized from blob[0] to ternary sand stored in cpu_binary()
/*
    Dtype alpha = (Dtype) this->blobs_[0]->get_alpha();

for(int i=0; i<bottom.size(); i++){
  Blob<Dtype>* blob = bottom[i];
  caffe_cpu_scale(blob->count(), alpha, blob->cpu_data(), blob->mutable_cpu_data());
}
*/

}
  const Dtype* weight = (BINARY || TERNARY) ? this->blobs_[0]->cpu_binary() : this->blobs_[0]->cpu_data();
...
}

步骤11~19，weight的Update是在full precision上，而计算gradient则是用ternary weight：

//conv_layer.cpp
template <typename Dtype>
void ConvolutionLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  const Dtype* weight = this->blobs_[0]->cpu_data();
  Dtype* weight_diff = this->blobs_[0]->mutable_cpu_diff();
  for (int i = 0; i < top.size(); ++i) {
    ...
    if (this->param_propagate_down_[0] || propagate_down[i]) {
      for (int n = 0; n < this->num_; ++n) {
        // gradient w.r.t. weight. Note that we will accumulate diffs.
        if (this->param_propagate_down_[0]) {
          this->weight_cpu_gemm(bottom_data + n * this->bottom_dim_,
              top_diff + n * this->top_dim_, weight_diff);
        }
        // gradient w.r.t. bottom data, if necessary.
        if (propagate_down[i]) {
          this->backward_cpu_gemm(top_diff + n * this->top_dim_, weight,
              bottom_diff + n * this->bottom_dim_);
        }
      }
    }
  }
}