def multi_dimensional_attention(rep_tensor,
rep_mask,
scope=None,
keep_prob=1.,
is_train=None,
wd=0.,
activation='elu',
tensor_dict=None,
name=None,
reuse=None):
bs, sl, vec = tf.shape(rep_tensor)[0], tf.shape(rep_tensor)[1], tf.shape(
rep_tensor)[2]
ivec = rep_tensor.get_shape()[2]
with tf.variable_scope(scope or 'multi_dimensional_attention',
reuse=reuse):
map1 = bn_dense_layer(rep_tensor, ivec, True, 0., 'bn_dense_map1',
activation, False, wd, keep_prob, is_train)
map2 = bn_dense_layer(map1, ivec, True, 0., 'bn_dense_map2', 'linear',
False, wd, keep_prob, is_train)
map2_masked = exp_mask_for_high_rank(map2, rep_mask)
soft = tf.nn.softmax(map2_masked, 1) # bs,sl,vec
attn_output = tf.reduce_sum(soft * rep_tensor, 1) # bs, vec
# save attn
if tensor_dict is not None and name is not None:
tensor_dict[name] = soft
return attn_output
def traditional_attention(rep_tensor,
rep_mask,
scope=None,
keep_prob=1.,
is_train=None,
wd=0.,
activation='elu',
tensor_dict=None,
name=None,
reuse=None):
bs, sl, vec = tf.shape(rep_tensor)[0], tf.shape(rep_tensor)[1], tf.shape(
rep_tensor)[2]
ivec = rep_tensor.get_shape()[2]
with tf.variable_scope(scope or 'traditional_attention', reuse=reuse):
rep_tensor_map = bn_dense_layer(rep_tensor, ivec, True, 0.,
'bn_dense_map', activation, False, wd,
keep_prob, is_train)
rep_tensor_logits = get_logits([rep_tensor_map],
None,
False,
scope='self_attn_logits',
mask=rep_mask,
input_keep_prob=keep_prob,
is_train=is_train) # bs,sl
attn_res = softsel(rep_tensor, rep_tensor_logits, rep_mask) # bs,vec
# save attn
if tensor_dict is not None and name is not None:
tensor_dict[name] = tf.nn.softmax(rep_tensor_logits)
return attn_res
def bi_sru_recurrent_network(
rep_tensor, rep_mask, is_train=None, keep_prob=1., wd=0.,
scope=None, hn=None, reuse=None):
"""
:param rep_tensor: [Tensor/tf.float32] rank is 3 with shape [batch_size/bs, max_sent_len/sl, vec]
:param rep_mask: [Tensor/tf.bool]rank is 2 with shape [bs,sl]
:param is_train: [Scalar Tensor/tf.bool]scalar tensor to indicate whether the mode is training or not
:param keep_prob: [float] dropout keep probability in the range of (0,1)
:param wd: [float]for L2 regularization, if !=0, add tensors to tf collection "reg_vars"
:param scope: [str]variable scope name
:param hn:
:param
:return: [Tensor/tf.float32] with shape [bs, sl, 2vec] for forward and backward
"""
bs, sl, vec = tf.shape(rep_tensor)[0], tf.shape(rep_tensor)[1], tf.shape(rep_tensor)[2]
ivec = rep_tensor.get_shape().as_list()[2]
ivec = hn or ivec
with tf.variable_scope(scope or 'bi_sru_recurrent_network'):
# U_d = bn_dense_layer([rep_tensor], 6 * ivec, False, 0., 'get_frc', 'linear',
# False, wd, keep_prob, is_train) # bs, sl, 6vec
# U_d_fw, U_d_bw = tf.split(U_d, 2, 2)
with tf.variable_scope('forward'):
U_d_fw = bn_dense_layer([rep_tensor], 3 * ivec, False, 0., 'get_frc_fw', 'linear',
False, wd, keep_prob, is_train) # bs, sl, 6vec
U_fw = tf.concat([rep_tensor, U_d_fw], -1)
fw_SRUCell = SwitchableDropoutWrapper(SRUCell(ivec, tf.nn.tanh, reuse), is_train, keep_prob)
fw_output, _ = dynamic_rnn(
fw_SRUCell, U_fw, tf.reduce_sum(tf.cast(rep_mask, tf.int32), -1),
dtype=tf.float32, scope='forward_sru') # bs, sl, vec
with tf.variable_scope('backward'):
U_d_bw = bn_dense_layer([rep_tensor], 3 * ivec, False, 0., 'get_frc_bw', 'linear',
False, wd, keep_prob, is_train) # bs, sl, 6vec
U_bw = tf.concat([rep_tensor, U_d_bw], -1)
bw_SRUCell = SwitchableDropoutWrapper(SRUCell(ivec, tf.nn.tanh, reuse), is_train, keep_prob)
bw_output, _ = bw_dynamic_rnn(
bw_SRUCell, U_bw, tf.reduce_sum(tf.cast(rep_mask, tf.int32), -1),
dtype=tf.float32, scope='backward_sru') # bs, sl, vec
all_output = tf.concat([fw_output, bw_output], -1) # bs, sl, 2vec
return all_output
def __call__(self, inputs, state, scope=None):
"""
:param inputs: [bs, vec]
:param state:
:param scope:
:return:
"""
with tf.variable_scope(scope or "SRU_cell"):
b_f = tf.get_variable('b_f', [self._num_units], dtype=tf.float32,
initializer=tf.constant_initializer(0))
b_r = tf.get_variable('b_r', [self._num_units], dtype=tf.float32,
initializer=tf.constant_initializer(0))
U_d = bn_dense_layer(inputs, 3 * self._num_units, False, 0., 'get_frc',
'linear', keep_prob=self.keep_prob,
is_train=self.is_train) # bs, 3vec
x_t = tf.identity(inputs, 'x_t')
x_dt, f_t, r_t = tf.split(U_d, 3, 1)
f_t = tf.nn.sigmoid(f_t + b_f)
r_t = tf.nn.sigmoid(r_t + b_r)
c_t = f_t * state + (1 - f_t) * x_dt
h_t = r_t * self._activation(c_t) + (1 - r_t) * x_t
return h_t, c_t
def directional_attention_with_dense(rep_tensor,
rep_mask,
direction=None,
scope=None,
keep_prob=1.,
is_train=None,
wd=0.,
activation='elu',
tensor_dict=None,
name=None,
hn=None):
def scaled_tanh(x, scale=5.):
return scale * tf.nn.tanh(1. / scale * x)
bs, sl, vec = tf.shape(rep_tensor)[0], tf.shape(rep_tensor)[1], tf.shape(
rep_tensor)[2]
ivec = rep_tensor.get_shape().as_list()[2]
ivec = hn or ivec
with tf.variable_scope(scope or 'directional_attention_%s' % direction
or 'diag'):
# mask generation
sl_indices = tf.range(sl, dtype=tf.int32)
sl_col, sl_row = tf.meshgrid(sl_indices, sl_indices)
if direction is None:
direct_mask = tf.cast(
tf.diag(-tf.ones([sl], tf.int32)) + 1, tf.bool)
else:
if direction == 'forward':
direct_mask = tf.greater(sl_row, sl_col)
else:
direct_mask = tf.greater(sl_col, sl_row)
direct_mask_tile = tf.tile(tf.expand_dims(direct_mask, 0),
[bs, 1, 1]) # bs,sl,sl
rep_mask_tile = tf.tile(tf.expand_dims(rep_mask, 1),
[1, sl, 1]) # bs,sl,sl
attn_mask = tf.logical_and(direct_mask_tile, rep_mask_tile) # bs,sl,sl
# non-linear
rep_map = bn_dense_layer(rep_tensor, ivec, True, 0., 'bn_dense_map',
activation, False, wd, keep_prob, is_train)
rep_map_tile = tf.tile(tf.expand_dims(rep_map, 1),
[1, sl, 1, 1]) # bs,sl,sl,vec
rep_map_dp = dropout(rep_map, keep_prob, is_train)
# attention
with tf.variable_scope('attention'): # bs,sl,sl,vec
f_bias = tf.get_variable('f_bias', [ivec], tf.float32,
tf.constant_initializer(0.))
dependent = linear(rep_map_dp,
ivec,
False,
scope='linear_dependent',
is_train=is_train) # bs,sl,vec
dependent_etd = tf.expand_dims(dependent, 1) # bs,1,sl,vec
head = linear(rep_map_dp,
ivec,
False,
scope='linear_head',
is_train=is_train) # bs,sl,vec
head_etd = tf.expand_dims(head, 2) # bs,sl,1,vec
logits = scaled_tanh(dependent_etd + head_etd + f_bias,
5.0) # bs,sl,sl,vec
logits_masked = exp_mask_for_high_rank(logits, attn_mask)
attn_score = tf.nn.softmax(logits_masked, 2) # bs,sl,sl,vec
attn_score = mask_for_high_rank(attn_score, attn_mask)
attn_result = tf.reduce_sum(attn_score * rep_map_tile,
2) # bs,sl,vec
with tf.variable_scope('output'):
o_bias = tf.get_variable('o_bias', [ivec], tf.float32,
tf.constant_initializer(0.))
# input gate
fusion_gate = tf.nn.sigmoid(
linear(rep_map, ivec, True, 0., 'linear_fusion_i', False, wd,
keep_prob, is_train) +
linear(attn_result, ivec, True, 0., 'linear_fusion_a', False,
wd, keep_prob, is_train) + o_bias)
output = fusion_gate * rep_map + (1 - fusion_gate) * attn_result
output = mask_for_high_rank(output, rep_mask)
# save attn
if tensor_dict is not None and name is not None:
tensor_dict[name + '_dependent'] = dependent
tensor_dict[name + '_head'] = head
tensor_dict[name] = attn_score
tensor_dict[name + '_gate'] = fusion_gate
return output
def simple_block_attention(rep_tensor,
rep_mask,
block_len=5,
scope=None,
direction=None,
keep_prob=1.,
is_train=None,
wd=0.,
activation='elu',
hn=None):
assert direction is not None
def scaled_tanh(x, scale=5.):
return scale * tf.nn.tanh(1. / scale * x)
bs, sl, vec = tf.shape(rep_tensor)[0], tf.shape(rep_tensor)[1], tf.shape(
rep_tensor)[2]
org_ivec = rep_tensor.get_shape().as_list()[2]
ivec = hn or org_ivec
with tf.variable_scope(scope or 'block_simple'):
# @1. split sequence
with tf.variable_scope('split_seq'):
block_num = tf.cast(
tf.ceil(
tf.divide(tf.cast(sl, tf.float32),
tf.cast(block_len, tf.float32))), tf.int32)
comp_len = block_num * block_len - sl
rep_tensor_comp = tf.concat(
[rep_tensor,
tf.zeros([bs, comp_len, org_ivec], tf.float32)], 1)
rep_mask_comp = tf.concat([
rep_mask,
tf.cast(tf.zeros([bs, comp_len], tf.int32), tf.bool)
], 1)
rep_tensor_split = tf.reshape(
rep_tensor_comp,
[bs, block_num, block_len, org_ivec]) # bs,bn,bl,d
rep_mask_split = tf.reshape(rep_mask_comp,
[bs, block_num, block_len]) # bs,bn,bl
# non-linear
rep_map = bn_dense_layer(rep_tensor_split, ivec, True, 0.,
'bn_dense_map', activation, False, wd,
keep_prob, is_train) # bs,bn,bl,vec
rep_map_tile = tf.tile(tf.expand_dims(rep_map, 2),
[1, 1, block_len, 1, 1]) # bs,bn,bl,bl,vec
# rep_map_dp = dropout(rep_map, keep_prob, is_train)
bn = block_num
bl = block_len
with tf.variable_scope('self_attention'):
# @2.self-attention in block
# mask generation
sl_indices = tf.range(block_len, dtype=tf.int32)
sl_col, sl_row = tf.meshgrid(sl_indices, sl_indices)
if direction == 'forward':
direct_mask = tf.greater(sl_row, sl_col) # bl,bl
else:
direct_mask = tf.greater(sl_col, sl_row) # bl,bl
direct_mask_tile = tf.tile(
tf.expand_dims(tf.expand_dims(direct_mask, 0), 0),
[bs, bn, 1, 1]) # bs,bn,bl,bl
rep_mask_tile_1 = tf.tile(tf.expand_dims(rep_mask_split, 2),
[1, 1, bl, 1]) # bs,bn,bl,bl
rep_mask_tile_2 = tf.tile(tf.expand_dims(rep_mask_split, 3),
[1, 1, 1, bl]) # bs,bn,bl,bl
rep_mask_tile = tf.logical_and(rep_mask_tile_1, rep_mask_tile_2)
attn_mask = tf.logical_and(direct_mask_tile,
rep_mask_tile,
name='attn_mask') # bs,bn,bl,bl
# attention
f_bias = tf.get_variable('f_bias', [ivec], tf.float32,
tf.constant_initializer(0.))
dependent_head = linear(rep_map, 2 * ivec, False, 0.,
'linear_dependent_head', False, wd,
keep_prob, is_train) # bs,bn,bl,2vec
dependent, head = tf.split(dependent_head, 2, 3)
dependent_etd = tf.expand_dims(dependent, 2) # bs,bn,1,bl,vec
head_etd = tf.expand_dims(head, 3) # bs,bn,bl,1,vec
logits = scaled_tanh(dependent_etd + head_etd + f_bias,
5.0) # bs,bn,bl,bl,vec
logits_masked = exp_mask_for_high_rank(logits, attn_mask)
attn_score = tf.nn.softmax(logits_masked, 3) # bs,bn,bl,bl,vec
attn_score = mask_for_high_rank(attn_score,
attn_mask) # bs,bn,bl,bl,vec
self_attn_result = tf.reduce_sum(attn_score * rep_map_tile,
3) # bs,bn,bl,vec
with tf.variable_scope('source2token_self_attn'):
inter_block_logits = bn_dense_layer(self_attn_result, ivec, True,
0., 'bn_dense_map', 'linear',
False, wd, keep_prob,
is_train) # bs,bn,bl,vec
inter_block_logits_masked = exp_mask_for_high_rank(
inter_block_logits, rep_mask_split) # bs,bn,bl,vec
inter_block_soft = tf.nn.softmax(inter_block_logits_masked,
2) # bs,bn,bl,vec
inter_block_attn_output = tf.reduce_sum(
self_attn_result * inter_block_soft, 2) # bs,bn,vec
with tf.variable_scope('self_attn_inter_block'):
inter_block_attn_output_mask = tf.cast(tf.ones([bs, bn], tf.int32),
tf.bool)
block_ct_res = directional_attention_with_dense(
inter_block_attn_output, inter_block_attn_output_mask,
direction, 'disa', keep_prob, is_train, wd,
activation) # [bs,bn,vec]
block_ct_res_tile = tf.tile(tf.expand_dims(
block_ct_res, 2), [1, 1, bl, 1]) #[bs,bn,vec]->[bs,bn,bl,vec]
with tf.variable_scope('combination'):
# input:1.rep_map[bs,bn,bl,vec]; 2.self_attn_result[bs,bn,bl,vec]; 3.rnn_res_tile[bs,bn,bl,vec]
rep_tensor_with_ct = tf.concat(
[rep_map, self_attn_result, block_ct_res_tile],
-1) # [bs,bn,bl,3vec]
new_context_and_gate = linear(rep_tensor_with_ct, 2 * ivec, True,
0., 'linear_new_context_and_gate',
False, wd, keep_prob,
is_train) # [bs,bn,bl,2vec]
new_context, gate = tf.split(new_context_and_gate, 2,
3) # bs,bn,bl,vec
if activation == "relu":
new_context_act = tf.nn.relu(new_context)
elif activation == "elu":
new_context_act = tf.nn.elu(new_context)
elif activation == "linear":
new_context_act = tf.identity(new_context)
else:
raise RuntimeError
gate_sig = tf.nn.sigmoid(gate)
combination_res = gate_sig * new_context_act + (
1 - gate_sig) * rep_map # bs,bn,bl,vec
with tf.variable_scope('restore_original_length'):
combination_res_reshape = tf.reshape(
combination_res, [bs, bn * bl, ivec]) # bs,bn*bl,vec
output = combination_res_reshape[:, :sl, :]
return output