1 /*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20 /*!
21 * \file hawkes_ll-inl.h
22 * \brief Log likelihood of a marked self-exciting Hawkes process
23 * \author Caner Turkmen <turkmen.ac@gmail.com>
24 */
25 #ifndef MXNET_OPERATOR_CONTRIB_HAWKES_LL_INL_H_
26 #define MXNET_OPERATOR_CONTRIB_HAWKES_LL_INL_H_
27
28 #include <mxnet/operator.h>
29 #include <vector>
30
31 #include "../operator_common.h"
32 #include "../mshadow_op.h"
33 #include "../mxnet_op.h"
34
35 namespace mxnet {
36 namespace op {
37
38 namespace hawkesll {
39 enum HawkesLLOpInputs {kMu, kAlpha, kBeta, kState, kIATimes, kMarks,
40 kValidLength, kMaxTime};
41 enum HawkesLLGradInputs {kOutGradLL, kOutGradStates, kGradMu, kGradAlpha,
42 kGradBeta, kGradState, kGradIATimes, kGradMarks,
43 kGradValidLength, kGradMaxTime};
44 enum HawkesLLOpOutputs {kOutLL, kOutStates};
45 enum HawkesLLOpResource {kTempSpace};
46 } // namespace hawkesll
47
HawkesLLOpType(const nnvm::NodeAttrs & attrs,std::vector<int> * in_attrs,std::vector<int> * out_attrs)48 inline bool HawkesLLOpType(const nnvm::NodeAttrs& attrs,
49 std::vector<int>* in_attrs,
50 std::vector<int>* out_attrs) {
51 // check dimensions of the type vectors
52 CHECK_EQ(in_attrs->size(), 8U);
53 CHECK_EQ(out_attrs->size(), 2U);
54
55 TYPE_ASSIGN_CHECK(*out_attrs, hawkesll::kOutLL, in_attrs->at(0))
56 TYPE_ASSIGN_CHECK(*out_attrs, hawkesll::kOutStates, in_attrs->at(0))
57
58 for (index_t j = 0; j < 8; ++j) {
59 if (j != hawkesll::kMarks) {
60 TYPE_ASSIGN_CHECK(*in_attrs, j, out_attrs->at(0))
61 }
62 }
63 TYPE_ASSIGN_CHECK(*in_attrs, hawkesll::kMarks, 4) // int32
64
65 return out_attrs->at(hawkesll::kOutLL) != -1;
66 }
67
HawkesLLOpShape(const nnvm::NodeAttrs & attrs,std::vector<TShape> * in_attrs,std::vector<TShape> * out_attrs)68 inline bool HawkesLLOpShape(const nnvm::NodeAttrs& attrs,
69 std::vector<TShape>* in_attrs,
70 std::vector<TShape>* out_attrs) {
71 using namespace mshadow;
72 int N, T, K;
73
74 CHECK_EQ(in_attrs->size(), 8U);
75 CHECK_EQ(out_attrs->size(), 2U);
76
77 // check ndims
78 CHECK_EQ(in_attrs->at(hawkesll::kMu).ndim(), 2); // mu (N, K)
79 CHECK_EQ(in_attrs->at(hawkesll::kAlpha).ndim(), 1); // branching ratio (K,)
80 CHECK_EQ(in_attrs->at(hawkesll::kBeta).ndim(), 1); // decay exponent (K,)
81 CHECK_EQ(in_attrs->at(hawkesll::kState).ndim(), 2); // Hawkes states (N, K)
82 CHECK_EQ(in_attrs->at(hawkesll::kIATimes).ndim(), 2); // i.a. times (N, T)
83 CHECK_EQ(in_attrs->at(hawkesll::kMarks).ndim(), 2); // marks (N, T)
84 CHECK_EQ(in_attrs->at(hawkesll::kValidLength).ndim(), 1); // valid len (N,)
85 CHECK_EQ(in_attrs->at(hawkesll::kMaxTime).ndim(), 1); // max_time (N,)
86
87 N = in_attrs->at(hawkesll::kIATimes)[0]; // number of samples in batch
88 T = in_attrs->at(hawkesll::kIATimes)[1]; // time length
89 K = in_attrs->at(hawkesll::kMu)[1]; // number of marks
90
91 // check inputs consistent
92 CHECK_EQ(in_attrs->at(hawkesll::kMu)[0], N);
93 CHECK_EQ(in_attrs->at(hawkesll::kMu)[1], K);
94 CHECK_EQ(in_attrs->at(hawkesll::kAlpha)[0], K);
95 CHECK_EQ(in_attrs->at(hawkesll::kBeta)[0], K);
96 CHECK_EQ(in_attrs->at(hawkesll::kState)[0], N);
97 CHECK_EQ(in_attrs->at(hawkesll::kState)[1], K);
98 CHECK_EQ(in_attrs->at(hawkesll::kMarks)[0], N);
99 CHECK_EQ(in_attrs->at(hawkesll::kMarks)[1], T);
100 CHECK_EQ(in_attrs->at(hawkesll::kValidLength)[0], N);
101 CHECK_EQ(in_attrs->at(hawkesll::kMaxTime)[0], N);
102
103 // infer output type
104 SHAPE_ASSIGN_CHECK(*out_attrs, hawkesll::kOutLL, Shape1(N))
105 SHAPE_ASSIGN_CHECK(*out_attrs, hawkesll::kOutStates, Shape2(N, K))
106
107 return out_attrs->at(hawkesll::kOutLL).ndim() != 0U &&
108 out_attrs->at(hawkesll::kOutStates).Size() != 0U;
109 }
110
111 template<int req>
112 struct hawkesll_forward {
113 template<typename DType>
Maphawkesll_forward114 MSHADOW_XINLINE static void Map(int i,
115 DType* out_loglike,
116 DType* out_state,
117 const DType* mu,
118 const DType* alpha,
119 const DType* beta,
120 DType* state,
121 const DType* lags,
122 const int32_t* marks,
123 DType* valid_length,
124 DType* max_time,
125 int K,
126 int T,
127 DType* temp_register
128 ) {
129 int32_t ci; // current mark
130 DType ll = 0; // log likelihood
131 DType t = 0; // current time
132 DType d, ed, lda, comp;
133 DType *last_ = &temp_register[i * K];
134
135 const DType *mu_ = &mu[i * K];
136 const DType *lag_ = &lags[i * T];
137 const int32_t *mark_ = &marks[i * T];
138 DType *state_ = &out_state[i * K];
139
140 // iterate over points in sequence
141 for (index_t j = 0; j < valid_length[i]; ++j) {
142 ci = mark_[j];
143 t += lag_[j];
144 d = t - last_[ci];
145 ed = expf(-beta[ci] * d);
146
147 lda = mu_[ci] + alpha[ci] * beta[ci] * state_[ci] * ed;
148 comp = mu_[ci] * d + alpha[ci] * state_[ci] * (1 - ed);
149
150 ll += logf(lda) - comp;
151
152 KERNEL_ASSIGN(state_[ci], req, 1 + (state_[ci] * ed))
153
154 last_[ci] = t;
155 }
156
157 KERNEL_ASSIGN(out_loglike[i], req, ll)
158 }
159 };
160
161 template<int req>
162 struct hawkesll_forward_compensator {
163 template<typename DType>
Maphawkesll_forward_compensator164 MSHADOW_XINLINE static void Map(int i,
165 DType* rem_comp,
166 DType* out_state,
167 const DType* mu,
168 const DType* alpha,
169 const DType* beta,
170 const DType* max_time,
171 const int K,
172 const DType* last_buffer
173 ) {
174 DType d, ed;
175 int m = i % K; // mark
176 int j = i / K; // particle
177
178 // take care of the remaining compensators and state update
179 d = max_time[j] - last_buffer[i];
180 ed = expf(-beta[m] * d);
181
182 // return the remaining compensator
183 KERNEL_ASSIGN(rem_comp[i], req,
184 mu[i] * d + alpha[m] * out_state[i] * (1 - ed))
185
186 // update the state
187 KERNEL_ASSIGN(out_state[i], req, ed * out_state[i])
188 }
189 };
190
191 template<typename xpu>
HawkesLLForward(const nnvm::NodeAttrs & attrs,const OpContext & ctx,const std::vector<TBlob> & inputs,const std::vector<OpReqType> & req,const std::vector<TBlob> & outputs)192 void HawkesLLForward(const nnvm::NodeAttrs& attrs,
193 const OpContext& ctx,
194 const std::vector<TBlob>& inputs,
195 const std::vector<OpReqType>& req,
196 const std::vector<TBlob>& outputs) {
197 using namespace mshadow;
198 using namespace mxnet_op;
199
200 Stream<xpu> *s = ctx.get_stream<xpu>();
201
202 CHECK_EQ(inputs.size(), 8U);
203 CHECK_EQ(outputs.size(), 2U);
204
205 const TBlob& out_loglike = outputs[hawkesll::kOutLL];
206 const TBlob& out_state = outputs[hawkesll::kOutStates];
207
208 int K = inputs[hawkesll::kMu].shape_[1];
209 int N = inputs[hawkesll::kIATimes].shape_[0];
210 int T = inputs[hawkesll::kIATimes].shape_[1];
211
212 MSHADOW_TYPE_SWITCH(out_loglike.type_flag_, DType, {
213 Tensor<xpu, 2, DType> temp_space = ctx.requested[hawkesll::kTempSpace]
214 .get_space_typed<xpu, 2, DType>(
215 Shape2(2*N, K),
216 s);
217
218 Tensor<xpu, 2, DType> last_buffer =
219 Tensor<xpu, 2, DType>(&temp_space.dptr_[0], Shape2(N, K), s);
220 Tensor<xpu, 2, DType> rem_comp =
221 Tensor<xpu, 2, DType>(&temp_space.dptr_[N*K], Shape2(N, K), s);
222
223 Tensor<xpu, 1, DType> out_loglike_ts =
224 out_loglike.get_with_shape<xpu, 1, DType>(Shape1(N), s);
225
226 last_buffer = DType(0.0);
227 rem_comp = DType(0.0);
228
229 Tensor<xpu, 2, DType> out_state_ts =
230 out_state.get_with_shape<xpu, 2, DType>(Shape2(N, K), s);
231 Tensor<xpu, 2, DType> in_state_ts =
232 inputs[hawkesll::kState].get_with_shape<xpu, 2, DType>(Shape2(N, K), s);
233
234 mshadow::Copy(out_state_ts, in_state_ts, s);
235
236 MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
237 Kernel<hawkesll_forward<req_type>, xpu>::Launch(
238 s, N,
239 out_loglike.dptr<DType>(),
240 out_state.dptr<DType>(),
241 inputs[hawkesll::kMu].dptr<DType>(), // mu
242 inputs[hawkesll::kAlpha].dptr<DType>(), // alpha
243 inputs[hawkesll::kBeta].dptr<DType>(), // beta
244 inputs[hawkesll::kState].dptr<DType>(), // states
245 inputs[hawkesll::kIATimes].dptr<DType>(), // interarrival times
246 inputs[hawkesll::kMarks].dptr<int32_t>(), // marks
247 inputs[hawkesll::kValidLength].dptr<DType>(), // valid_length
248 inputs[hawkesll::kMaxTime].dptr<DType>(), // max_time
249 K,
250 T,
251 last_buffer.dptr_);
252 });
253
254 // in parallel, we take care of the remaining compensators
255 MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
256 Kernel<hawkesll_forward_compensator<req_type>, xpu>::Launch(
257 s, N * K,
258 rem_comp.dptr_,
259 out_state.dptr<DType>(),
260 inputs[hawkesll::kMu].dptr<DType>(), // mu
261 inputs[hawkesll::kAlpha].dptr<DType>(), // alpha
262 inputs[hawkesll::kBeta].dptr<DType>(), // beta
263 inputs[hawkesll::kMaxTime].dptr<DType>(), // max_time
264 K,
265 last_buffer.dptr_);
266 });
267 out_loglike_ts -= mshadow::expr::sumall_except_dim<0>(rem_comp);
268 })
269 }
270
271 template<int req>
272 struct hawkesll_backward {
273 template<typename DType>
Maphawkesll_backward274 MSHADOW_XINLINE static void Map(int i, // indexes the sample (particle)
275 DType* mu_gbfr,
276 DType* alpha_gbfr,
277 DType* beta_gbfr, // (N, K)
278 const DType* mu, // (N, K)
279 const DType* alpha, // (K,)
280 const DType* beta, // (K,)
281 const DType* lags, // (N, T)
282 const int32_t* marks, // (N, T)
283 const DType* valid_length, // (N,)
284 const DType* max_time, // (N,)
285 const int K,
286 const int T,
287 DType* last_buffer,
288 DType* phi_buffer,
289 DType* phig_buffer
290 ) {
291 int32_t ci;
292 int32_t part_ix_K = i*K, part_ix_T = i*T;
293
294 DType t = 0, d, lda, ed;
295 DType* last_ = &last_buffer[part_ix_K];
296 DType* state_ = &phi_buffer[part_ix_K];
297 DType* dstate_ = &phig_buffer[part_ix_K];
298
299 DType* mug_ = &mu_gbfr[part_ix_K];
300 DType* alphag_ = &alpha_gbfr[part_ix_K];
301 DType* betag_ = &beta_gbfr[part_ix_K];
302
303 const DType* lag_ = &lags[part_ix_T];
304 const int32_t* mark_ = &marks[part_ix_T];
305
306 // iterate over points
307 for (index_t j = 0; j < valid_length[i]; ++j){
308 ci = mark_[j];
309 t += lag_[j];
310 d = t - last_[ci];
311 ed = expf(-beta[ci] * d);
312
313 lda = mu[part_ix_K + ci] + alpha[ci] * beta[ci] * state_[ci] * ed;
314
315 KERNEL_ASSIGN(mug_[ci], req, mug_[ci] + (1 / lda) - d)
316 KERNEL_ASSIGN(alphag_[ci], req,
317 (
318 alphag_[ci]
319 + (beta[ci] * state_[ci] * ed) / lda
320 - state_[ci] * (1 - ed)
321 )
322 )
323 KERNEL_ASSIGN(betag_[ci], req,
324 betag_[ci]
325 + alpha[ci] * ed
326 * (state_[ci] * (1 - beta[ci] * d) + beta[ci] * dstate_[ci])
327 / lda
328 - alpha[ci]
329 * (dstate_[ci] * (1 - ed) + state_[ci] * d * ed)
330 )
331
332 KERNEL_ASSIGN(dstate_[ci], req, ed * (-d * state_[ci] + dstate_[ci]))
333 KERNEL_ASSIGN(state_[ci], req, 1 + (state_[ci] * ed))
334
335 last_[ci] = t;
336 }
337 }
338 };
339
340
341 template<int req>
342 struct hawkesll_backward_compensator {
343 template<typename DType>
Maphawkesll_backward_compensator344 MSHADOW_XINLINE static void Map(int i,
345 DType* mu_gbfr,
346 DType* alpha_gbfr,
347 DType* beta_gbfr, // (N, K)
348 DType* out_grad, // read this (N,)
349 const DType* mu, // (N, K)
350 const DType* alpha, // (K,)
351 const DType* beta, // (K,)
352 const DType* max_time, // (N,)
353 const int K,
354 DType* last_buffer,
355 DType* phi_buffer,
356 DType* phig_buffer
357 ) {
358 DType d, ed;
359 int m = i % K; // mark
360 int j = i / K; // particle
361 int32_t part_ix_K = j*K;
362 DType* mug_ = &mu_gbfr[part_ix_K];
363 DType* alphag_ = &alpha_gbfr[part_ix_K];
364 DType* betag_ = &beta_gbfr[part_ix_K];
365
366 // take care of the remaining compensators and state update
367 d = max_time[j] - last_buffer[i];
368 ed = expf(-beta[m] * d);
369
370 // take care of the gradients of the remaining compensator
371 KERNEL_ASSIGN(mug_[m], req, mug_[m] - d)
372 KERNEL_ASSIGN(alphag_[m], req,
373 alphag_[m] - phi_buffer[i] * (1 - ed)
374 )
375 KERNEL_ASSIGN(betag_[m], req,
376 betag_[m] - alpha[m] * (
377 phig_buffer[i] * (1 - ed)
378 + phi_buffer[i] * d * ed
379 )
380 )
381
382 // // correct the gradients with respect to output gradients
383 KERNEL_ASSIGN(mug_[m], req, out_grad[j] * mug_[m])
384 KERNEL_ASSIGN(alphag_[m], req, out_grad[j] * alphag_[m])
385 KERNEL_ASSIGN(betag_[m], req, out_grad[j] * betag_[m])
386 }
387 };
388
389 template<typename xpu>
HawkesLLBackward(const nnvm::NodeAttrs & attrs,const OpContext & ctx,const std::vector<TBlob> & inputs,const std::vector<OpReqType> & req,const std::vector<TBlob> & outputs)390 void HawkesLLBackward(const nnvm::NodeAttrs& attrs,
391 const OpContext& ctx,
392 const std::vector<TBlob>& inputs,
393 const std::vector<OpReqType>& req,
394 const std::vector<TBlob>& outputs) {
395 CHECK_EQ(inputs.size(), 10U);
396 CHECK_EQ(outputs.size(), 8U);
397 CHECK_EQ(req.size(), 8U);
398
399 mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
400
401 int K = inputs[hawkesll::kGradMu].shape_[1]; // mu data
402 int N = inputs[hawkesll::kGradIATimes].shape_[0];
403 int T = inputs[hawkesll::kGradIATimes].shape_[1];
404
405 CHECK_EQ(inputs[hawkesll::kOutGradLL].shape_[0], N); // grad of out 0 (LL)
406 CHECK_EQ(inputs[hawkesll::kOutGradStates].shape_[0], N); // grad out 1-states
407 CHECK_EQ(inputs[hawkesll::kOutGradStates].shape_[1], K);
408
409 // sufficient statistics are not differentiated w.r.t.
410 CHECK_EQ(req[hawkesll::kIATimes], OpReqType::kNullOp);
411 CHECK_EQ(req[hawkesll::kMarks], OpReqType::kNullOp);
412 CHECK_EQ(req[hawkesll::kValidLength], OpReqType::kNullOp);
413 CHECK_EQ(req[hawkesll::kMaxTime], OpReqType::kNullOp);
414
415 const TBlob& out_grad = inputs[hawkesll::kOutGradLL];
416
417 using namespace mshadow;
418 using namespace mxnet_op;
419 MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, {
420 // allocate gradient buffers
421 Tensor<xpu, 2, DType> bfr =
422 ctx.requested[hawkesll::kTempSpace]
423 .get_space_typed<xpu, 2, DType>(Shape2(6*N, K), s);
424
425 Tensor<xpu, 2, DType> alpha_gbfr =
426 Tensor<xpu, 2, DType>(&bfr.dptr_[N*K], Shape2(N, K), s);
427 Tensor<xpu, 2, DType> beta_gbfr =
428 Tensor<xpu, 2, DType>(&bfr.dptr_[2*N*K], Shape2(N, K), s);
429 Tensor<xpu, 2, DType> last_buffer =
430 Tensor<xpu, 2, DType>(&bfr.dptr_[3*N*K], Shape2(N, K), s);
431 Tensor<xpu, 2, DType> phig_buffer =
432 Tensor<xpu, 2, DType>(&bfr.dptr_[4*N*K], Shape2(N, K), s);
433 Tensor<xpu, 2, DType> phi_buffer =
434 Tensor<xpu, 2, DType>(&bfr.dptr_[5*N*K], Shape2(N, K), s);
435
436 alpha_gbfr = DType(0.0);
437 beta_gbfr = DType(0.0);
438 last_buffer = DType(0.0);
439 phig_buffer = DType(0.0);
440
441 mshadow::Copy(phi_buffer,
442 inputs[hawkesll::kGradState]
443 .get_with_shape<xpu, 2, DType>(Shape2(N, K), s),
444 s);
445
446 // get the gradient to be output
447 Tensor<xpu, 2, DType> in_grad_mu =
448 outputs[hawkesll::kMu].get_with_shape<xpu, 2, DType>(Shape2(N, K), s);
449 Tensor<xpu, 1, DType> in_grad_alpha =
450 outputs[hawkesll::kAlpha].get_with_shape<xpu, 1, DType>(Shape1(K), s);
451 Tensor<xpu, 1, DType> in_grad_beta =
452 outputs[hawkesll::kBeta].get_with_shape<xpu, 1, DType>(Shape1(K), s);
453
454 in_grad_mu = DType(0.0);
455
456 MXNET_ASSIGN_REQ_SWITCH(req[hawkesll::kMu], req_type, {
457 Kernel<hawkesll_backward<req_type>, xpu>::Launch(
458 s,
459 N,
460 in_grad_mu.dptr_, alpha_gbfr.dptr_, beta_gbfr.dptr_, // gradients
461 inputs[hawkesll::kGradMu].dptr<DType>(), // mu data
462 inputs[hawkesll::kGradAlpha].dptr<DType>(), // alpha data
463 inputs[hawkesll::kGradBeta].dptr<DType>(), // beta data
464 inputs[hawkesll::kGradIATimes].dptr<DType>(), // lags data
465 inputs[hawkesll::kGradMarks].dptr<int32_t>(), // marks data
466 inputs[hawkesll::kGradValidLength].dptr<DType>(), // valid_length data
467 inputs[hawkesll::kGradMaxTime].dptr<DType>(), // max_time data
468 K,
469 T,
470 last_buffer.dptr_, // buffer to keep timestamp of last item
471 phi_buffer.dptr_, // "states"
472 phig_buffer.dptr_); // derivatives of "states"
473 });
474
475 MXNET_ASSIGN_REQ_SWITCH(req[hawkesll::kMu], req_type, {
476 Kernel<hawkesll_backward_compensator<req_type>, xpu>::Launch(
477 s,
478 N * K,
479 in_grad_mu.dptr_, alpha_gbfr.dptr_, beta_gbfr.dptr_, // gradients
480 out_grad.dptr<DType>(),
481 inputs[hawkesll::kGradMu].dptr<DType>(), // mu data
482 inputs[hawkesll::kGradAlpha].dptr<DType>(), // alpha data
483 inputs[hawkesll::kGradBeta].dptr<DType>(), // beta data
484 inputs[hawkesll::kGradMaxTime].dptr<DType>(), // max_time data
485 K,
486 last_buffer.dptr_, // buffer to keep timestamp of last item
487 phi_buffer.dptr_, // "states"
488 phig_buffer.dptr_); // derivatives of "states"
489 });
490
491 // reduce the gradients
492 Assign(in_grad_alpha, req[hawkesll::kAlpha],
493 mshadow::expr::sumall_except_dim<1>(alpha_gbfr)
494 )
495
496 Assign(in_grad_beta, req[hawkesll::kBeta],
497 mshadow::expr::sumall_except_dim<1>(beta_gbfr)
498 )
499 })
500 }
501
502 } // namespace op
503 } // namespace mxnet
504
505 #endif // MXNET_OPERATOR_CONTRIB_HAWKES_LL_INL_H_
506