Skip to content

Commit d912554

Browse files
xingniuxingniu
authored
Add blocking cross-attention between decoder and encoded prepended tokens (#1085)
* Add blocking cross-attention between decoder and encoded prepended tokens * Use a new dictionary-based prepared data format
1 parent 13c63be commit d912554

25 files changed

Lines changed: 427 additions & 100 deletions

‎CHANGELOG.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,21 @@ Note that Sockeye has checks in place to not translate with an old model that wa
1111

1212
Each version section may have subsections for: _Added_, _Changed_, _Removed_, _Deprecated_, and _Fixed_.
1313

14+
## [3.1.32]
15+
16+
### Added
17+
18+
- Sockeye now supports blocking cross-attention between decoder and encoded prepended tokens.
19+
- If the source contains prepended text and a tag indicating the end of prepended text,
20+
Sockeye supports blocking the cross-attention between decoder and encoded prepended tokens (including the tag).
21+
To enable this operation, specify `--end-of-prepending-tag` for training or data preparation,
22+
and `--transformer-block-prepended-cross-attention` for training.
23+
24+
### Changed
25+
26+
- Sockeye uses a new dictionary-based prepared data format that supports storing length of prepended source tokens
27+
(version 7). The previous format (version 6) is still supported.
28+
1429
## [3.1.31]
1530

1631
### Fixed

‎docs/training.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,3 +185,10 @@ This is similar to using `--restrict-lexicon` for `sockeye-translate` with the a
185185
To use NVS simply specify `--neural-vocab-selection` to `sockeye-train`.
186186
This will train a model with NVS that is automatically used by `sockeye-translate`.
187187
If you want look at translations without vocabulary selection specify `--skip-nvs` as an argument to `sockeye-translate`.
188+
189+
## Prepended Source Text
190+
191+
If the source contains prepended text and a tag indicating the end of prepended text,
192+
Sockeye supports blocking the cross-attention between decoder and encoded prepended tokens (including the tag).
193+
To enable this operation, specify `--end-of-prepending-tag` for training or data preparation,
194+
and `--transformer-block-prepended-cross-attention` for training.

‎sockeye/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,4 @@
1111
# express or implied. See the License for the specific language governing
1212
# permissions and limitations under the License.
1313

14-
__version__ = '3.1.31'
14+
__version__ = '3.1.32'

‎sockeye/arguments.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,12 @@ def add_training_data_args(params, required=False):
443443
required=required,
444444
type=regular_file(),
445445
help='Target side of parallel training data.')
446+
params.add_argument('--end-of-prepending-tag',
447+
type=str,
448+
default=None,
449+
help='Tag indicating the end of prepended text. Prepended tokens before this tag (inclusive) '
450+
'will be marked, and they will not be counted toward source length when calculating '
451+
'maximum output length for beam search.')
446452

447453

448454
def add_validation_data_params(params):
@@ -687,6 +693,11 @@ def add_model_parameters(params):
687693
choices=C.POSITIONAL_EMBEDDING_TYPES,
688694
default=C.FIXED_POSITIONAL_EMBEDDING,
689695
help='The type of positional embedding. Default: %(default)s.')
696+
model_params.add_argument('--transformer-block-prepended-cross-attention',
697+
action='store_true',
698+
default=False,
699+
help='Block cross-attention between decoder and encoded prepended tokens. '
700+
'Default: %(default)s.')
690701
model_params.add_argument('--transformer-preprocess',
691702
type=multiple_values(num_values=2, greater_or_equal=None, data_type=str),
692703
default=('n', 'n'),

‎sockeye/constants.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
UNK_ID = VOCAB_SYMBOLS.index(UNK_SYMBOL)
2929
BOS_ID = VOCAB_SYMBOLS.index(BOS_SYMBOL)
3030
EOS_ID = VOCAB_SYMBOLS.index(EOS_SYMBOL)
31+
INVALID_ID = -1 # an example of invalid ids (i.e., negative integers)
3132
# reserve extra space for the EOS or BOS symbol that is added to both source and target
3233
SPACE_FOR_XOS = 1
3334

@@ -336,13 +337,18 @@
336337
FIXED_PARAM_STRATEGY_ENCODER_HALF_AND_SOURCE_EMBEDDINGS]
337338

338339
# data sharding
340+
DATA_KEY_SOURCE = 'source'
341+
DATA_KEY_TARGET = 'target'
342+
DATA_KEY_PREPENDED_SOURCE_LENGTH = 'prepended_source_length'
339343
SHARD_NAME = "shard.%05d"
340344
SHARD_SOURCE = SHARD_NAME + ".source"
341345
SHARD_TARGET = SHARD_NAME + ".target"
346+
SHARD_PREPENDED_SOURCE_LENGTH = SHARD_NAME + ".prepended_source_length"
342347
DATA_INFO = "data.info"
343348
DATA_CONFIG = "data.config"
344349
PREPARED_DATA_VERSION_FILE = "data.version"
345-
PREPARED_DATA_VERSION = 6
350+
PREPARED_DATA_VERSION = 7
351+
PREPARED_DATA_LEGACY_VERSION = 6
346352

347353
# reranking metric options
348354
RERANK_BLEU = "bleu"

‎sockeye/data_io.py‎

Lines changed: 156 additions & 29 deletions
Large diffs are not rendered by default.

‎sockeye/decoder.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -200,26 +200,23 @@ def init_state_from_encoder(self,
200200
[autoregressive state dummies] * num_layers.
201201
202202
:param encoder_outputs: Encoder outputs. Shape: (batch, source_length, encoder_dim).
203-
:param encoder_valid_length: Valid lengths of encoder outputs. Shape: (batch,).
203+
:param encoder_valid_length: Valid lengths of encoder outputs. Shape: (batch, 2).
204204
:param target_embed: Target-side embedding layer output. Shape: (batch, target_length, target_embedding_dim).
205205
:return: Initial states.
206206
"""
207207
source_max_len = encoder_outputs.size()[1]
208+
# (batch * heads, 1, source_max_len)
209+
source_mask = layers.prepare_source_length_mask(encoder_valid_length, self.config.attention_heads,
210+
source_max_len, mask_prepended_tokens=
211+
self.config.block_prepended_cross_attention)
208212
if target_embed is None: # Inference: initial step = 0. Shape: (batch_size, 1)
209-
steps = pt.zeros_like(encoder_valid_length).unsqueeze(1)
210-
# (batch * heads, 1, source_max_len)
211-
source_mask = layers.prepare_source_length_mask(encoder_valid_length, self.config.attention_heads,
212-
source_max_len)
213+
steps = pt.zeros_like(encoder_valid_length[:, :1])
213214
# Shape: (batch, heads, 1, src_max_len)
214215
source_mask = source_mask.view(-1, self.config.attention_heads, 1, source_max_len)
215216
else: # Training: steps up to target length. Shape: (1, target_length)
216217
target_length = target_embed.size()[1]
217218
steps = pt.arange(0, target_length, device=target_embed.device).unsqueeze(0)
218-
# (batch * heads, 1, source_max_len)
219-
source_mask = layers.prepare_source_length_mask(encoder_valid_length, self.config.attention_heads,
220-
source_max_len)
221219
source_mask = source_mask.expand(-1, target_length, -1) # Shape: (batch * heads, trg_max_len, src_max_len)
222-
223220
# Shape: (batch, heads, trg_max_len, src_max_len)
224221
source_mask = source_mask.view(-1, self.config.attention_heads, target_length, source_max_len)
225222

‎sockeye/generate_decoder_states.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -152,13 +152,15 @@ def init_store_file(self, initial_size: int) -> None:
152152
def generate_states_and_store(self,
153153
sources: List[str],
154154
targets: List[str],
155-
batch_size: int) -> None:
155+
batch_size: int,
156+
eop_id: int = C.INVALID_ID) -> None:
156157
"""
157158
Generate decoder states by force-decoding the sentence pairs in `sources` and `targets` with a NMT model.
158159
159160
:param sources: list of source segments.
160161
:param targets: list of target segments.
161162
:param batch_size: number of sentence pairs to decode at once.
163+
:param eop_id: End-of-prepending tag id.
162164
"""
163165
assert self.state_store_file != None, \
164166
"You should call probe_token_count first to initialize the store files."
@@ -171,7 +173,8 @@ def generate_states_and_store(self,
171173
target_vocabs=self.target_vocabs,
172174
batch_size=batch_size,
173175
max_seq_len_source=self.max_seq_len_source,
174-
max_seq_len_target=self.max_seq_len_target
176+
max_seq_len_target=self.max_seq_len_target,
177+
eop_id=eop_id
175178
)
176179

177180
with pt.inference_mode():
@@ -254,7 +257,7 @@ def store(args: argparse.Namespace):
254257
args.state_dtype, C.KNN_WORD_DATA_STORE_DTYPE, device)
255258
generator.num_states = DecoderStateGenerator.probe_token_count(targets[0], max_seq_len_target)
256259
generator.init_store_file(generator.num_states)
257-
generator.generate_states_and_store(sources, targets, args.batch_size)
260+
generator.generate_states_and_store(sources, targets, args.batch_size, model.eop_id)
258261
generator.save_config()
259262

260263

@@ -272,6 +275,8 @@ def main():
272275
level=args.loglevel) # pylint: disable=no-member
273276

274277
utils.log_basic_info(args)
278+
if args.end_of_prepending_tag is not None:
279+
logger.warning("The end-of-prepending tag defined in the model will be used.")
275280

276281
store(args)
277282

‎sockeye/inference.py‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
from . import utils
3232
from . import vocab
3333
from .beam_search import CandidateScorer, get_search_algorithm, GreedySearch, SearchResult
34-
from .data_io import tokens2ids
34+
from .data_io import tokens2ids, get_prepended_token_length
3535
from .model import SockeyeModel
3636

3737
logger = logging.getLogger(__name__)
@@ -858,6 +858,10 @@ def num_source_factors(self) -> int:
858858
def num_target_factors(self) -> int:
859859
return self.models[0].num_target_factors
860860

861+
@property
862+
def eop_id(self) -> int:
863+
return self.models[0].eop_id
864+
861865
def translate(self, trans_inputs: List[TranslatorInput], fill_up_batches: bool = True) -> List[TranslatorOutput]:
862866
"""
863867
Batch-translates a list of TranslatorInputs, returns a list of TranslatorOutputs.
@@ -1001,13 +1005,14 @@ def _get_inference_input(self,
10011005
optional target prefix, and optional target prefix factors.
10021006
"""
10031007
batch_size = len(trans_inputs)
1004-
lengths = [len(inp) for inp in trans_inputs]
10051008

10061009
max_target_prefix_length = max(inp.num_target_prefix_tokens for inp in trans_inputs)
10071010
max_target_prefix_factors_length = max(inp.num_target_prefix_factors for inp in trans_inputs)
10081011
max_length = max(len(inp) for inp in trans_inputs)
10091012
# assembling source ids on cpu array (faster) and copy to Translator.device (potentially GPU) in one go below.
10101013
source_np = np.zeros((batch_size, max_length, self.num_source_factors), dtype='int32')
1014+
# total token length and prepended token length
1015+
length_np = np.zeros((batch_size, 2), dtype='int32')
10111016

10121017
target_prefix_np = np.zeros((batch_size, max_target_prefix_length), dtype='int32') \
10131018
if max_target_prefix_length > 0 else None
@@ -1019,9 +1024,13 @@ def _get_inference_input(self,
10191024
max_output_lengths = [] # type: List[int]
10201025
for j, trans_input in enumerate(trans_inputs):
10211026
num_tokens = len(trans_input) # includes eos
1022-
max_output_lengths.append(self._get_max_output_length(num_tokens))
1023-
source_np[j, :num_tokens, 0] = tokens2ids(itertools.chain(trans_input.get_source_prefix_tokens(),
1024-
trans_input.tokens), self.source_vocabs[0])
1027+
primary_source_ids = tokens2ids(itertools.chain(trans_input.get_source_prefix_tokens(),
1028+
trans_input.tokens), self.source_vocabs[0])
1029+
source_np[j, :num_tokens, 0] = primary_source_ids
1030+
length_np[j, 0] = num_tokens
1031+
length_np[j, 1] = get_prepended_token_length(primary_source_ids, self.eop_id)
1032+
# the effective source length excludes prepended tokens
1033+
max_output_lengths.append(self._get_max_output_length(length_np[j, 0] - length_np[j, 1]))
10251034
if target_prefix_np is not None and trans_input.num_target_prefix_tokens > 0:
10261035
target_prefix_np[j, :trans_input.num_target_prefix_tokens] = \
10271036
tokens2ids(trans_input.get_target_prefix_tokens(), self.vocab_targets[0])
@@ -1068,7 +1077,7 @@ def _get_inference_input(self,
10681077
"will default to not using a restrict lexicon.")
10691078

10701079
source = pt.tensor(source_np, device=self.device, dtype=pt.int32)
1071-
source_length = pt.tensor(lengths, device=self.device, dtype=pt.int32) # shape: (batch_size,)
1080+
source_length = pt.tensor(length_np, device=self.device, dtype=pt.int32) # shape: (batch_size, 2)
10721081
max_out_lengths = pt.tensor(max_output_lengths, device=self.device, dtype=pt.int32)
10731082
target_prefix = pt.tensor(target_prefix_np, device=self.device, dtype=pt.int32) \
10741083
if target_prefix_np is not None else None

‎sockeye/layers.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -320,16 +320,26 @@ def forward(self,
320320
return interleaved_matmul_encdec_valatt(key_values, probs, heads=self.heads)
321321

322322

323-
def prepare_source_length_mask(lengths: pt.Tensor, heads: int, max_length: int, expand=True) -> pt.Tensor:
323+
def prepare_source_length_mask(lengths: pt.Tensor, heads: int, max_length: int, expand: bool = True,
324+
mask_prepended_tokens: bool = True) -> pt.Tensor:
324325
"""
325-
lengths: (batch_size,)
326-
expand: Expand to the heads.
326+
Prepare source length masks where positions of invalid tokens are marked as True.
327+
328+
:param lengths: Total source length and prepended source length. Shape: (batch_size, 2)
329+
:param heads: Number of attention heads.
330+
:param max_length: Maximum sequence length.
331+
:param expand: Expand to the heads.
332+
:param mask_prepended_tokens: Mask prepended tokens.
333+
:return: Source length mask.
327334
"""
328335
# (batch_size, max_len)
329-
mask = ~(pt.arange(max_length, device=lengths.device).unsqueeze(0) < lengths.reshape((-1, 1)))
336+
mask = ~(pt.arange(max_length, device=lengths.device).unsqueeze(0) < lengths[:, :1])
337+
if mask_prepended_tokens:
338+
prepended_token_mask = pt.arange(max_length, device=lengths.device).unsqueeze(0) < lengths[:, 1:2]
339+
mask |= prepended_token_mask
330340
if expand:
331-
# (batch_size*heads, 1, max_len)
332-
mask = mask.unsqueeze(1).expand(-1, heads, -1).reshape((-1, max_length)).unsqueeze(1)
341+
# (batch_size * heads, 1, max_len)
342+
mask = mask.unsqueeze(1).expand(-1, heads, -1).reshape((-1, max_length)).unsqueeze(1)
333343
return mask
334344

335345

@@ -673,7 +683,7 @@ def separate_kv(module: pt.nn.Module):
673683
def get_positional_embeddings(length: int, depth: int) -> pt.Tensor:
674684
utils.check_condition(depth % 2 == 0, "Positional embeddings require an even embedding size it "
675685
"is however %d." % depth)
676-
# (1, depth)
686+
# (1, depth/2)
677687
channels = pt.arange(depth // 2).unsqueeze(0)
678688

679689
# (length, 1)
@@ -683,7 +693,7 @@ def get_positional_embeddings(length: int, depth: int) -> pt.Tensor:
683693
sin = pt.sin(scaled_positions)
684694
# cosines:
685695
cos = pt.cos(scaled_positions)
686-
# interleave: (length, num_embed)
696+
# stack sin and cos: (length, depth)
687697
encodings = pt.hstack([sin, cos])
688698
return encodings
689699

0 commit comments

Comments
 (0)