Skip to content

Transformer Decoder: make cross attention config optional #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
NeoLegends merged 1 commit into from
Jan 19, 2026
Merged

Transformer Decoder: make cross attention config optional #90

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions i6_models/assemblies/transformer/transformer_decoder_v1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ class TransformerDecoderBlockV1Config(ModelConfiguration):
Attributes:
ff_cfg: Configuration for ConformerPositionwiseFeedForwardV1
mhsa_cfg: Configuration for CausalSelfAttentionV1
cross_cfg: Configuration for CrossAttentionV1
cross_cfg: Configuration for CrossAttentionV1.
Can be set to None in case there is no cross attention block (e.g. for LM usage).
modules: List of modules to use for ConformerBlockV2:
- "ff" for feed forward module
- "mhcsa" for multi-head causal self attention module
Expand All @@ -58,7 +59,7 @@ class TransformerDecoderBlockV1Config(ModelConfiguration):

ff_cfg: ConformerPositionwiseFeedForwardV2Config
mhsa_cfg: CausalSelfAttentionV1Config
cross_cfg: CrossAttentionV1Config
cross_cfg: Optional[CrossAttentionV1Config]
modules: List[str] = field(default_factory=lambda: ["mhcsa", "cross", "ff"])
scales: List[float] = field(default_factory=lambda: [1.0, 1.0, 1.0])

Expand All @@ -67,6 +68,9 @@ def __post__init__(self):

assert len(self.modules) == len(self.scales), "modules and scales must have same length"
assert all(name in ["ff", "mhcsa", "cross"] for name in self.modules), "module type not supported"
assert "cross" not in self.modules or self.cross_cfg is not None, (
"must specify cross attention config when enabling the cross attention module"
)


class TransformerDecoderBlockV1State(TypedDict):
Expand All @@ -88,6 +92,9 @@ def __init__(self, cfg: TransformerDecoderBlockV1Config):
elif module_name == "mhcsa":
modules.append(CausalSelfAttentionV1(cfg.mhsa_cfg))
elif module_name == "cross":
assert cfg.cross_cfg is not None, (
"must specify cross attention config when enabling the cross attention module"
)
modules.append(CrossAttentionV1(cfg.cross_cfg))
else:
raise NotImplementedError
Expand Down