| """Controlled encoder-scale ablation; fixed frame-wise LayerNorm, no extra bits. | |
| The failed baseline's continuous features grew to hundreds in RMS, saturating | |
| the tanh quantizer. Normalize before both continuous decoding and quantization | |
| so the warmup and transition have a consistent feature-scale constraint. | |
| This changes the encoder architecture and must be identified in checkpoints. | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from tinyasr.tinycall import TinyCallCodec, TinyCallEncoder | |
| ARCHITECTURE = 'unit_layernorm_encoder_v1' | |
| class UnitScaleEncoder(TinyCallEncoder): | |
| def forward(self, audio): | |
| features = super().forward(audio) | |
| return F.layer_norm(features.transpose(1, 2), (self.dimension,), eps=1e-5).transpose(1, 2) | |
| class UnitScaleCodec(TinyCallCodec): | |
| def __init__(self): | |
| super().__init__() | |
| # Preserve the baseline's exact initial weights and RNG consumption. | |
| # Fixed LayerNorm has neither learned parameters nor running statistics. | |
| with torch.random.fork_rng(devices=[]): | |
| encoder = UnitScaleEncoder() | |
| encoder.load_state_dict(self.encoder.state_dict(), strict=True) | |
| self.encoder = encoder | |